diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..0412f9155080bf6be91d0d6bde1be6c5149b0e04
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+# BEGIN oh-my-opencode-slim clonedeps
+.slim/clonedeps/repos/
+# END oh-my-opencode-slim clonedeps
diff --git a/.ignore b/.ignore
new file mode 100644
index 0000000000000000000000000000000000000000..32dbf4fbaccb437ffb678b1d500966f6577eea01
--- /dev/null
+++ b/.ignore
@@ -0,0 +1,9 @@
+# BEGIN oh-my-opencode-slim clonedeps
+!.slim/
+!.slim/clonedeps.json
+!.slim/clonedeps/
+!.slim/clonedeps/repos/
+!.slim/clonedeps/repos/**
+.slim/clonedeps/repos/**/.git/
+.slim/clonedeps/repos/**/.git/**
+# END oh-my-opencode-slim clonedeps
diff --git a/.slim/clonedeps.json b/.slim/clonedeps.json
new file mode 100644
index 0000000000000000000000000000000000000000..32b7a16acaee268e19463f3ed5408c6bf6f9345e
--- /dev/null
+++ b/.slim/clonedeps.json
@@ -0,0 +1,23 @@
+{
+ "version": "1.0.0",
+ "updatedAt": "2026-08-08T14:28:19.000Z",
+ "dependencies": [
+ {
+ "name": "diffusers",
+ "resolvedVersion": "665f578278365ea4a3318cb8c9b66ce6c01204b9 (refs/pull/14371/head pin)",
+ "repoUrl": "https://github.com/huggingface/diffusers.git",
+ "ref": "665f578278365ea4a3318cb8c9b66ce6c01204b9",
+ "path": ".slim/clonedeps/repos/huggingface__diffusers",
+ "packagePath": "src/diffusers/modular_pipelines",
+ "reason": "Pinned WIP-PR diffusers this repo subclasses (SequentialPipelineBlocks + minimax_h3 block classes); source for h3_split_blocks.py wiring and h3_aoti.py block export"
+ },
+ {
+ "name": "multimodalart/qwen3vl-conditioner",
+ "resolvedVersion": "20d92357b94f1ee6a93b4fcaeaa06b9d151122d7 (main)",
+ "repoUrl": "https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner",
+ "ref": "main",
+ "path": ".slim/clonedeps/repos/multimodalart__qwen3vl-conditioner",
+ "reason": "The other half of the split: defines the cross-Space contract (CANVASES labels, wire format) this repo must match; hosts h3_core / prompt_rewrite"
+ }
+ ]
+}
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/.gitignore b/.slim/clonedeps/repos/huggingface__diffusers/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..79889b2b3d4591186fc1ad152fb5a6a78c860e9a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/.gitignore
@@ -0,0 +1,188 @@
+# Initially taken from GitHub's Python gitignore file
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# tests and logs
+tests/fixtures/cached_*_text.txt
+logs/
+lightning_logs/
+lang_code_data/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a Python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vs
+.vscode
+
+# Cursor
+.cursor
+
+# Pycharm
+.idea
+
+# TF code
+tensorflow_code
+
+# Models
+proc_data
+
+# examples
+runs
+/runs_old
+/wandb
+/examples/runs
+/examples/**/*.args
+/examples/rag/sweep
+
+# data
+/data
+serialization_dir
+
+# emacs
+*.*~
+debug.env
+
+# vim
+.*.swp
+
+# ctags
+tags
+
+# pre-commit
+.pre-commit*
+
+# .lock
+*.lock
+
+# DS_Store (MacOS)
+.DS_Store
+
+# RL pipelines may produce mp4 outputs
+*.mp4
+*.jpg
+*.jepg
+*.wav
+
+# dependencies
+/transformers
+
+# ruff
+.ruff_cache
+
+# wandb
+wandb
+
+# AI agent generated symlinks
+/.agents/skills
+/.claude/skills
\ No newline at end of file
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/AGENTS.md b/.slim/clonedeps/repos/huggingface__diffusers/AGENTS.md
new file mode 100644
index 0000000000000000000000000000000000000000..b28461c924a4217b9478062bea99d2306d7d4d05
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/AGENTS.md
@@ -0,0 +1 @@
+.ai/AGENTS.md
\ No newline at end of file
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/CITATION.cff b/.slim/clonedeps/repos/huggingface__diffusers/CITATION.cff
new file mode 100644
index 0000000000000000000000000000000000000000..2f3658ef2cb8db5dbfd78c3413a336a5d6c957e3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/CITATION.cff
@@ -0,0 +1,52 @@
+cff-version: 1.2.0
+title: 'Diffusers: State-of-the-art diffusion models'
+message: >-
+ If you use this software, please cite it using the
+ metadata from this file.
+type: software
+authors:
+ - given-names: Patrick
+ family-names: von Platen
+ - given-names: Suraj
+ family-names: Patil
+ - given-names: Anton
+ family-names: Lozhkov
+ - given-names: Pedro
+ family-names: Cuenca
+ - given-names: Nathan
+ family-names: Lambert
+ - given-names: Kashif
+ family-names: Rasul
+ - given-names: Mishig
+ family-names: Davaadorj
+ - given-names: Dhruv
+ family-names: Nair
+ - given-names: Sayak
+ family-names: Paul
+ - given-names: Steven
+ family-names: Liu
+ - given-names: William
+ family-names: Berman
+ - given-names: Yiyi
+ family-names: Xu
+ - given-names: Thomas
+ family-names: Wolf
+repository-code: 'https://github.com/huggingface/diffusers'
+abstract: >-
+ Diffusers provides pretrained diffusion models across
+ multiple modalities, such as vision and audio, and serves
+ as a modular toolbox for inference and training of
+ diffusion models.
+keywords:
+ - deep-learning
+ - pytorch
+ - image-generation
+ - hacktoberfest
+ - diffusion
+ - text2image
+ - image2image
+ - score-based-generative-modeling
+ - stable-diffusion
+ - stable-diffusion-diffusers
+license: Apache-2.0
+version: 0.12.1
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md b/.slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md
new file mode 100644
index 0000000000000000000000000000000000000000..b28461c924a4217b9478062bea99d2306d7d4d05
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md
@@ -0,0 +1 @@
+.ai/AGENTS.md
\ No newline at end of file
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md b/.slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md
new file mode 100644
index 0000000000000000000000000000000000000000..b22302f022718637031a93f6221a78fb3bfbcb43
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md
@@ -0,0 +1,130 @@
+
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+ and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+ overall Diffusers community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+ advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+ address, without their explicit permission
+* Spamming issues or PRs with links to projects unrelated to this library
+* Other conduct which could reasonably be considered inappropriate in a
+ professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+feedback@huggingface.co.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior, harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.1, available at
+https://www.contributor-covenant.org/version/2/1/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md b/.slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md
new file mode 100644
index 0000000000000000000000000000000000000000..53de38ca21e31642bddee9c55c2cb8d398b5418d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md
@@ -0,0 +1 @@
+docs/source/en/conceptual/contribution.md
\ No newline at end of file
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/LICENSE b/.slim/clonedeps/repos/huggingface__diffusers/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..440eadcf5fffa9c81b9bfb121d3adc822862fb5e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/LICENSE
@@ -0,0 +1,201 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, Any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in b/.slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in
new file mode 100644
index 0000000000000000000000000000000000000000..ea8aa974629d3ae9e4284448a89bcc08729e6878
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in
@@ -0,0 +1,2 @@
+include LICENSE
+include src/diffusers/utils/model_card_template.md
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/Makefile b/.slim/clonedeps/repos/huggingface__diffusers/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..187f16495c4fe8fd9a31404de7bc9f3f7cfc9c4b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/Makefile
@@ -0,0 +1,120 @@
+.PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples codex claude clean-ai
+
+# make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!)
+export PYTHONPATH = src
+
+check_dirs := examples scripts src tests utils benchmarks
+
+modified_only_fixup:
+ $(eval modified_py_files := $(shell python utils/get_modified_files.py $(check_dirs)))
+ @if test -n "$(modified_py_files)"; then \
+ echo "Checking/fixing $(modified_py_files)"; \
+ ruff check $(modified_py_files) --fix; \
+ ruff format $(modified_py_files);\
+ else \
+ echo "No library .py files were modified"; \
+ fi
+
+# Update src/diffusers/dependency_versions_table.py
+
+deps_table_update:
+ @python setup.py deps_table_update
+
+deps_table_check_updated:
+ @md5sum src/diffusers/dependency_versions_table.py > md5sum.saved
+ @python setup.py deps_table_update
+ @md5sum -c --quiet md5sum.saved || (printf "\nError: the version dependency table is outdated.\nPlease run 'make fixup' or 'make style' and commit the changes.\n\n" && exit 1)
+ @rm md5sum.saved
+
+# autogenerating code
+
+autogenerate_code: deps_table_update
+
+# Check that the repo is in a good state
+
+repo-consistency:
+ python utils/check_dummies.py
+ python utils/check_repo.py
+ python utils/check_inits.py
+ python utils/check_forward_call_docstrings.py
+
+# this target runs checks on all files
+
+quality:
+ ruff check $(check_dirs) setup.py
+ ruff format --check $(check_dirs) setup.py
+ doc-builder style src/diffusers docs/source --max_len 119 --check_only
+ python utils/check_doc_toc.py
+
+# Format source code automatically and check is there are any problems left that need manual fixing
+
+extra_style_checks:
+ python utils/custom_init_isort.py
+ python utils/check_doc_toc.py --fix_and_overwrite
+
+# this target runs checks on all files and potentially modifies some of them
+
+style:
+ ruff check $(check_dirs) setup.py --fix
+ ruff format $(check_dirs) setup.py
+ doc-builder style src/diffusers docs/source --max_len 119
+ ${MAKE} autogenerate_code
+ ${MAKE} extra_style_checks
+
+# Super fast fix and check target that only works on relevant modified files since the branch was made
+
+fixup: modified_only_fixup extra_style_checks autogenerate_code repo-consistency
+
+# Make marked copies of snippets of codes conform to the original
+
+fix-copies:
+ python utils/check_copies.py --fix_and_overwrite
+ python utils/check_dummies.py --fix_and_overwrite
+
+# Auto docstrings in modular blocks
+modular-autodoctrings:
+ python utils/modular_auto_docstring.py
+
+# Verify forward() / __call__() arguments are documented in their docstrings
+check-forward-call-docstrings:
+ python utils/check_forward_call_docstrings.py
+
+# Run tests for the library
+
+test:
+ python -m pytest -n auto --dist=loadfile -s -v ./tests/
+
+# Run tests for examples
+
+test-examples:
+ python -m pytest -n auto --dist=loadfile -s -v ./examples/
+
+
+# Release stuff
+
+pre-release:
+ python utils/release.py
+
+pre-patch:
+ python utils/release.py --patch
+
+post-release:
+ python utils/release.py --post_release
+
+post-patch:
+ python utils/release.py --post_release --patch
+
+# AI agent symlinks
+
+codex:
+ mkdir -p .agents
+ rm -rf .agents/skills
+ ln -snf ../.ai/skills .agents/skills
+
+claude:
+ mkdir -p .claude
+ rm -rf .claude/skills
+ ln -snf ../.ai/skills .claude/skills
+
+clean-ai:
+ rm -rf .agents/skills .claude/skills
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md b/.slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md
new file mode 100644
index 0000000000000000000000000000000000000000..2fb658a8ca546c6c7d1755d95be0689b1a749622
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md
@@ -0,0 +1 @@
+docs/source/en/conceptual/philosophy.md
\ No newline at end of file
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/README.md b/.slim/clonedeps/repos/huggingface__diffusers/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..53c639264de1bf0c727d8c1ad1a096cbd3f97d9d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/README.md
@@ -0,0 +1,232 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+🤗 Diffusers is the go-to library for state-of-the-art pretrained diffusion models for generating images, audio, and even 3D structures of molecules. Whether you're looking for a simple inference solution or training your own diffusion models, 🤗 Diffusers is a modular toolbox that supports both. Our library is designed with a focus on [usability over performance](https://huggingface.co/docs/diffusers/conceptual/philosophy#usability-over-performance), [simple over easy](https://huggingface.co/docs/diffusers/conceptual/philosophy#simple-over-easy), and [customizability over abstractions](https://huggingface.co/docs/diffusers/conceptual/philosophy#tweakable-contributorfriendly-over-abstraction).
+
+🤗 Diffusers offers three core components:
+
+- State-of-the-art [diffusion pipelines](https://huggingface.co/docs/diffusers/api/pipelines/overview) that can be run in inference with just a few lines of code.
+- Interchangeable noise [schedulers](https://huggingface.co/docs/diffusers/api/schedulers/overview) for different diffusion speeds and output quality.
+- Pretrained [models](https://huggingface.co/docs/diffusers/api/models/overview) that can be used as building blocks, and combined with schedulers, for creating your own end-to-end diffusion systems.
+
+## Installation
+
+We recommend installing 🤗 Diffusers in a virtual environment from PyPI or Conda. For more details about installing [PyTorch](https://pytorch.org/get-started/locally/), please refer to their official documentation.
+
+### PyTorch
+
+With `pip` (official package):
+
+```bash
+pip install --upgrade diffusers[torch]
+```
+
+With `conda` (maintained by the community):
+
+```sh
+conda install -c conda-forge diffusers
+```
+
+### Apple Silicon (M1/M2) support
+
+Please refer to the [How to use Stable Diffusion in Apple Silicon](https://huggingface.co/docs/diffusers/optimization/mps) guide.
+
+## Quickstart
+
+Generating outputs is super easy with 🤗 Diffusers. To generate an image from text, use the `from_pretrained` method to load any pretrained diffusion model (browse the [Hub](https://huggingface.co/models?library=diffusers&sort=downloads) for 30,000+ checkpoints):
+
+```python
+from diffusers import DiffusionPipeline
+import torch
+
+pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16)
+pipeline.to("cuda")
+pipeline("An image of a squirrel in Picasso style").images[0]
+```
+
+You can also dig into the models and schedulers toolbox to build your own diffusion system:
+
+```python
+from diffusers import DDPMScheduler, UNet2DModel
+from PIL import Image
+import torch
+
+scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
+model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
+scheduler.set_timesteps(50)
+
+sample_size = model.config.sample_size
+noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
+input = noise
+
+for t in scheduler.timesteps:
+ with torch.no_grad():
+ noisy_residual = model(input, t).sample
+ prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
+ input = prev_noisy_sample
+
+image = (input / 2 + 0.5).clamp(0, 1)
+image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+image = Image.fromarray((image * 255).round().astype("uint8"))
+image
+```
+
+Check out the [Quickstart](https://huggingface.co/docs/diffusers/quicktour) to launch your diffusion journey today!
+
+## How to navigate the documentation
+
+| **Documentation** | **What can I learn?** |
+|---------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| [Tutorial](https://huggingface.co/docs/diffusers/tutorials/tutorial_overview) | A basic crash course for learning how to use the library's most important features like using models and schedulers to build your own diffusion system, and training your own diffusion model. |
+| [Loading](https://huggingface.co/docs/diffusers/using-diffusers/loading) | Guides for how to load and configure all the components (pipelines, models, and schedulers) of the library, as well as how to use different schedulers. |
+| [Pipelines for inference](https://huggingface.co/docs/diffusers/using-diffusers/overview_techniques) | Guides for how to use pipelines for different inference tasks, batched generation, controlling generated outputs and randomness, and how to contribute a pipeline to the library. |
+| [Optimization](https://huggingface.co/docs/diffusers/optimization/fp16) | Guides for how to optimize your diffusion model to run faster and consume less memory. |
+| [Training](https://huggingface.co/docs/diffusers/training/overview) | Guides for how to train a diffusion model for different tasks with different training techniques. |
+## Contribution
+
+We ❤️ contributions from the open-source community!
+If you want to contribute to this library, please check out our [Contribution guide](https://huggingface.co/docs/diffusers/main/en/conceptual/contribution).
+If you are using an AI agent, please point it at the project conventions in [`.ai/`](https://github.com/huggingface/diffusers/tree/main/.ai) first (run `make claude` or `make codex`) — see [Coding with AI agents](https://huggingface.co/docs/diffusers/main/en/conceptual/contribution#coding-with-ai-agents).
+You can look out for [issues](https://github.com/huggingface/diffusers/issues) you'd like to tackle to contribute to the library.
+- See [Good first issues](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22) for general opportunities to contribute
+- See [New model/pipeline](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+pipeline%2Fmodel%22) to contribute exciting new diffusion models / diffusion pipelines
+- See [New scheduler](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+scheduler%22)
+
+Also, say 👋 in our public Discord channel . We discuss the hottest trends about diffusion models, help each other with contributions, personal projects or just hang out ☕.
+
+
+## Popular Tasks & Pipelines
+
+
+
+## Popular libraries using 🧨 Diffusers
+
+- https://github.com/microsoft/TaskMatrix
+- https://github.com/invoke-ai/InvokeAI
+- https://github.com/InstantID/InstantID
+- https://github.com/apple/ml-stable-diffusion
+- https://github.com/Sanster/lama-cleaner
+- https://github.com/IDEA-Research/Grounded-Segment-Anything
+- https://github.com/ashawkey/stable-dreamfusion
+- https://github.com/deep-floyd/IF
+- https://github.com/bentoml/BentoML
+- https://github.com/bmaltais/kohya_ss
+- +14,000 other amazing GitHub repositories 💪
+
+Thank you for using us ❤️.
+
+## Credits
+
+This library concretizes previous work by many different authors and would not have been possible without their great research and implementations. We'd like to thank, in particular, the following implementations which have helped us in our development and without which the API could not have been as polished today:
+
+- @CompVis' latent diffusion models library, available [here](https://github.com/CompVis/latent-diffusion)
+- @hojonathanho original DDPM implementation, available [here](https://github.com/hojonathanho/diffusion) as well as the extremely useful translation into PyTorch by @pesser, available [here](https://github.com/pesser/pytorch_diffusion)
+- @ermongroup's DDIM implementation, available [here](https://github.com/ermongroup/ddim)
+- @yang-song's Score-VE and Score-VP implementations, available [here](https://github.com/yang-song/score_sde_pytorch)
+
+We also want to thank @heejkoo for the very helpful overview of papers, code and resources on diffusion models, available [here](https://github.com/heejkoo/Awesome-Diffusion-Models) as well as @crowsonkb and @rromb for useful discussions and insights.
+
+## Citation
+
+```bibtex
+@misc{von-platen-etal-2022-diffusers,
+ author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Dhruv Nair and Sayak Paul and William Berman and Yiyi Xu and Steven Liu and Thomas Wolf},
+ title = {Diffusers: State-of-the-art diffusion models},
+ year = {2022},
+ publisher = {GitHub},
+ journal = {GitHub repository},
+ howpublished = {\url{https://github.com/huggingface/diffusers}}
+}
+```
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/SECURITY.md b/.slim/clonedeps/repos/huggingface__diffusers/SECURITY.md
new file mode 100644
index 0000000000000000000000000000000000000000..6212e94050e3faf40f6dc431547d2dd43f10d3b8
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/SECURITY.md
@@ -0,0 +1,106 @@
+# Security Policy
+
+## Supported Versions
+
+We ship security fixes in the latest release. Please reproduce on the current released version before reporting — we do not backport fixes to older releases.
+
+## Reporting a Vulnerability
+
+Report privately — **do not open a public issue or PR for a suspected vulnerability.**
+
+- **Preferred:** GitHub private vulnerability reporting — the **"Report a vulnerability"** button under this repo's **Security** tab. This routes to the maintainers, keeps the report private until a fix is ready, and lets us issue a CVE through GitHub if warranted.
+- **Email:** security@huggingface.co
+
+We acknowledge valid, in-scope reports and will keep you updated on remediation. Please give us a reasonable window to fix before any public disclosure.
+
+## Recognition
+
+We do not offer a monetary bounty. For a valid, in-scope report we credit you on the published GitHub Security Advisory and name you as the reporter in the associated CVE. Let us know how you'd like to be credited (name or handle).
+
+## What your report must include
+
+We receive a high volume of reports. To be triaged, a report **must** follow the structure below. Copy this block into your submission and fill in every field. Reports missing the version, the proof of concept, or the impact are returned as incomplete and are not investigated until provided.
+
+```markdown
+### Summary
+One sentence: what the vulnerability is and where.
+
+### Affected version / commit
+Exact released version or commit SHA you reproduced on (e.g. v4.57.0 / a1b2c3d).
+Not "latest" or "main".
+
+### Affected component
+The public API, module, or entry point involved (e.g. `DiffusionPipeline.from_pretrained`, `AutoModel.from_pretrained`).
+
+### Vulnerability class
+Type and CWE if known (e.g. deserialization / CWE-502, path traversal / CWE-22).
+
+### Attack vector & preconditions
+- How is the vulnerable code reached? (which API call / input / config)
+- Who is the attacker and what do they control?
+- What must be true for the attack to work? (auth, a user action, a non-default
+ setting, a malicious file being loaded, etc.)
+
+### Proof of concept
+A minimal, self-contained script or step sequence that runs on a clean install
+of the version above. Include:
+- the exact commands / code to run,
+- any input files needed (attach them, or give a script that generates them),
+- the **expected** behavior vs. the **actual** behavior you observed.
+A snippet showing that a function *exists* or *could* be misused is not a PoC.
+
+### Impact
+What an attacker gains in a realistic deployment. "Could theoretically…"
+without a working chain is not an impact.
+
+### Scope
+Which trust boundary (see below) does this cross? If your finding touches
+anything in the "Out of scope" list, name which item and explain why it is
+nonetheless a violation of a guarantee we make.
+
+### Suggested severity (optional)
+We assign the final severity. Include a CVSS v3.1 vector only if you have one.
+
+### Suggested fix (optional)
+```
+
+> The bar is a **reproducible PoC against a supported version, with a concrete impact that crosses a trust boundary we actually defend** (see scope below). Reports that are theoretical, auto-generated by a scanner or LLM, or that restate documented behavior will be closed without detailed review.
+
+## Threat model & trust boundaries
+
+Understanding these saves everyone time — most of what we close as not-a-vulnerability falls inside one of them.
+
+**Loading artifacts you did not create is a trust decision.** Models, tokenizers, and configuration files can carry code or instructions that run when loaded. That a maliciously crafted artifact can execute code, read local files, or otherwise act when you load it is the **documented risk of loading untrusted content** — not a vulnerability in this library. Protect yourself by:
+- preferring the `safetensors` format over pickle-based formats (`use_safetensors=True`),
+- pinning a specific, reviewed `revision`,
+- setting `trust_remote_code=True` only for repositories whose modeling files you have inspected.
+
+We *will* treat as a vulnerability anything that breaks one of these protections — e.g. code executing despite `safetensors`-only loading, or a pinned revision being bypassed.
+
+**Conversion, CLI, and developer utilities are operator tools.** Scripts you run yourself against inputs you chose (format converters, training/utility CLIs, dev helpers) are not a library API attack surface. Behaviors such as a conversion script deserializing a checkpoint you pointed it at are within the operator's trust.
+
+## In scope
+
+We treat as vulnerabilities issues in the **published package code** — the library's own API surface — that an attacker can trigger without the victim having opted into a documented risk. For example:
+
+- code execution, memory corruption, or file access reachable through a normal API call on input that is **not** an untrusted model/artifact the user chose to load;
+- a control we advertise being bypassed (e.g. code running despite `safetensors`-only loading, or a pinned revision being ignored);
+- exposure or mishandling of credentials, tokens, or another user's data by the library;
+- CI/CD or supply-chain issues in this repository.
+
+## Out of scope
+
+The following are **not** treated as vulnerabilities in transformers. If your finding touches one of these, the report must explain why it is nonetheless a violation of a guarantee we make — otherwise it will be closed.
+
+- Issues that require loading an untrusted artifact and amount to the documented load-time risk above (code execution / file access on load of a malicious model, config, or pickle).
+- Executable code or class references declared in a model's own configuration or metadata and resolved when that model is loaded — loading a repository that sets these is the same documented load-time trust decision, not a library flaw. (A control we advertise being bypassed — for example remote or config-declared code running without the `trust_remote_code` gate — remains in scope.)
+- Findings in documentation, tests, or other non-packaged reference material.
+- Local denial-of-service from feeding pathological input to a function on your own machine (high memory, slow parse, panic), absent a multi-tenant or remote-service impact.
+- Model behavior: jailbreaks, alignment failures, prompt injection, or harmful generations. Model weights are authored by their uploaders; report these to the model owner. This includes indirect prompt injection — transformers is a modeling library with no agent- or tool-execution code path, so there is no orchestration layer here to inject into.
+- Vulnerabilities in third-party dependencies we do not vendor — report upstream (we'll bump once fixed).
+- Theoretical issues without a working proof of concept, and reports auto-generated from scanners or LLMs without a verified, reproducible chain.
+- Best-practice or hardening suggestions with no demonstrated impact — missing email-authentication or transport records (MTA-STS, TLS-RPT, DMARC/SPF tuning), missing HTTP security headers, TLS configuration preferences, and similar scanner or config-checker output presented without a working exploit chain.
+
+## Safe harbor
+
+Good-faith research that respects these guidelines, avoids privacy violations and service disruption, and gives us a reasonable disclosure window will not be pursued by us. Do not access data that isn't yours and do not run tests against Hugging Face production infrastructure.
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/_typos.toml b/.slim/clonedeps/repos/huggingface__diffusers/_typos.toml
new file mode 100644
index 0000000000000000000000000000000000000000..9c373e6722607ad5134915cf1ec15379bda4dc7b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/_typos.toml
@@ -0,0 +1,13 @@
+# Files for typos
+# Instruction: https://github.com/marketplace/actions/typos-action#getting-started
+
+[default.extend-identifiers]
+
+[default.extend-words]
+NIN="NIN" # NIN is used in scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py
+nd="np" # nd may be np (numpy)
+parms="parms" # parms is used in scripts/convert_original_stable_diffusion_to_diffusers.py
+
+
+[files]
+extend-exclude = ["_typos.toml"]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/pyproject.toml b/.slim/clonedeps/repos/huggingface__diffusers/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..4fd1ad75360c19b884cb92d1d5b93cc1ab578aab
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/pyproject.toml
@@ -0,0 +1,32 @@
+[tool.ruff]
+line-length = 119
+extend-exclude = [
+ "src/diffusers/pipelines/flux2/system_messages.py",
+]
+
+[tool.ruff.lint]
+# Never enforce `E501` (line length violations).
+ignore = ["C901", "E501", "E721", "E741", "F402", "F823"]
+select = ["C", "E", "F", "I", "W"]
+
+# Ignore import violations in all `__init__.py` files.
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+"src/diffusers/utils/dummy_*.py" = ["F401"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+known-first-party = ["diffusers"]
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/setup.py b/.slim/clonedeps/repos/huggingface__diffusers/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..787c4dd0ceac82b73e34c4328f0b9427d172638f
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/setup.py
@@ -0,0 +1,305 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Simple check list from AllenNLP repo: https://github.com/allenai/allennlp/blob/main/setup.py
+
+To create the package for PyPI.
+
+1. Run `make pre-release` (or `make pre-patch` for a patch release) then run `make fix-copies` to fix the index of the
+ documentation.
+
+ If releasing on a special branch, copy the updated README.md on the main branch for the commit you will make
+ for the post-release and run `make fix-copies` on the main branch as well.
+
+2. Unpin specific versions from setup.py that use a git install.
+
+3. Checkout the release branch (v-release, for example v4.19-release), and commit these changes with the
+ message: "Release: " and push.
+
+4. Manually trigger the "Nightly and release tests on main/release branch" workflow from the release branch. Wait for
+ the tests to complete. We can safely ignore the known test failures.
+
+5. Wait for the tests on main to be completed and be green (otherwise revert and fix bugs).
+
+6. Add a tag in git to mark the release: "git tag v -m 'Adds tag v for PyPI'"
+ Push the tag to git: git push --tags origin v-release
+
+7. Build both the sources and the wheel. Do not change anything in setup.py between
+ creating the wheel and the source distribution (obviously).
+
+ For the wheel, run: "python setup.py bdist_wheel" in the top level directory
+ (This will build a wheel for the Python version you use to build it).
+
+ For the sources, run: "python setup.py sdist"
+ You should now have a /dist directory with both .whl and .tar.gz source versions.
+
+ Long story cut short, you need to run both before you can upload the distribution to the
+ test PyPI and the actual PyPI servers:
+
+ python setup.py bdist_wheel && python setup.py sdist
+
+8. Check that everything looks correct by uploading the package to the PyPI test server:
+
+ twine upload dist/* -r pypitest
+ (pypi suggests using twine as other methods upload files via plaintext.)
+ You may have to specify the repository url, use the following command then:
+ twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
+
+ Check that you can install it in a virtualenv by running:
+ pip install -i https://testpypi.python.org/pypi diffusers
+
+ If you are testing from a Colab Notebook, for instance, then do:
+ pip install diffusers && pip uninstall diffusers
+ pip install -i https://testpypi.python.org/pypi diffusers
+
+ Check you can run the following commands:
+ python -c "from diffusers import __version__; print(__version__)"
+ python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('fusing/unet-ldm-dummy-update'); pipe()"
+ python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe', safety_checker=None); pipe('ah suh du')"
+ python -c "from diffusers import *"
+
+9. Upload the final version to the actual PyPI:
+ twine upload dist/* -r pypi
+
+10. Prepare the release notes and publish them on GitHub once everything is looking hunky-dory. You can use the following
+ Space to fetch all the commits applicable for the release: https://huggingface.co/spacmes/sayakpaul/auto-release-notes-diffusers.
+ It automatically fetches the correct tag and branch but also provides the option to configure them.
+ `tag` should be the previous release tag (v0.26.1, for example), and `branch` should be
+ the latest release branch (v0.27.0-release, for example). It denotes all commits that have happened on branch
+ v0.27.0-release after the tag v0.26.1 was created.
+
+11. Run `make post-release` (or, for a patch release, `make post-patch`). If you were on a branch for the release,
+ you need to go back to main before executing this.
+"""
+
+import re
+import sys
+
+from setuptools import Command, find_packages, setup
+
+
+# IMPORTANT:
+# 1. all dependencies should be listed here with their version requirements if any
+# 2. once modified, run: `make deps_table_update` to update src/diffusers/dependency_versions_table.py
+_deps = [
+ "Pillow", # keep the PIL.Image.Resampling deprecation away
+ "accelerate>=0.31.0",
+ "datasets",
+ "filelock",
+ "ftfy",
+ "hf-doc-builder>=0.3.0",
+ "httpx<1.0.0",
+ "huggingface-hub>=1.23.0,<2.0",
+ "requests-mock==1.10.0",
+ "importlib_metadata",
+ "invisible-watermark>=0.2.0",
+ "isort>=5.5.4",
+ "Jinja2",
+ "torchsde",
+ "note_seq",
+ "librosa",
+ # librosa pulls in numba/llvmlite; pin to versions that support Python 3.10+
+ "llvmlite>=0.40.0",
+ "numba>=0.57.0",
+ "numpy",
+ "parameterized",
+ "peft>=0.17.0",
+ "protobuf>=3.20.3,<4",
+ "pytest",
+ "pytest-timeout",
+ "pytest-xdist",
+ "python>=3.10.0",
+ "ruff==0.9.10",
+ "safetensors>=0.8.0",
+ "sentencepiece>=0.1.91,!=0.1.92",
+ "GitPython<3.1.19",
+ "scipy",
+ "onnx",
+ "optimum_quanto>=0.2.6",
+ "gguf>=0.10.0",
+ "auto-round>=0.13.0",
+ "torchao>=0.7.0",
+ "bitsandbytes>=0.43.3",
+ "nvidia_modelopt[hf]>=0.33.1",
+ "sdnq>=0.2.2",
+ "regex!=2019.12.17",
+ "requests",
+ "tensorboard",
+ "tiktoken>=0.7.0",
+ "torch>=2.6",
+ "torchvision",
+ "transformers>=4.41.2",
+ "urllib3<=2.0.0",
+ "black",
+ "phonemizer",
+ "opencv-python",
+ "timm",
+ "flashpack",
+]
+
+# this is a lookup table with items like:
+#
+# tokenizers: "huggingface-hub==0.8.0"
+# packaging: "packaging"
+#
+# some of the values are versioned whereas others aren't.
+deps = {b: a for a, b in (re.findall(r"^(([^!=<>~]+)(?:[!=<>~].*)?$)", x)[0] for x in _deps)}
+
+# since we save this data in src/diffusers/dependency_versions_table.py it can be easily accessed from
+# anywhere. If you need to quickly access the data from this table in a shell, you can do so easily with:
+#
+# python -c 'import sys; from diffusers.dependency_versions_table import deps; \
+# print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets
+#
+# Just pass the desired package names to that script as it's shown with 2 packages above.
+#
+# If diffusers is not yet installed and the work is done from the cloned repo remember to add `PYTHONPATH=src` to the script above
+#
+# You can then feed this for example to `pip`:
+#
+# pip install -U $(python -c 'import sys; from diffusers.dependency_versions_table import deps; \
+# print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets)
+#
+
+
+def deps_list(*pkgs):
+ return [deps[pkg] for pkg in pkgs]
+
+
+class DepsTableUpdateCommand(Command):
+ """
+ A custom command that updates the dependency table.
+ usage: python setup.py deps_table_update
+ """
+
+ description = "build runtime dependency table"
+ user_options = [
+ # format: (long option, short option, description).
+ (
+ "dep-table-update",
+ None,
+ "updates src/diffusers/dependency_versions_table.py",
+ ),
+ ]
+
+ def initialize_options(self):
+ pass
+
+ def finalize_options(self):
+ pass
+
+ def run(self):
+ entries = "\n".join([f' "{k}": "{v}",' for k, v in deps.items()])
+ content = [
+ "# THIS FILE HAS BEEN AUTOGENERATED. To update:",
+ "# 1. modify the `_deps` dict in setup.py",
+ "# 2. run `make deps_table_update`",
+ "deps = {",
+ entries,
+ "}",
+ "",
+ ]
+ target = "src/diffusers/dependency_versions_table.py"
+ print(f"updating {target}")
+ with open(target, "w", encoding="utf-8", newline="\n") as f:
+ f.write("\n".join(content))
+
+
+extras = {}
+extras["quality"] = deps_list("urllib3", "isort", "ruff", "hf-doc-builder")
+extras["docs"] = deps_list("hf-doc-builder")
+extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft", "timm")
+extras["test"] = deps_list(
+ "ftfy",
+ "GitPython",
+ "datasets",
+ "Jinja2",
+ "invisible-watermark",
+ "librosa",
+ "llvmlite",
+ "numba",
+ "parameterized",
+ "protobuf",
+ "pytest",
+ "pytest-timeout",
+ "pytest-xdist",
+ "requests-mock",
+ "safetensors",
+ "sentencepiece",
+ "scipy",
+ "tiktoken",
+ "torchsde",
+ "torchvision",
+ "transformers",
+ "phonemizer",
+)
+extras["torch"] = deps_list("torch", "accelerate")
+
+extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate")
+extras["gguf"] = deps_list("gguf", "accelerate")
+extras["optimum_quanto"] = deps_list("optimum_quanto", "accelerate")
+extras["torchao"] = deps_list("torchao", "accelerate")
+extras["nvidia_modelopt"] = deps_list("nvidia_modelopt[hf]")
+extras["sdnq"] = deps_list("sdnq")
+extras["flashpack"] = deps_list("flashpack")
+
+extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"]
+
+install_requires = [
+ deps["importlib_metadata"],
+ deps["filelock"],
+ deps["httpx"],
+ deps["huggingface-hub"],
+ deps["numpy"],
+ deps["regex"],
+ deps["requests"],
+ deps["safetensors"],
+ deps["Pillow"],
+]
+
+version_range_max = max(sys.version_info[1], 10) + 1
+
+setup(
+ name="diffusers",
+ version="0.40.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
+ description="State-of-the-art diffusion in PyTorch.",
+ long_description=open("README.md", "r", encoding="utf-8").read(),
+ long_description_content_type="text/markdown",
+ keywords="deep learning diffusion pytorch stable diffusion audioldm",
+ license="Apache 2.0 License",
+ author="The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)",
+ author_email="diffusers@huggingface.co",
+ url="https://github.com/huggingface/diffusers",
+ package_dir={"": "src"},
+ packages=find_packages("src"),
+ package_data={"diffusers": ["py.typed"]},
+ include_package_data=True,
+ python_requires=">=3.10.0",
+ install_requires=list(install_requires),
+ extras_require=extras,
+ entry_points={"console_scripts": ["diffusers-cli=diffusers.commands.diffusers_cli:main"]},
+ classifiers=[
+ "Development Status :: 5 - Production/Stable",
+ "Intended Audience :: Developers",
+ "Intended Audience :: Education",
+ "Intended Audience :: Science/Research",
+ "License :: OSI Approved :: Apache Software License",
+ "Operating System :: OS Independent",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Programming Language :: Python :: 3",
+ ]
+ + [f"Programming Language :: Python :: 3.{i}" for i in range(8, version_range_max)],
+ cmdclass={"deps_table_update": DepsTableUpdateCommand},
+)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d68cc9329314e8d01b0edd8fcd7f378b722c9e60
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py
@@ -0,0 +1,1746 @@
+__version__ = "0.40.0.dev0"
+
+from typing import TYPE_CHECKING
+
+from .utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ is_accelerate_available,
+ is_auto_round_available,
+ is_bitsandbytes_available,
+ is_gguf_available,
+ is_librosa_available,
+ is_note_seq_available,
+ is_nvidia_modelopt_available,
+ is_onnx_available,
+ is_opencv_available,
+ is_optimum_quanto_available,
+ is_scipy_available,
+ is_sdnq_available,
+ is_sentencepiece_available,
+ is_torch_available,
+ is_torchao_available,
+ is_torchsde_available,
+ is_transformers_available,
+ is_transformers_version,
+)
+
+
+# Lazy Import based on
+# https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
+
+# When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names,
+# and is used to defer the actual importing for when the objects are requested.
+# This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends).
+
+_import_structure = {
+ "configuration_utils": ["ConfigMixin"],
+ "guiders": [],
+ "hooks": [],
+ "loaders": ["FromOriginalModelMixin"],
+ "models": [],
+ "modular_pipelines": [],
+ "pipelines": [],
+ "quantizers.pipe_quant_config": ["PipelineQuantizationConfig"],
+ "quantizers.quantization_config": [],
+ "schedulers": [],
+ "utils": [
+ "OptionalDependencyNotAvailable",
+ "is_inflect_available",
+ "is_invisible_watermark_available",
+ "is_librosa_available",
+ "is_note_seq_available",
+ "is_onnx_available",
+ "is_scipy_available",
+ "is_torch_available",
+ "is_torchsde_available",
+ "is_transformers_available",
+ "is_transformers_version",
+ "is_unidecode_available",
+ "logging",
+ ],
+}
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_bitsandbytes_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_bitsandbytes_objects
+
+ _import_structure["utils.dummy_bitsandbytes_objects"] = [
+ name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig")
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_gguf_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_gguf_objects
+
+ _import_structure["utils.dummy_gguf_objects"] = [
+ name for name in dir(dummy_gguf_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_torchao_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torchao_objects
+
+ _import_structure["utils.dummy_torchao_objects"] = [
+ name for name in dir(dummy_torchao_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("TorchAoConfig")
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_optimum_quanto_objects
+
+ _import_structure["utils.dummy_optimum_quanto_objects"] = [
+ name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("QuantoConfig")
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_nvidia_modelopt_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_nvidia_modelopt_objects
+
+ _import_structure["utils.dummy_nvidia_modelopt_objects"] = [
+ name for name in dir(dummy_nvidia_modelopt_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig")
+
+try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_nunchaku_lite_objects
+
+ _import_structure["utils.dummy_nunchaku_lite_objects"] = [
+ name for name in dir(dummy_nunchaku_lite_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig")
+
+try:
+ if not is_auto_round_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_auto_round_objects
+
+ _import_structure["utils.dummy_auto_round_objects"] = [
+ name for name in dir(dummy_auto_round_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("AutoRoundConfig")
+
+try:
+ if not is_torch_available() and not is_accelerate_available() and not is_sdnq_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_sdnq_objects
+
+ _import_structure["utils.dummy_sdnq_objects"] = [
+ name for name in dir(dummy_sdnq_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("SDNQConfig")
+
+try:
+ if not is_onnx_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_onnx_objects # noqa F403
+
+ _import_structure["utils.dummy_onnx_objects"] = [
+ name for name in dir(dummy_onnx_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["pipelines"].extend(["OnnxRuntimeModel"])
+
+try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_pt_objects # noqa F403
+
+ _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
+
+else:
+ _import_structure["guiders"].extend(
+ [
+ "AdaptiveProjectedGuidance",
+ "AdaptiveProjectedMixGuidance",
+ "AutoGuidance",
+ "BaseGuidance",
+ "ClassifierFreeGuidance",
+ "ClassifierFreeZeroStarGuidance",
+ "FrequencyDecoupledGuidance",
+ "PerturbedAttentionGuidance",
+ "SkipLayerGuidance",
+ "SmoothedEnergyGuidance",
+ "TangentialClassifierFreeGuidance",
+ ]
+ )
+ _import_structure["hooks"].extend(
+ [
+ "FasterCacheConfig",
+ "FirstBlockCacheConfig",
+ "HookRegistry",
+ "LayerSkipConfig",
+ "MagCacheConfig",
+ "PyramidAttentionBroadcastConfig",
+ "SmoothedEnergyGuidanceConfig",
+ "TaylorSeerCacheConfig",
+ "TextKVCacheConfig",
+ "apply_faster_cache",
+ "apply_first_block_cache",
+ "apply_layer_skip",
+ "apply_mag_cache",
+ "apply_pyramid_attention_broadcast",
+ "apply_taylorseer_cache",
+ "apply_text_kv_cache",
+ ]
+ )
+ _import_structure["image_processor"] = [
+ "InpaintProcessor",
+ "IPAdapterMaskProcessor",
+ "PixArtImageProcessor",
+ "VaeImageProcessor",
+ "VaeImageProcessorLDM3D",
+ ]
+ _import_structure["models"].extend(
+ [
+ "AceStepTransformer1DModel",
+ "AllegroTransformer3DModel",
+ "AnimaTextConditioner",
+ "AnyFlowFARTransformer3DModel",
+ "AnyFlowTransformer3DModel",
+ "AsymmetricAutoencoderKL",
+ "AttentionBackendName",
+ "AuraFlowTransformer2DModel",
+ "AutoencoderDC",
+ "AutoencoderKL",
+ "AutoencoderKLAllegro",
+ "AutoencoderKLCogVideoX",
+ "AutoencoderKLCosmos",
+ "AutoencoderKLFlux2",
+ "AutoencoderKLHunyuanImage",
+ "AutoencoderKLHunyuanImageRefiner",
+ "AutoencoderKLHunyuanVideo",
+ "AutoencoderKLHunyuanVideo15",
+ "AutoencoderKLKVAE",
+ "AutoencoderKLKVAEVideo",
+ "AutoencoderKLLTX2Audio",
+ "AutoencoderKLLTX2Video",
+ "AutoencoderKLLTXVideo",
+ "AutoencoderKLMagvit",
+ "AutoencoderKLMiniMaxH3",
+ "AutoencoderKLMiniMaxH3Audio",
+ "AutoencoderKLMochi",
+ "AutoencoderKLQwenImage",
+ "AutoencoderKLTemporalDecoder",
+ "AutoencoderKLWan",
+ "AutoencoderOobleck",
+ "AutoencoderRAE",
+ "AutoencoderTiny",
+ "AutoencoderVidTok",
+ "AutoModel",
+ "BriaFiboTransformer2DModel",
+ "BriaTransformer2DModel",
+ "CacheMixin",
+ "ChromaTransformer2DModel",
+ "ChronoEditTransformer3DModel",
+ "CogVideoXTransformer3DModel",
+ "CogView3PlusTransformer2DModel",
+ "CogView4Transformer2DModel",
+ "ConsisIDTransformer3DModel",
+ "ConsistencyDecoderVAE",
+ "ContextParallelConfig",
+ "ControlNetModel",
+ "ControlNetUnionModel",
+ "ControlNetXSAdapter",
+ "Cosmos3AVAEAudioTokenizer",
+ "Cosmos3OmniTransformer",
+ "CosmosControlNetModel",
+ "CosmosTransformer3DModel",
+ "DiTTransformer2DModel",
+ "DreamLiteTransformer2DModel",
+ "DreamLiteUNetModel",
+ "EasyAnimateTransformer3DModel",
+ "ErnieImageTransformer2DModel",
+ "Flux2Transformer2DModel",
+ "FluxControlNetModel",
+ "FluxMultiControlNetModel",
+ "FluxTransformer2DModel",
+ "GlmImageTransformer2DModel",
+ "HeliosTransformer3DModel",
+ "HiDreamImageTransformer2DModel",
+ "HunyuanDiT2DControlNetModel",
+ "HunyuanDiT2DModel",
+ "HunyuanDiT2DMultiControlNetModel",
+ "HunyuanImageTransformer2DModel",
+ "HunyuanVideo15Transformer3DModel",
+ "HunyuanVideoFramepackTransformer3DModel",
+ "HunyuanVideoTransformer3DModel",
+ "I2VGenXLUNet",
+ "Ideogram4Transformer2DModel",
+ "JoyImageEditPlusTransformer3DModel",
+ "JoyImageEditTransformer3DModel",
+ "Kandinsky3UNet",
+ "Kandinsky5Transformer3DModel",
+ "Krea2Transformer2DModel",
+ "LatteTransformer3DModel",
+ "LongCatAudioDiTTransformer",
+ "LongCatAudioDiTVae",
+ "LongCatImageTransformer2DModel",
+ "LTX2VideoTransformer3DModel",
+ "LTXVideoTransformer3DModel",
+ "Lumina2Transformer2DModel",
+ "LuminaNextDiT2DModel",
+ "MiniMaxH3Transformer3DModel",
+ "MochiTransformer3DModel",
+ "ModelMixin",
+ "MotifVideoTransformer3DModel",
+ "MotionAdapter",
+ "MultiAdapter",
+ "MultiControlNetModel",
+ "NucleusMoEImageTransformer2DModel",
+ "OmniGenTransformer2DModel",
+ "OvisImageTransformer2DModel",
+ "ParallelConfig",
+ "PixArtTransformer2DModel",
+ "PriorTransformer",
+ "PRXTransformer2DModel",
+ "QwenImageControlNetModel",
+ "QwenImageMultiControlNetModel",
+ "QwenImageTransformer2DModel",
+ "SanaControlNetModel",
+ "SanaTransformer2DModel",
+ "SanaVideoTransformer3DModel",
+ "SD3ControlNetModel",
+ "SD3MultiControlNetModel",
+ "SD3Transformer2DModel",
+ "SkyReelsV2Transformer3DModel",
+ "SparseControlNetModel",
+ "StableAudioDiTModel",
+ "StableCascadeUNet",
+ "T2IAdapter",
+ "T5FilmDecoder",
+ "Transformer2DModel",
+ "TransformerTemporalModel",
+ "UNet1DModel",
+ "UNet2DConditionModel",
+ "UNet2DModel",
+ "UNet3DConditionModel",
+ "UNetControlNetXSModel",
+ "UNetMotionModel",
+ "UNetSpatioTemporalConditionModel",
+ "UVit2DModel",
+ "VQModel",
+ "WanAnimateTransformer3DModel",
+ "WanTransformer3DModel",
+ "WanVACETransformer3DModel",
+ "ZImageControlNetModel",
+ "ZImageTransformer2DModel",
+ "attention_backend",
+ ]
+ )
+ _import_structure["modular_pipelines"].extend(
+ [
+ "AutoPipelineBlocks",
+ "ComponentsManager",
+ "ComponentSpec",
+ "ConditionalPipelineBlocks",
+ "ConfigSpec",
+ "InputParam",
+ "LoopSequentialPipelineBlocks",
+ "ModularPipeline",
+ "ModularPipelineBlocks",
+ "OutputParam",
+ "SequentialPipelineBlocks",
+ ]
+ )
+ _import_structure["optimization"] = [
+ "get_constant_schedule",
+ "get_constant_schedule_with_warmup",
+ "get_cosine_schedule_with_warmup",
+ "get_cosine_with_hard_restarts_schedule_with_warmup",
+ "get_linear_schedule_with_warmup",
+ "get_polynomial_decay_schedule_with_warmup",
+ "get_scheduler",
+ ]
+ _import_structure["pipelines"].extend(
+ [
+ "AudioPipelineOutput",
+ "AutoPipelineForImage2Image",
+ "AutoPipelineForInpainting",
+ "AutoPipelineForText2Audio",
+ "AutoPipelineForText2Image",
+ "ConsistencyModelPipeline",
+ "DanceDiffusionPipeline",
+ "DDIMPipeline",
+ "DDPMPipeline",
+ "DiffusionPipeline",
+ "DiTPipeline",
+ "ImagePipelineOutput",
+ "KarrasVePipeline",
+ "LDMPipeline",
+ "LDMSuperResolutionPipeline",
+ "PNDMPipeline",
+ "RePaintPipeline",
+ "ScoreSdeVePipeline",
+ "StableDiffusionMixin",
+ ]
+ )
+ _import_structure["quantizers"] = ["DiffusersQuantizer"]
+ _import_structure["schedulers"].extend(
+ [
+ "AmusedScheduler",
+ "BlockRefinementScheduler",
+ "BlockRefinementSchedulerOutput",
+ "CMStochasticIterativeScheduler",
+ "CogVideoXDDIMScheduler",
+ "CogVideoXDPMScheduler",
+ "DDIMInverseScheduler",
+ "DDIMParallelScheduler",
+ "DDIMScheduler",
+ "DDPMParallelScheduler",
+ "DDPMScheduler",
+ "DDPMWuerstchenScheduler",
+ "DEISMultistepScheduler",
+ "DiscreteDDIMScheduler",
+ "DiscreteDDIMSchedulerOutput",
+ "DPMSolverMultistepInverseScheduler",
+ "DPMSolverMultistepScheduler",
+ "DPMSolverSinglestepScheduler",
+ "EDMDPMSolverMultistepScheduler",
+ "EDMEulerScheduler",
+ "EntropyBoundScheduler",
+ "EntropyBoundSchedulerOutput",
+ "EulerAncestralDiscreteScheduler",
+ "EulerDiscreteScheduler",
+ "FlowMapEulerDiscreteScheduler",
+ "FlowMatchEulerDiscreteScheduler",
+ "FlowMatchHeunDiscreteScheduler",
+ "FlowMatchLCMScheduler",
+ "HeliosDMDScheduler",
+ "HeliosScheduler",
+ "HeunDiscreteScheduler",
+ "IPNDMScheduler",
+ "KarrasVeScheduler",
+ "KDPM2AncestralDiscreteScheduler",
+ "KDPM2DiscreteScheduler",
+ "LCMScheduler",
+ "LTXEulerAncestralRFScheduler",
+ "MiniMaxH3Scheduler",
+ "PNDMScheduler",
+ "RePaintScheduler",
+ "SASolverScheduler",
+ "SchedulerMixin",
+ "SCMScheduler",
+ "ScoreSdeVeScheduler",
+ "TCDScheduler",
+ "UnCLIPScheduler",
+ "UniPCMultistepScheduler",
+ "VQDiffusionScheduler",
+ ]
+ )
+ _import_structure["training_utils"] = ["EMAModel"]
+ _import_structure["video_processor"] = ["VideoProcessor"]
+
+try:
+ if not (is_torch_available() and is_scipy_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_scipy_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_scipy_objects"] = [
+ name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["schedulers"].extend(["LMSDiscreteScheduler"])
+
+try:
+ if not (is_torch_available() and is_torchsde_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_torchsde_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_torchsde_objects"] = [
+ name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"])
+
+try:
+ if not (is_torch_available() and is_transformers_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_transformers_objects"] = [
+ name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["modular_pipelines"].extend(
+ [
+ "AnimaAutoBlocks",
+ "AnimaModularPipeline",
+ "Cosmos3DistilledBlocks",
+ "Cosmos3DistilledModularPipeline",
+ "Cosmos3OmniBlocks",
+ "Cosmos3OmniModularPipeline",
+ "ErnieImageAutoBlocks",
+ "ErnieImageModularPipeline",
+ "Flux2AutoBlocks",
+ "Flux2KleinAutoBlocks",
+ "Flux2KleinBaseAutoBlocks",
+ "Flux2KleinBaseModularPipeline",
+ "Flux2KleinModularPipeline",
+ "Flux2ModularPipeline",
+ "FluxAutoBlocks",
+ "FluxKontextAutoBlocks",
+ "FluxKontextModularPipeline",
+ "FluxModularPipeline",
+ "HeliosAutoBlocks",
+ "HeliosModularPipeline",
+ "HeliosPyramidAutoBlocks",
+ "HeliosPyramidDistilledAutoBlocks",
+ "HeliosPyramidDistilledModularPipeline",
+ "HeliosPyramidModularPipeline",
+ "HunyuanVideo15AutoBlocks",
+ "HunyuanVideo15ModularPipeline",
+ "Ideogram4AutoBlocks",
+ "Ideogram4ModularPipeline",
+ "Krea2AutoBlocks",
+ "Krea2ModularPipeline",
+ "Krea2TurboAutoBlocks",
+ "Krea2TurboModularPipeline",
+ "LTXAutoBlocks",
+ "LTXModularPipeline",
+ "MiniMaxH3Blocks",
+ "MiniMaxH3ModularPipeline",
+ "QwenImageAutoBlocks",
+ "QwenImageEditAutoBlocks",
+ "QwenImageEditModularPipeline",
+ "QwenImageEditPlusAutoBlocks",
+ "QwenImageEditPlusModularPipeline",
+ "QwenImageLayeredAutoBlocks",
+ "QwenImageLayeredModularPipeline",
+ "QwenImageModularPipeline",
+ "StableDiffusion3AutoBlocks",
+ "StableDiffusion3ModularPipeline",
+ "StableDiffusionXLAutoBlocks",
+ "StableDiffusionXLModularPipeline",
+ "Wan22Blocks",
+ "Wan22Image2VideoBlocks",
+ "Wan22Image2VideoModularPipeline",
+ "Wan22ModularPipeline",
+ "WanBlocks",
+ "WanImage2VideoAutoBlocks",
+ "WanImage2VideoModularPipeline",
+ "WanModularPipeline",
+ "ZImageAutoBlocks",
+ "ZImageModularPipeline",
+ ]
+ )
+ _import_structure["pipelines"].extend(
+ [
+ "AceStepAudioTokenDetokenizer",
+ "AceStepAudioTokenizer",
+ "AceStepConditionEncoder",
+ "AceStepPipeline",
+ "AllegroPipeline",
+ "AltDiffusionImg2ImgPipeline",
+ "AltDiffusionPipeline",
+ "AmusedImg2ImgPipeline",
+ "AmusedInpaintPipeline",
+ "AmusedPipeline",
+ "AnimateDiffControlNetPipeline",
+ "AnimateDiffPAGPipeline",
+ "AnimateDiffPipeline",
+ "AnimateDiffSDXLPipeline",
+ "AnimateDiffSparseControlNetPipeline",
+ "AnimateDiffVideoToVideoControlNetPipeline",
+ "AnimateDiffVideoToVideoPipeline",
+ "AnyFlowFARPipeline",
+ "AnyFlowPipeline",
+ "AudioLDM2Pipeline",
+ "AudioLDM2ProjectionModel",
+ "AudioLDM2UNet2DConditionModel",
+ "AudioLDMPipeline",
+ "AuraFlowPipeline",
+ "BlipDiffusionControlNetPipeline",
+ "BlipDiffusionPipeline",
+ "BriaFiboEditPipeline",
+ "BriaFiboPipeline",
+ "BriaPipeline",
+ "ChromaImg2ImgPipeline",
+ "ChromaInpaintPipeline",
+ "ChromaPipeline",
+ "ChronoEditPipeline",
+ "CLIPImageProjection",
+ "CogVideoXFunControlPipeline",
+ "CogVideoXImageToVideoPipeline",
+ "CogVideoXPipeline",
+ "CogVideoXVideoToVideoPipeline",
+ "CogView3PlusPipeline",
+ "CogView4ControlPipeline",
+ "CogView4Pipeline",
+ "ConsisIDPipeline",
+ "Cosmos2_5_PredictBasePipeline",
+ "Cosmos2_5_TransferPipeline",
+ "Cosmos2TextToImagePipeline",
+ "Cosmos2VideoToWorldPipeline",
+ "Cosmos3OmniPipeline",
+ "CosmosActionCondition",
+ "CosmosTextToWorldPipeline",
+ "CosmosVideoToWorldPipeline",
+ "CycleDiffusionPipeline",
+ "DiffusionGemmaPipeline",
+ "DiffusionGemmaPipelineOutput",
+ "DreamLiteMobilePipeline",
+ "DreamLitePipeline",
+ "DreamLitePipelineOutput",
+ "EasyAnimateControlPipeline",
+ "EasyAnimateInpaintPipeline",
+ "EasyAnimatePipeline",
+ "ErnieImagePipeline",
+ "Flux2KleinInpaintPipeline",
+ "Flux2KleinKVPipeline",
+ "Flux2KleinPipeline",
+ "Flux2Pipeline",
+ "FluxControlImg2ImgPipeline",
+ "FluxControlInpaintPipeline",
+ "FluxControlNetImg2ImgPipeline",
+ "FluxControlNetInpaintPipeline",
+ "FluxControlNetPipeline",
+ "FluxControlPipeline",
+ "FluxFillPipeline",
+ "FluxImg2ImgPipeline",
+ "FluxInpaintPipeline",
+ "FluxKontextInpaintPipeline",
+ "FluxKontextPipeline",
+ "FluxPipeline",
+ "FluxPriorReduxPipeline",
+ "GlmImagePipeline",
+ "HeliosPipeline",
+ "HeliosPyramidPipeline",
+ "HiDreamImagePipeline",
+ "HunyuanDiTControlNetPipeline",
+ "HunyuanDiTPAGPipeline",
+ "HunyuanDiTPipeline",
+ "HunyuanImagePipeline",
+ "HunyuanImageRefinerPipeline",
+ "HunyuanSkyreelsImageToVideoPipeline",
+ "HunyuanVideo15ImageToVideoPipeline",
+ "HunyuanVideo15Pipeline",
+ "HunyuanVideoFramepackPipeline",
+ "HunyuanVideoImageToVideoPipeline",
+ "HunyuanVideoPipeline",
+ "I2VGenXLPipeline",
+ "Ideogram4Pipeline",
+ "Ideogram4PromptEnhancerHead",
+ "IFImg2ImgPipeline",
+ "IFImg2ImgSuperResolutionPipeline",
+ "IFInpaintingPipeline",
+ "IFInpaintingSuperResolutionPipeline",
+ "IFPipeline",
+ "IFSuperResolutionPipeline",
+ "ImageTextPipelineOutput",
+ "JoyImageEditPipeline",
+ "JoyImageEditPipelineOutput",
+ "JoyImageEditPlusPipeline",
+ "JoyImageEditPlusPipelineOutput",
+ "Kandinsky3Img2ImgPipeline",
+ "Kandinsky3Pipeline",
+ "Kandinsky5I2IPipeline",
+ "Kandinsky5I2VPipeline",
+ "Kandinsky5T2IPipeline",
+ "Kandinsky5T2VPipeline",
+ "KandinskyCombinedPipeline",
+ "KandinskyImg2ImgCombinedPipeline",
+ "KandinskyImg2ImgPipeline",
+ "KandinskyInpaintCombinedPipeline",
+ "KandinskyInpaintPipeline",
+ "KandinskyPipeline",
+ "KandinskyPriorPipeline",
+ "KandinskyV22CombinedPipeline",
+ "KandinskyV22ControlnetImg2ImgPipeline",
+ "KandinskyV22ControlnetPipeline",
+ "KandinskyV22Img2ImgCombinedPipeline",
+ "KandinskyV22Img2ImgPipeline",
+ "KandinskyV22InpaintCombinedPipeline",
+ "KandinskyV22InpaintPipeline",
+ "KandinskyV22Pipeline",
+ "KandinskyV22PriorEmb2EmbPipeline",
+ "KandinskyV22PriorPipeline",
+ "Krea2Pipeline",
+ "LatentConsistencyModelImg2ImgPipeline",
+ "LatentConsistencyModelPipeline",
+ "LattePipeline",
+ "LDMTextToImagePipeline",
+ "LEditsPPPipelineStableDiffusion",
+ "LEditsPPPipelineStableDiffusionXL",
+ "LLaDA2Pipeline",
+ "LLaDA2PipelineOutput",
+ "LongCatAudioDiTPipeline",
+ "LongCatImageEditPipeline",
+ "LongCatImagePipeline",
+ "LTX2ConditionPipeline",
+ "LTX2HDRPipeline",
+ "LTX2ImageToVideoPipeline",
+ "LTX2InContextPipeline",
+ "LTX2LatentUpsamplePipeline",
+ "LTX2Pipeline",
+ "LTXConditionPipeline",
+ "LTXI2VLongMultiPromptPipeline",
+ "LTXImageToVideoPipeline",
+ "LTXLatentUpsamplePipeline",
+ "LTXPipeline",
+ "LucyEditPipeline",
+ "Lumina2Pipeline",
+ "Lumina2Text2ImgPipeline",
+ "LuminaPipeline",
+ "LuminaText2ImgPipeline",
+ "MarigoldDepthPipeline",
+ "MarigoldIntrinsicsPipeline",
+ "MarigoldNormalsPipeline",
+ "MochiPipeline",
+ "MotifVideoImage2VideoPipeline",
+ "MotifVideoPipeline",
+ "MotifVideoPipelineOutput",
+ "MusicLDMPipeline",
+ "NucleusMoEImagePipeline",
+ "OmniGenPipeline",
+ "OvisImagePipeline",
+ "PaintByExamplePipeline",
+ "PIAPipeline",
+ "PixArtAlphaPipeline",
+ "PixArtSigmaPAGPipeline",
+ "PixArtSigmaPipeline",
+ "PRXPipeline",
+ "PRXPixelPipeline",
+ "QwenImageControlNetInpaintPipeline",
+ "QwenImageControlNetPipeline",
+ "QwenImageEditInpaintPipeline",
+ "QwenImageEditPipeline",
+ "QwenImageEditPlusPipeline",
+ "QwenImageImg2ImgPipeline",
+ "QwenImageInpaintPipeline",
+ "QwenImageLayeredPipeline",
+ "QwenImagePipeline",
+ "ReduxImageEncoder",
+ "SanaControlNetPipeline",
+ "SanaImageToVideoPipeline",
+ "SanaPAGPipeline",
+ "SanaPipeline",
+ "SanaSprintImg2ImgPipeline",
+ "SanaSprintPipeline",
+ "SanaVideoPipeline",
+ "SanaVideoPipeline",
+ "SemanticStableDiffusionPipeline",
+ "ShapEImg2ImgPipeline",
+ "ShapEPipeline",
+ "SkyReelsV2DiffusionForcingImageToVideoPipeline",
+ "SkyReelsV2DiffusionForcingPipeline",
+ "SkyReelsV2DiffusionForcingVideoToVideoPipeline",
+ "SkyReelsV2ImageToVideoPipeline",
+ "SkyReelsV2Pipeline",
+ "StableAudioPipeline",
+ "StableAudioProjectionModel",
+ "StableCascadeCombinedPipeline",
+ "StableCascadeDecoderPipeline",
+ "StableCascadePriorPipeline",
+ "StableDiffusion3ControlNetInpaintingPipeline",
+ "StableDiffusion3ControlNetPipeline",
+ "StableDiffusion3Img2ImgPipeline",
+ "StableDiffusion3InpaintPipeline",
+ "StableDiffusion3PAGImg2ImgPipeline",
+ "StableDiffusion3PAGImg2ImgPipeline",
+ "StableDiffusion3PAGPipeline",
+ "StableDiffusion3Pipeline",
+ "StableDiffusionAdapterPipeline",
+ "StableDiffusionAttendAndExcitePipeline",
+ "StableDiffusionControlNetImg2ImgPipeline",
+ "StableDiffusionControlNetInpaintPipeline",
+ "StableDiffusionControlNetPAGInpaintPipeline",
+ "StableDiffusionControlNetPAGPipeline",
+ "StableDiffusionControlNetPipeline",
+ "StableDiffusionControlNetXSPipeline",
+ "StableDiffusionDepth2ImgPipeline",
+ "StableDiffusionDiffEditPipeline",
+ "StableDiffusionGLIGENPipeline",
+ "StableDiffusionGLIGENTextImagePipeline",
+ "StableDiffusionImageVariationPipeline",
+ "StableDiffusionImg2ImgPipeline",
+ "StableDiffusionInpaintPipeline",
+ "StableDiffusionInpaintPipelineLegacy",
+ "StableDiffusionInstructPix2PixPipeline",
+ "StableDiffusionLatentUpscalePipeline",
+ "StableDiffusionLDM3DPipeline",
+ "StableDiffusionModelEditingPipeline",
+ "StableDiffusionPAGImg2ImgPipeline",
+ "StableDiffusionPAGInpaintPipeline",
+ "StableDiffusionPAGPipeline",
+ "StableDiffusionPanoramaPipeline",
+ "StableDiffusionParadigmsPipeline",
+ "StableDiffusionPipeline",
+ "StableDiffusionPipelineSafe",
+ "StableDiffusionPix2PixZeroPipeline",
+ "StableDiffusionSAGPipeline",
+ "StableDiffusionUpscalePipeline",
+ "StableDiffusionXLAdapterPipeline",
+ "StableDiffusionXLControlNetImg2ImgPipeline",
+ "StableDiffusionXLControlNetInpaintPipeline",
+ "StableDiffusionXLControlNetPAGImg2ImgPipeline",
+ "StableDiffusionXLControlNetPAGPipeline",
+ "StableDiffusionXLControlNetPipeline",
+ "StableDiffusionXLControlNetUnionImg2ImgPipeline",
+ "StableDiffusionXLControlNetUnionInpaintPipeline",
+ "StableDiffusionXLControlNetUnionPipeline",
+ "StableDiffusionXLControlNetXSPipeline",
+ "StableDiffusionXLImg2ImgPipeline",
+ "StableDiffusionXLInpaintPipeline",
+ "StableDiffusionXLInstructPix2PixPipeline",
+ "StableDiffusionXLPAGImg2ImgPipeline",
+ "StableDiffusionXLPAGInpaintPipeline",
+ "StableDiffusionXLPAGPipeline",
+ "StableDiffusionXLPipeline",
+ "StableUnCLIPImg2ImgPipeline",
+ "StableUnCLIPPipeline",
+ "StableVideoDiffusionPipeline",
+ "TextToVideoSDPipeline",
+ "TextToVideoZeroPipeline",
+ "TextToVideoZeroSDXLPipeline",
+ "UnCLIPImageVariationPipeline",
+ "UnCLIPPipeline",
+ "UniDiffuserModel",
+ "UniDiffuserPipeline",
+ "UniDiffuserTextDecoder",
+ "VersatileDiffusionDualGuidedPipeline",
+ "VersatileDiffusionImageVariationPipeline",
+ "VersatileDiffusionPipeline",
+ "VersatileDiffusionTextToImagePipeline",
+ "VideoToVideoSDPipeline",
+ "VisualClozeGenerationPipeline",
+ "VisualClozePipeline",
+ "VQDiffusionPipeline",
+ "WanAnimatePipeline",
+ "WanImageToVideoPipeline",
+ "WanPipeline",
+ "WanVACEPipeline",
+ "WanVideoToVideoPipeline",
+ "WuerstchenCombinedPipeline",
+ "WuerstchenDecoderPipeline",
+ "WuerstchenPriorPipeline",
+ "ZImageControlNetInpaintPipeline",
+ "ZImageControlNetPipeline",
+ "ZImageImg2ImgPipeline",
+ "ZImageInpaintPipeline",
+ "ZImageOmniPipeline",
+ "ZImagePipeline",
+ ]
+ )
+
+
+try:
+ if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_transformers_and_opencv_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_transformers_and_opencv_objects"] = [
+ name for name in dir(dummy_torch_and_transformers_and_opencv_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["pipelines"].extend(["ConsisIDPipeline"])
+
+try:
+ if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_transformers_and_sentencepiece_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_transformers_and_sentencepiece_objects"] = [
+ name for name in dir(dummy_torch_and_transformers_and_sentencepiece_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["pipelines"].extend(["KolorsImg2ImgPipeline", "KolorsPAGPipeline", "KolorsPipeline"])
+
+try:
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [
+ name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["pipelines"].extend(
+ [
+ "OnnxStableDiffusionImg2ImgPipeline",
+ "OnnxStableDiffusionInpaintPipeline",
+ "OnnxStableDiffusionInpaintPipelineLegacy",
+ "OnnxStableDiffusionPipeline",
+ "OnnxStableDiffusionUpscalePipeline",
+ "StableDiffusionOnnxPipeline",
+ ]
+ )
+
+try:
+ if not (is_torch_available() and is_librosa_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torch_and_librosa_objects # noqa F403
+
+ _import_structure["utils.dummy_torch_and_librosa_objects"] = [
+ name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_")
+ ]
+
+else:
+ _import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"])
+
+try:
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403
+
+ _import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [
+ name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_")
+ ]
+
+
+else:
+ _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"])
+
+try:
+ if not (is_note_seq_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_note_seq_objects # noqa F403
+
+ _import_structure["utils.dummy_note_seq_objects"] = [
+ name for name in dir(dummy_note_seq_objects) if not name.startswith("_")
+ ]
+
+
+else:
+ _import_structure["pipelines"].extend(["MidiProcessor"])
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ from .configuration_utils import ConfigMixin
+ from .quantizers import PipelineQuantizationConfig
+
+ try:
+ if not is_bitsandbytes_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_bitsandbytes_objects import *
+ else:
+ from .quantizers.quantization_config import BitsAndBytesConfig
+
+ try:
+ if not is_gguf_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_gguf_objects import *
+ else:
+ from .quantizers.quantization_config import GGUFQuantizationConfig
+
+ try:
+ if not is_torchao_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torchao_objects import *
+ else:
+ from .quantizers.quantization_config import TorchAoConfig
+
+ try:
+ if not is_optimum_quanto_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_optimum_quanto_objects import *
+ else:
+ from .quantizers.quantization_config import QuantoConfig
+
+ try:
+ if not is_nvidia_modelopt_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_nvidia_modelopt_objects import *
+ else:
+ from .quantizers.quantization_config import NVIDIAModelOptConfig
+
+ try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_nunchaku_lite_objects import *
+ else:
+ from .quantizers.quantization_config import NunchakuLiteQuantizationConfig
+
+ try:
+ if not is_auto_round_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_auto_round_objects import *
+ else:
+ from .quantizers.quantization_config import AutoRoundConfig
+
+ try:
+ if not is_sdnq_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_sdnq_objects import *
+ else:
+ from .quantizers.quantization_config import SDNQConfig
+
+ try:
+ if not is_onnx_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_onnx_objects import * # noqa F403
+ else:
+ from .pipelines import OnnxRuntimeModel
+
+ try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_pt_objects import * # noqa F403
+ else:
+ from .guiders import (
+ AdaptiveProjectedGuidance,
+ AdaptiveProjectedMixGuidance,
+ AutoGuidance,
+ BaseGuidance,
+ ClassifierFreeGuidance,
+ ClassifierFreeZeroStarGuidance,
+ FrequencyDecoupledGuidance,
+ PerturbedAttentionGuidance,
+ SkipLayerGuidance,
+ SmoothedEnergyGuidance,
+ TangentialClassifierFreeGuidance,
+ )
+ from .hooks import (
+ FasterCacheConfig,
+ FirstBlockCacheConfig,
+ HookRegistry,
+ LayerSkipConfig,
+ MagCacheConfig,
+ PyramidAttentionBroadcastConfig,
+ SmoothedEnergyGuidanceConfig,
+ TaylorSeerCacheConfig,
+ TextKVCacheConfig,
+ apply_faster_cache,
+ apply_first_block_cache,
+ apply_layer_skip,
+ apply_mag_cache,
+ apply_pyramid_attention_broadcast,
+ apply_taylorseer_cache,
+ apply_text_kv_cache,
+ )
+ from .image_processor import (
+ InpaintProcessor,
+ IPAdapterMaskProcessor,
+ PixArtImageProcessor,
+ VaeImageProcessor,
+ VaeImageProcessorLDM3D,
+ )
+ from .models import (
+ AceStepTransformer1DModel,
+ AllegroTransformer3DModel,
+ AnimaTextConditioner,
+ AnyFlowFARTransformer3DModel,
+ AnyFlowTransformer3DModel,
+ AsymmetricAutoencoderKL,
+ AttentionBackendName,
+ AuraFlowTransformer2DModel,
+ AutoencoderDC,
+ AutoencoderKL,
+ AutoencoderKLAllegro,
+ AutoencoderKLCogVideoX,
+ AutoencoderKLCosmos,
+ AutoencoderKLFlux2,
+ AutoencoderKLHunyuanImage,
+ AutoencoderKLHunyuanImageRefiner,
+ AutoencoderKLHunyuanVideo,
+ AutoencoderKLHunyuanVideo15,
+ AutoencoderKLKVAE,
+ AutoencoderKLKVAEVideo,
+ AutoencoderKLLTX2Audio,
+ AutoencoderKLLTX2Video,
+ AutoencoderKLLTXVideo,
+ AutoencoderKLMagvit,
+ AutoencoderKLMiniMaxH3,
+ AutoencoderKLMiniMaxH3Audio,
+ AutoencoderKLMochi,
+ AutoencoderKLQwenImage,
+ AutoencoderKLTemporalDecoder,
+ AutoencoderKLWan,
+ AutoencoderOobleck,
+ AutoencoderRAE,
+ AutoencoderTiny,
+ AutoencoderVidTok,
+ AutoModel,
+ BriaFiboTransformer2DModel,
+ BriaTransformer2DModel,
+ CacheMixin,
+ ChromaTransformer2DModel,
+ ChronoEditTransformer3DModel,
+ CogVideoXTransformer3DModel,
+ CogView3PlusTransformer2DModel,
+ CogView4Transformer2DModel,
+ ConsisIDTransformer3DModel,
+ ConsistencyDecoderVAE,
+ ContextParallelConfig,
+ ControlNetModel,
+ ControlNetUnionModel,
+ ControlNetXSAdapter,
+ Cosmos3AVAEAudioTokenizer,
+ Cosmos3OmniTransformer,
+ CosmosControlNetModel,
+ CosmosTransformer3DModel,
+ DiTTransformer2DModel,
+ DreamLiteTransformer2DModel,
+ DreamLiteUNetModel,
+ EasyAnimateTransformer3DModel,
+ ErnieImageTransformer2DModel,
+ Flux2Transformer2DModel,
+ FluxControlNetModel,
+ FluxMultiControlNetModel,
+ FluxTransformer2DModel,
+ GlmImageTransformer2DModel,
+ HeliosTransformer3DModel,
+ HiDreamImageTransformer2DModel,
+ HunyuanDiT2DControlNetModel,
+ HunyuanDiT2DModel,
+ HunyuanDiT2DMultiControlNetModel,
+ HunyuanImageTransformer2DModel,
+ HunyuanVideo15Transformer3DModel,
+ HunyuanVideoFramepackTransformer3DModel,
+ HunyuanVideoTransformer3DModel,
+ I2VGenXLUNet,
+ Ideogram4Transformer2DModel,
+ JoyImageEditPlusTransformer3DModel,
+ JoyImageEditTransformer3DModel,
+ Kandinsky3UNet,
+ Kandinsky5Transformer3DModel,
+ Krea2Transformer2DModel,
+ LatteTransformer3DModel,
+ LongCatAudioDiTTransformer,
+ LongCatAudioDiTVae,
+ LongCatImageTransformer2DModel,
+ LTX2VideoTransformer3DModel,
+ LTXVideoTransformer3DModel,
+ Lumina2Transformer2DModel,
+ LuminaNextDiT2DModel,
+ MiniMaxH3Transformer3DModel,
+ MochiTransformer3DModel,
+ ModelMixin,
+ MotifVideoTransformer3DModel,
+ MotionAdapter,
+ MultiAdapter,
+ MultiControlNetModel,
+ NucleusMoEImageTransformer2DModel,
+ OmniGenTransformer2DModel,
+ OvisImageTransformer2DModel,
+ ParallelConfig,
+ PixArtTransformer2DModel,
+ PriorTransformer,
+ PRXTransformer2DModel,
+ QwenImageControlNetModel,
+ QwenImageMultiControlNetModel,
+ QwenImageTransformer2DModel,
+ SanaControlNetModel,
+ SanaTransformer2DModel,
+ SanaVideoTransformer3DModel,
+ SD3ControlNetModel,
+ SD3MultiControlNetModel,
+ SD3Transformer2DModel,
+ SkyReelsV2Transformer3DModel,
+ SparseControlNetModel,
+ StableAudioDiTModel,
+ T2IAdapter,
+ T5FilmDecoder,
+ Transformer2DModel,
+ TransformerTemporalModel,
+ UNet1DModel,
+ UNet2DConditionModel,
+ UNet2DModel,
+ UNet3DConditionModel,
+ UNetControlNetXSModel,
+ UNetMotionModel,
+ UNetSpatioTemporalConditionModel,
+ UVit2DModel,
+ VQModel,
+ WanAnimateTransformer3DModel,
+ WanTransformer3DModel,
+ WanVACETransformer3DModel,
+ ZImageControlNetModel,
+ ZImageTransformer2DModel,
+ attention_backend,
+ )
+ from .modular_pipelines import (
+ AutoPipelineBlocks,
+ ComponentsManager,
+ ComponentSpec,
+ ConditionalPipelineBlocks,
+ ConfigSpec,
+ InputParam,
+ LoopSequentialPipelineBlocks,
+ ModularPipeline,
+ ModularPipelineBlocks,
+ OutputParam,
+ SequentialPipelineBlocks,
+ )
+ from .optimization import (
+ get_constant_schedule,
+ get_constant_schedule_with_warmup,
+ get_cosine_schedule_with_warmup,
+ get_cosine_with_hard_restarts_schedule_with_warmup,
+ get_linear_schedule_with_warmup,
+ get_polynomial_decay_schedule_with_warmup,
+ get_scheduler,
+ )
+ from .pipelines import (
+ AudioPipelineOutput,
+ AutoPipelineForImage2Image,
+ AutoPipelineForInpainting,
+ AutoPipelineForText2Audio,
+ AutoPipelineForText2Image,
+ BlipDiffusionControlNetPipeline,
+ BlipDiffusionPipeline,
+ CLIPImageProjection,
+ ConsistencyModelPipeline,
+ DanceDiffusionPipeline,
+ DDIMPipeline,
+ DDPMPipeline,
+ DiffusionPipeline,
+ DiTPipeline,
+ ImagePipelineOutput,
+ KarrasVePipeline,
+ LDMPipeline,
+ LDMSuperResolutionPipeline,
+ PNDMPipeline,
+ RePaintPipeline,
+ ScoreSdeVePipeline,
+ StableDiffusionMixin,
+ )
+ from .quantizers import DiffusersQuantizer
+ from .schedulers import (
+ AmusedScheduler,
+ BlockRefinementScheduler,
+ BlockRefinementSchedulerOutput,
+ CMStochasticIterativeScheduler,
+ CogVideoXDDIMScheduler,
+ CogVideoXDPMScheduler,
+ DDIMInverseScheduler,
+ DDIMParallelScheduler,
+ DDIMScheduler,
+ DDPMParallelScheduler,
+ DDPMScheduler,
+ DDPMWuerstchenScheduler,
+ DEISMultistepScheduler,
+ DiscreteDDIMScheduler,
+ DiscreteDDIMSchedulerOutput,
+ DPMSolverMultistepInverseScheduler,
+ DPMSolverMultistepScheduler,
+ DPMSolverSinglestepScheduler,
+ EDMDPMSolverMultistepScheduler,
+ EDMEulerScheduler,
+ EntropyBoundScheduler,
+ EntropyBoundSchedulerOutput,
+ EulerAncestralDiscreteScheduler,
+ EulerDiscreteScheduler,
+ FlowMapEulerDiscreteScheduler,
+ FlowMatchEulerDiscreteScheduler,
+ FlowMatchHeunDiscreteScheduler,
+ FlowMatchLCMScheduler,
+ HeliosDMDScheduler,
+ HeliosScheduler,
+ HeunDiscreteScheduler,
+ IPNDMScheduler,
+ KarrasVeScheduler,
+ KDPM2AncestralDiscreteScheduler,
+ KDPM2DiscreteScheduler,
+ LCMScheduler,
+ LTXEulerAncestralRFScheduler,
+ MiniMaxH3Scheduler,
+ PNDMScheduler,
+ RePaintScheduler,
+ SASolverScheduler,
+ SchedulerMixin,
+ SCMScheduler,
+ ScoreSdeVeScheduler,
+ TCDScheduler,
+ UnCLIPScheduler,
+ UniPCMultistepScheduler,
+ VQDiffusionScheduler,
+ )
+ from .training_utils import EMAModel
+ from .video_processor import VideoProcessor
+
+ try:
+ if not (is_torch_available() and is_scipy_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_scipy_objects import * # noqa F403
+ else:
+ from .schedulers import LMSDiscreteScheduler
+
+ try:
+ if not (is_torch_available() and is_torchsde_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_torchsde_objects import * # noqa F403
+ else:
+ from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler
+
+ try:
+ if not (is_torch_available() and is_transformers_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_pipelines import (
+ AnimaAutoBlocks,
+ AnimaModularPipeline,
+ Cosmos3DistilledBlocks,
+ Cosmos3DistilledModularPipeline,
+ Cosmos3OmniBlocks,
+ Cosmos3OmniModularPipeline,
+ ErnieImageAutoBlocks,
+ ErnieImageModularPipeline,
+ Flux2AutoBlocks,
+ Flux2KleinAutoBlocks,
+ Flux2KleinBaseAutoBlocks,
+ Flux2KleinBaseModularPipeline,
+ Flux2KleinModularPipeline,
+ Flux2ModularPipeline,
+ FluxAutoBlocks,
+ FluxKontextAutoBlocks,
+ FluxKontextModularPipeline,
+ FluxModularPipeline,
+ HeliosAutoBlocks,
+ HeliosModularPipeline,
+ HeliosPyramidAutoBlocks,
+ HeliosPyramidDistilledAutoBlocks,
+ HeliosPyramidDistilledModularPipeline,
+ HeliosPyramidModularPipeline,
+ HunyuanVideo15AutoBlocks,
+ HunyuanVideo15ModularPipeline,
+ Ideogram4AutoBlocks,
+ Ideogram4ModularPipeline,
+ Krea2AutoBlocks,
+ Krea2ModularPipeline,
+ Krea2TurboAutoBlocks,
+ Krea2TurboModularPipeline,
+ LTXAutoBlocks,
+ LTXModularPipeline,
+ MiniMaxH3Blocks,
+ MiniMaxH3ModularPipeline,
+ QwenImageAutoBlocks,
+ QwenImageEditAutoBlocks,
+ QwenImageEditModularPipeline,
+ QwenImageEditPlusAutoBlocks,
+ QwenImageEditPlusModularPipeline,
+ QwenImageLayeredAutoBlocks,
+ QwenImageLayeredModularPipeline,
+ QwenImageModularPipeline,
+ StableDiffusion3AutoBlocks,
+ StableDiffusion3ModularPipeline,
+ StableDiffusionXLAutoBlocks,
+ StableDiffusionXLModularPipeline,
+ Wan22Blocks,
+ Wan22Image2VideoBlocks,
+ Wan22Image2VideoModularPipeline,
+ Wan22ModularPipeline,
+ WanBlocks,
+ WanImage2VideoAutoBlocks,
+ WanImage2VideoModularPipeline,
+ WanModularPipeline,
+ ZImageAutoBlocks,
+ ZImageModularPipeline,
+ )
+ from .pipelines import (
+ AceStepAudioTokenDetokenizer,
+ AceStepAudioTokenizer,
+ AceStepConditionEncoder,
+ AceStepPipeline,
+ AllegroPipeline,
+ AltDiffusionImg2ImgPipeline,
+ AltDiffusionPipeline,
+ AmusedImg2ImgPipeline,
+ AmusedInpaintPipeline,
+ AmusedPipeline,
+ AnimateDiffControlNetPipeline,
+ AnimateDiffPAGPipeline,
+ AnimateDiffPipeline,
+ AnimateDiffSDXLPipeline,
+ AnimateDiffSparseControlNetPipeline,
+ AnimateDiffVideoToVideoControlNetPipeline,
+ AnimateDiffVideoToVideoPipeline,
+ AnyFlowFARPipeline,
+ AnyFlowPipeline,
+ AudioLDM2Pipeline,
+ AudioLDM2ProjectionModel,
+ AudioLDM2UNet2DConditionModel,
+ AudioLDMPipeline,
+ AuraFlowPipeline,
+ BriaFiboEditPipeline,
+ BriaFiboPipeline,
+ BriaPipeline,
+ ChromaImg2ImgPipeline,
+ ChromaInpaintPipeline,
+ ChromaPipeline,
+ ChronoEditPipeline,
+ CLIPImageProjection,
+ CogVideoXFunControlPipeline,
+ CogVideoXImageToVideoPipeline,
+ CogVideoXPipeline,
+ CogVideoXVideoToVideoPipeline,
+ CogView3PlusPipeline,
+ CogView4ControlPipeline,
+ CogView4Pipeline,
+ ConsisIDPipeline,
+ Cosmos2_5_PredictBasePipeline,
+ Cosmos2_5_TransferPipeline,
+ Cosmos2TextToImagePipeline,
+ Cosmos2VideoToWorldPipeline,
+ Cosmos3OmniPipeline,
+ CosmosActionCondition,
+ CosmosTextToWorldPipeline,
+ CosmosVideoToWorldPipeline,
+ CycleDiffusionPipeline,
+ DiffusionGemmaPipeline,
+ DiffusionGemmaPipelineOutput,
+ DreamLiteMobilePipeline,
+ DreamLitePipeline,
+ DreamLitePipelineOutput,
+ EasyAnimateControlPipeline,
+ EasyAnimateInpaintPipeline,
+ EasyAnimatePipeline,
+ ErnieImagePipeline,
+ Flux2KleinInpaintPipeline,
+ Flux2KleinKVPipeline,
+ Flux2KleinPipeline,
+ Flux2Pipeline,
+ FluxControlImg2ImgPipeline,
+ FluxControlInpaintPipeline,
+ FluxControlNetImg2ImgPipeline,
+ FluxControlNetInpaintPipeline,
+ FluxControlNetPipeline,
+ FluxControlPipeline,
+ FluxFillPipeline,
+ FluxImg2ImgPipeline,
+ FluxInpaintPipeline,
+ FluxKontextInpaintPipeline,
+ FluxKontextPipeline,
+ FluxPipeline,
+ FluxPriorReduxPipeline,
+ GlmImagePipeline,
+ HeliosPipeline,
+ HeliosPyramidPipeline,
+ HiDreamImagePipeline,
+ HunyuanDiTControlNetPipeline,
+ HunyuanDiTPAGPipeline,
+ HunyuanDiTPipeline,
+ HunyuanImagePipeline,
+ HunyuanImageRefinerPipeline,
+ HunyuanSkyreelsImageToVideoPipeline,
+ HunyuanVideo15ImageToVideoPipeline,
+ HunyuanVideo15Pipeline,
+ HunyuanVideoFramepackPipeline,
+ HunyuanVideoImageToVideoPipeline,
+ HunyuanVideoPipeline,
+ I2VGenXLPipeline,
+ Ideogram4Pipeline,
+ Ideogram4PromptEnhancerHead,
+ IFImg2ImgPipeline,
+ IFImg2ImgSuperResolutionPipeline,
+ IFInpaintingPipeline,
+ IFInpaintingSuperResolutionPipeline,
+ IFPipeline,
+ IFSuperResolutionPipeline,
+ ImageTextPipelineOutput,
+ JoyImageEditPipeline,
+ JoyImageEditPipelineOutput,
+ JoyImageEditPlusPipeline,
+ JoyImageEditPlusPipelineOutput,
+ Kandinsky3Img2ImgPipeline,
+ Kandinsky3Pipeline,
+ Kandinsky5I2IPipeline,
+ Kandinsky5I2VPipeline,
+ Kandinsky5T2IPipeline,
+ Kandinsky5T2VPipeline,
+ KandinskyCombinedPipeline,
+ KandinskyImg2ImgCombinedPipeline,
+ KandinskyImg2ImgPipeline,
+ KandinskyInpaintCombinedPipeline,
+ KandinskyInpaintPipeline,
+ KandinskyPipeline,
+ KandinskyPriorPipeline,
+ KandinskyV22CombinedPipeline,
+ KandinskyV22ControlnetImg2ImgPipeline,
+ KandinskyV22ControlnetPipeline,
+ KandinskyV22Img2ImgCombinedPipeline,
+ KandinskyV22Img2ImgPipeline,
+ KandinskyV22InpaintCombinedPipeline,
+ KandinskyV22InpaintPipeline,
+ KandinskyV22Pipeline,
+ KandinskyV22PriorEmb2EmbPipeline,
+ KandinskyV22PriorPipeline,
+ Krea2Pipeline,
+ LatentConsistencyModelImg2ImgPipeline,
+ LatentConsistencyModelPipeline,
+ LattePipeline,
+ LDMTextToImagePipeline,
+ LEditsPPPipelineStableDiffusion,
+ LEditsPPPipelineStableDiffusionXL,
+ LLaDA2Pipeline,
+ LLaDA2PipelineOutput,
+ LongCatAudioDiTPipeline,
+ LongCatImageEditPipeline,
+ LongCatImagePipeline,
+ LTX2ConditionPipeline,
+ LTX2HDRPipeline,
+ LTX2ImageToVideoPipeline,
+ LTX2InContextPipeline,
+ LTX2LatentUpsamplePipeline,
+ LTX2Pipeline,
+ LTXConditionPipeline,
+ LTXI2VLongMultiPromptPipeline,
+ LTXImageToVideoPipeline,
+ LTXLatentUpsamplePipeline,
+ LTXPipeline,
+ LucyEditPipeline,
+ Lumina2Pipeline,
+ Lumina2Text2ImgPipeline,
+ LuminaPipeline,
+ LuminaText2ImgPipeline,
+ MarigoldDepthPipeline,
+ MarigoldIntrinsicsPipeline,
+ MarigoldNormalsPipeline,
+ MochiPipeline,
+ MotifVideoImage2VideoPipeline,
+ MotifVideoPipeline,
+ MotifVideoPipelineOutput,
+ MusicLDMPipeline,
+ NucleusMoEImagePipeline,
+ OmniGenPipeline,
+ OvisImagePipeline,
+ PaintByExamplePipeline,
+ PIAPipeline,
+ PixArtAlphaPipeline,
+ PixArtSigmaPAGPipeline,
+ PixArtSigmaPipeline,
+ PRXPipeline,
+ PRXPixelPipeline,
+ QwenImageControlNetInpaintPipeline,
+ QwenImageControlNetPipeline,
+ QwenImageEditInpaintPipeline,
+ QwenImageEditPipeline,
+ QwenImageEditPlusPipeline,
+ QwenImageImg2ImgPipeline,
+ QwenImageInpaintPipeline,
+ QwenImageLayeredPipeline,
+ QwenImagePipeline,
+ ReduxImageEncoder,
+ SanaControlNetPipeline,
+ SanaImageToVideoPipeline,
+ SanaPAGPipeline,
+ SanaPipeline,
+ SanaSprintImg2ImgPipeline,
+ SanaSprintPipeline,
+ SanaVideoPipeline,
+ SemanticStableDiffusionPipeline,
+ ShapEImg2ImgPipeline,
+ ShapEPipeline,
+ SkyReelsV2DiffusionForcingImageToVideoPipeline,
+ SkyReelsV2DiffusionForcingPipeline,
+ SkyReelsV2DiffusionForcingVideoToVideoPipeline,
+ SkyReelsV2ImageToVideoPipeline,
+ SkyReelsV2Pipeline,
+ StableAudioPipeline,
+ StableAudioProjectionModel,
+ StableCascadeCombinedPipeline,
+ StableCascadeDecoderPipeline,
+ StableCascadePriorPipeline,
+ StableDiffusion3ControlNetInpaintingPipeline,
+ StableDiffusion3ControlNetPipeline,
+ StableDiffusion3Img2ImgPipeline,
+ StableDiffusion3InpaintPipeline,
+ StableDiffusion3PAGImg2ImgPipeline,
+ StableDiffusion3PAGPipeline,
+ StableDiffusion3Pipeline,
+ StableDiffusionAdapterPipeline,
+ StableDiffusionAttendAndExcitePipeline,
+ StableDiffusionControlNetImg2ImgPipeline,
+ StableDiffusionControlNetInpaintPipeline,
+ StableDiffusionControlNetPAGInpaintPipeline,
+ StableDiffusionControlNetPAGPipeline,
+ StableDiffusionControlNetPipeline,
+ StableDiffusionControlNetXSPipeline,
+ StableDiffusionDepth2ImgPipeline,
+ StableDiffusionDiffEditPipeline,
+ StableDiffusionGLIGENPipeline,
+ StableDiffusionGLIGENTextImagePipeline,
+ StableDiffusionImageVariationPipeline,
+ StableDiffusionImg2ImgPipeline,
+ StableDiffusionInpaintPipeline,
+ StableDiffusionInpaintPipelineLegacy,
+ StableDiffusionInstructPix2PixPipeline,
+ StableDiffusionLatentUpscalePipeline,
+ StableDiffusionLDM3DPipeline,
+ StableDiffusionModelEditingPipeline,
+ StableDiffusionPAGImg2ImgPipeline,
+ StableDiffusionPAGInpaintPipeline,
+ StableDiffusionPAGPipeline,
+ StableDiffusionPanoramaPipeline,
+ StableDiffusionParadigmsPipeline,
+ StableDiffusionPipeline,
+ StableDiffusionPipelineSafe,
+ StableDiffusionPix2PixZeroPipeline,
+ StableDiffusionSAGPipeline,
+ StableDiffusionUpscalePipeline,
+ StableDiffusionXLAdapterPipeline,
+ StableDiffusionXLControlNetImg2ImgPipeline,
+ StableDiffusionXLControlNetInpaintPipeline,
+ StableDiffusionXLControlNetPAGImg2ImgPipeline,
+ StableDiffusionXLControlNetPAGPipeline,
+ StableDiffusionXLControlNetPipeline,
+ StableDiffusionXLControlNetUnionImg2ImgPipeline,
+ StableDiffusionXLControlNetUnionInpaintPipeline,
+ StableDiffusionXLControlNetUnionPipeline,
+ StableDiffusionXLControlNetXSPipeline,
+ StableDiffusionXLImg2ImgPipeline,
+ StableDiffusionXLInpaintPipeline,
+ StableDiffusionXLInstructPix2PixPipeline,
+ StableDiffusionXLPAGImg2ImgPipeline,
+ StableDiffusionXLPAGInpaintPipeline,
+ StableDiffusionXLPAGPipeline,
+ StableDiffusionXLPipeline,
+ StableUnCLIPImg2ImgPipeline,
+ StableUnCLIPPipeline,
+ StableVideoDiffusionPipeline,
+ TextToVideoSDPipeline,
+ TextToVideoZeroPipeline,
+ TextToVideoZeroSDXLPipeline,
+ UnCLIPImageVariationPipeline,
+ UnCLIPPipeline,
+ UniDiffuserModel,
+ UniDiffuserPipeline,
+ UniDiffuserTextDecoder,
+ VersatileDiffusionDualGuidedPipeline,
+ VersatileDiffusionImageVariationPipeline,
+ VersatileDiffusionPipeline,
+ VersatileDiffusionTextToImagePipeline,
+ VideoToVideoSDPipeline,
+ VisualClozeGenerationPipeline,
+ VisualClozePipeline,
+ VQDiffusionPipeline,
+ WanAnimatePipeline,
+ WanImageToVideoPipeline,
+ WanPipeline,
+ WanVACEPipeline,
+ WanVideoToVideoPipeline,
+ WuerstchenCombinedPipeline,
+ WuerstchenDecoderPipeline,
+ WuerstchenPriorPipeline,
+ ZImageControlNetInpaintPipeline,
+ ZImageControlNetPipeline,
+ ZImageImg2ImgPipeline,
+ ZImageInpaintPipeline,
+ ZImageOmniPipeline,
+ ZImagePipeline,
+ )
+
+ try:
+ if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_transformers_and_sentencepiece_objects import * # noqa F403
+ else:
+ from .pipelines import KolorsImg2ImgPipeline, KolorsPAGPipeline, KolorsPipeline
+
+ try:
+ if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_transformers_and_opencv_objects import * # noqa F403
+ else:
+ from .pipelines import ConsisIDPipeline
+
+ try:
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403
+ else:
+ from .pipelines import (
+ OnnxStableDiffusionImg2ImgPipeline,
+ OnnxStableDiffusionInpaintPipeline,
+ OnnxStableDiffusionInpaintPipelineLegacy,
+ OnnxStableDiffusionPipeline,
+ OnnxStableDiffusionUpscalePipeline,
+ StableDiffusionOnnxPipeline,
+ )
+
+ try:
+ if not (is_torch_available() and is_librosa_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torch_and_librosa_objects import * # noqa F403
+ else:
+ from .pipelines import AudioDiffusionPipeline, Mel
+
+ try:
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403
+ else:
+ from .pipelines import SpectrogramDiffusionPipeline
+
+ try:
+ if not (is_note_seq_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_note_seq_objects import * # noqa F403
+ else:
+ from .pipelines import MidiProcessor
+
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ extra_objects={"__version__": __version__},
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py
new file mode 100644
index 0000000000000000000000000000000000000000..83a307899dd73db6e978414af5ed2f78b64d3849
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py
@@ -0,0 +1,244 @@
+from typing import Any
+
+from .configuration_utils import ConfigMixin, register_to_config
+from .utils import CONFIG_NAME
+
+
+class PipelineCallback(ConfigMixin):
+ """
+ Base class for all the official callbacks used in a pipeline. This class provides a structure for implementing
+ custom callbacks and ensures that all callbacks have a consistent interface.
+
+ Please implement the following:
+ `tensor_inputs`: This should return a list of tensor inputs specific to your callback. You will only be able to
+ include
+ variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
+ `callback_fn`: This method defines the core functionality of your callback.
+ """
+
+ config_name = CONFIG_NAME
+
+ @register_to_config
+ def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None):
+ super().__init__()
+
+ if (cutoff_step_ratio is None and cutoff_step_index is None) or (
+ cutoff_step_ratio is not None and cutoff_step_index is not None
+ ):
+ raise ValueError("Either cutoff_step_ratio or cutoff_step_index should be provided, not both or none.")
+
+ if cutoff_step_ratio is not None and (
+ not isinstance(cutoff_step_ratio, float) or not (0.0 <= cutoff_step_ratio <= 1.0)
+ ):
+ raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.")
+
+ @property
+ def tensor_inputs(self) -> list[str]:
+ raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}")
+
+ def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> dict[str, Any]:
+ raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}")
+
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ return self.callback_fn(pipeline, step_index, timestep, callback_kwargs)
+
+
+class MultiPipelineCallbacks:
+ """
+ This class is designed to handle multiple pipeline callbacks. It accepts a list of PipelineCallback objects and
+ provides a unified interface for calling all of them.
+ """
+
+ def __init__(self, callbacks: list[PipelineCallback]):
+ self.callbacks = callbacks
+
+ @property
+ def tensor_inputs(self) -> list[str]:
+ return [input for callback in self.callbacks for input in callback.tensor_inputs]
+
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ """
+ Calls all the callbacks in order with the given arguments and returns the final callback_kwargs.
+ """
+ for callback in self.callbacks:
+ callback_kwargs = callback(pipeline, step_index, timestep, callback_kwargs)
+
+ return callback_kwargs
+
+
+class SDCFGCutoffCallback(PipelineCallback):
+ """
+ Callback function for Stable Diffusion Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
+ `cutoff_step_index`), this callback will disable the CFG.
+
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
+ """
+
+ tensor_inputs = ["prompt_embeds"]
+
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ cutoff_step_ratio = self.config.cutoff_step_ratio
+ cutoff_step_index = self.config.cutoff_step_index
+
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
+ cutoff_step = (
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
+ )
+
+ if step_index == cutoff_step:
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
+
+ pipeline._guidance_scale = 0.0
+
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
+ return callback_kwargs
+
+
+class SDXLCFGCutoffCallback(PipelineCallback):
+ """
+ Callback function for the base Stable Diffusion XL Pipelines. After certain number of steps (set by
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
+
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
+ """
+
+ tensor_inputs = [
+ "prompt_embeds",
+ "add_text_embeds",
+ "add_time_ids",
+ ]
+
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ cutoff_step_ratio = self.config.cutoff_step_ratio
+ cutoff_step_index = self.config.cutoff_step_index
+
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
+ cutoff_step = (
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
+ )
+
+ if step_index == cutoff_step:
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
+
+ add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
+ add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens
+
+ add_time_ids = callback_kwargs[self.tensor_inputs[2]]
+ add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector
+
+ pipeline._guidance_scale = 0.0
+
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
+ callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
+ callback_kwargs[self.tensor_inputs[2]] = add_time_ids
+
+ return callback_kwargs
+
+
+class SDXLControlnetCFGCutoffCallback(PipelineCallback):
+ """
+ Callback function for the Controlnet Stable Diffusion XL Pipelines. After certain number of steps (set by
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
+
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
+ """
+
+ tensor_inputs = [
+ "prompt_embeds",
+ "add_text_embeds",
+ "add_time_ids",
+ "image",
+ ]
+
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ cutoff_step_ratio = self.config.cutoff_step_ratio
+ cutoff_step_index = self.config.cutoff_step_index
+
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
+ cutoff_step = (
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
+ )
+
+ if step_index == cutoff_step:
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
+
+ add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
+ add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens
+
+ add_time_ids = callback_kwargs[self.tensor_inputs[2]]
+ add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector
+
+ # For Controlnet
+ image = callback_kwargs[self.tensor_inputs[3]]
+ image = image[-1:]
+
+ pipeline._guidance_scale = 0.0
+
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
+ callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
+ callback_kwargs[self.tensor_inputs[2]] = add_time_ids
+ callback_kwargs[self.tensor_inputs[3]] = image
+
+ return callback_kwargs
+
+
+class IPAdapterScaleCutoffCallback(PipelineCallback):
+ """
+ Callback function for any pipeline that inherits `IPAdapterMixin`. After certain number of steps (set by
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will set the IP Adapter scale to `0.0`.
+
+ Note: This callback mutates the IP Adapter attention processors by setting the scale to 0.0 after the cutoff step.
+ """
+
+ tensor_inputs = []
+
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ cutoff_step_ratio = self.config.cutoff_step_ratio
+ cutoff_step_index = self.config.cutoff_step_index
+
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
+ cutoff_step = (
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
+ )
+
+ if step_index == cutoff_step:
+ pipeline.set_ip_adapter_scale(0.0)
+ return callback_kwargs
+
+
+class SD3CFGCutoffCallback(PipelineCallback):
+ """
+ Callback function for Stable Diffusion 3 Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
+ `cutoff_step_index`), this callback will disable the CFG.
+
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
+ """
+
+ tensor_inputs = ["prompt_embeds", "pooled_prompt_embeds"]
+
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
+ cutoff_step_ratio = self.config.cutoff_step_ratio
+ cutoff_step_index = self.config.cutoff_step_index
+
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
+ cutoff_step = (
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
+ )
+
+ if step_index == cutoff_step:
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
+
+ pooled_prompt_embeds = callback_kwargs[self.tensor_inputs[1]]
+ pooled_prompt_embeds = pooled_prompt_embeds[
+ -1:
+ ] # "-1" denotes the embeddings for conditional pooled text tokens.
+
+ pipeline._guidance_scale = 0.0
+
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
+ callback_kwargs[self.tensor_inputs[1]] = pooled_prompt_embeds
+ return callback_kwargs
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..123b39edf73088663c076efa10bb12ee4adf6526
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py
@@ -0,0 +1,752 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ConfigMixin base class and utilities."""
+
+import functools
+import importlib
+import inspect
+import json
+import os
+import re
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from huggingface_hub import DDUFEntry, create_repo, hf_hub_download
+from huggingface_hub.utils import (
+ EntryNotFoundError,
+ HfHubHTTPError,
+ RepositoryNotFoundError,
+ RevisionNotFoundError,
+ validate_hf_hub_args,
+)
+from typing_extensions import Self
+
+from . import __version__
+from .utils import (
+ HUGGINGFACE_CO_RESOLVE_ENDPOINT,
+ DummyObject,
+ deprecate,
+ extract_commit_hash,
+ http_user_agent,
+ logging,
+)
+
+
+logger = logging.get_logger(__name__)
+
+_re_configuration_file = re.compile(r"config\.(.*)\.json")
+
+
+class FrozenDict(OrderedDict):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ for key, value in self.items():
+ setattr(self, key, value)
+
+ self.__frozen = True
+
+ def __delitem__(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
+
+ def setdefault(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
+
+ def pop(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
+
+ def update(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
+
+ def __setattr__(self, name, value):
+ if hasattr(self, "__frozen") and self.__frozen:
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
+ super().__setattr__(name, value)
+
+ def __setitem__(self, name, value):
+ if hasattr(self, "__frozen") and self.__frozen:
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
+ super().__setitem__(name, value)
+
+
+class ConfigMixin:
+ r"""
+ Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
+ provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
+ saving classes that inherit from [`ConfigMixin`].
+
+ Class attributes:
+ - **config_name** (`str`) -- A filename under which the config should stored when calling
+ [`~ConfigMixin.save_config`] (should be overridden by parent class).
+ - **ignore_for_config** (`list[str]`) -- A list of attributes that should not be saved in the config (should be
+ overridden by subclass).
+ - **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
+ - **_deprecated_kwargs** (`list[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
+ should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
+ subclass).
+ """
+
+ config_name = None
+ ignore_for_config = []
+ has_compatibles = False
+
+ _deprecated_kwargs = []
+ _auto_class = None
+
+ @classmethod
+ def register_for_auto_class(cls, auto_class="AutoModel"):
+ """
+ Register this class with the given auto class so that it can be loaded with `AutoModel.from_pretrained(...,
+ trust_remote_code=True)`.
+
+ When the config is saved, the resulting `config.json` will include an `auto_map` entry mapping the auto class
+ to this class's module and class name.
+
+ Args:
+ auto_class (`str` or type, *optional*, defaults to `"AutoModel"`):
+ The auto class to register this class with. Can be a string (e.g. `"AutoModel"`) or the class itself.
+ Currently only `"AutoModel"` is supported.
+
+ Example:
+
+ ```python
+ from diffusers import ModelMixin, ConfigMixin
+
+
+ class MyCustomModel(ModelMixin, ConfigMixin): ...
+
+
+ MyCustomModel.register_for_auto_class("AutoModel")
+ ```
+ """
+ if auto_class != "AutoModel":
+ raise ValueError(f"Only 'AutoModel' is supported, got '{auto_class}'.")
+
+ cls._auto_class = auto_class
+
+ def register_to_config(self, **kwargs):
+ if self.config_name is None:
+ raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
+ # Special case for `kwargs` used in deprecation warning added to schedulers
+ # TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
+ # or solve in a more general way.
+ kwargs.pop("kwargs", None)
+
+ if not hasattr(self, "_internal_dict"):
+ internal_dict = kwargs
+ else:
+ previous_dict = dict(self._internal_dict)
+ internal_dict = {**self._internal_dict, **kwargs}
+ logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
+
+ self._internal_dict = FrozenDict(internal_dict)
+
+ def __getattr__(self, name: str) -> Any:
+ """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
+ config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
+
+ This function is mostly copied from PyTorch's __getattr__ overwrite:
+ https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
+ """
+
+ is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
+ is_attribute = name in self.__dict__
+
+ if is_in_config and not is_attribute:
+ deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'."
+ deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
+ return self._internal_dict[name]
+
+ raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+ def save_config(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
+ """
+ Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
+ [`~ConfigMixin.from_config`] class method.
+
+ Args:
+ save_directory (`str` or `os.PathLike`):
+ Directory where the configuration JSON file is saved (will be created if it does not exist).
+ push_to_hub (`bool`, *optional*, defaults to `False`):
+ Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
+ repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
+ namespace).
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
+ """
+ if os.path.isfile(save_directory):
+ raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+ os.makedirs(save_directory, exist_ok=True)
+
+ # If we save using the predefined names, we can load using `from_config`
+ output_config_file = os.path.join(save_directory, self.config_name)
+
+ self.to_json_file(output_config_file)
+ logger.info(f"Configuration saved in {output_config_file}")
+
+ if push_to_hub:
+ commit_message = kwargs.pop("commit_message", None)
+ private = kwargs.pop("private", None)
+ create_pr = kwargs.pop("create_pr", False)
+ token = kwargs.pop("token", None)
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+ repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
+ subfolder = kwargs.pop("subfolder", None)
+
+ self._upload_folder(
+ save_directory,
+ repo_id,
+ token=token,
+ commit_message=commit_message,
+ create_pr=create_pr,
+ subfolder=subfolder,
+ )
+
+ @classmethod
+ def from_config(
+ cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs
+ ) -> Self | tuple[Self, dict[str, Any]]:
+ r"""
+ Instantiate a Python class from a config dictionary.
+
+ Parameters:
+ config (`dict[str, Any]`):
+ A config dictionary from which the Python class is instantiated. Make sure to only load configuration
+ files of compatible classes.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ Whether kwargs that are not consumed by the Python class should be returned or not.
+ kwargs (remaining dictionary of keyword arguments, *optional*):
+ Can be used to update the configuration object (after it is loaded) and initiate the Python class.
+ `**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
+ overwrite the same named arguments in `config`.
+
+ Returns:
+ [`ModelMixin`] or [`SchedulerMixin`]:
+ A model or scheduler object instantiated from a config dictionary.
+
+ Examples:
+
+ ```python
+ >>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
+
+ >>> # Download scheduler from huggingface.co and cache.
+ >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
+
+ >>> # Instantiate DDIM scheduler class with same config as DDPM
+ >>> scheduler = DDIMScheduler.from_config(scheduler.config)
+
+ >>> # Instantiate PNDM scheduler class with same config as DDPM
+ >>> scheduler = PNDMScheduler.from_config(scheduler.config)
+ ```
+ """
+ # <===== TO BE REMOVED WITH DEPRECATION
+ # TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
+ if "pretrained_model_name_or_path" in kwargs:
+ config = kwargs.pop("pretrained_model_name_or_path")
+
+ if config is None:
+ raise ValueError("Please make sure to provide a config as the first positional argument.")
+ # ======>
+
+ if not isinstance(config, dict):
+ deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
+ if "Scheduler" in cls.__name__:
+ deprecation_message += (
+ f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
+ " Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
+ " be removed in v1.0.0."
+ )
+ elif "Model" in cls.__name__:
+ deprecation_message += (
+ f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
+ f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
+ " instead. This functionality will be removed in v1.0.0."
+ )
+ deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
+ config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
+
+ init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
+
+ # Allow dtype to be specified on initialization
+ if "dtype" in unused_kwargs:
+ init_dict["dtype"] = unused_kwargs.pop("dtype")
+
+ # add possible deprecated kwargs
+ for deprecated_kwarg in cls._deprecated_kwargs:
+ if deprecated_kwarg in unused_kwargs:
+ init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
+
+ # Return model and optionally state and/or unused_kwargs
+ model = cls(**init_dict)
+
+ # make sure to also save config parameters that might be used for compatible classes
+ # update _class_name
+ if "_class_name" in hidden_dict:
+ hidden_dict["_class_name"] = cls.__name__
+
+ model.register_to_config(**hidden_dict)
+
+ # add hidden kwargs of compatible classes to unused_kwargs
+ unused_kwargs = {**unused_kwargs, **hidden_dict}
+
+ if return_unused_kwargs:
+ return (model, unused_kwargs)
+ else:
+ return model
+
+ @classmethod
+ def get_config_dict(cls, *args, **kwargs):
+ deprecation_message = (
+ f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
+ " removed in version v1.0.0"
+ )
+ deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
+ return cls.load_config(*args, **kwargs)
+
+ @classmethod
+ @validate_hf_hub_args
+ def load_config(
+ cls,
+ pretrained_model_name_or_path: str | os.PathLike,
+ return_unused_kwargs=False,
+ return_commit_hash=False,
+ **kwargs,
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
+ r"""
+ Load a model or scheduler configuration.
+
+ Parameters:
+ pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
+ Can be either:
+
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
+ the Hub.
+ - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
+ [`~ConfigMixin.save_config`].
+
+ cache_dir (`str | os.PathLike`, *optional*):
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
+ is not used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+ cached versions if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ output_loading_info(`bool`, *optional*, defaults to `False`):
+ Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ Whether to only load local model weights and configuration files or not. If set to `True`, the model
+ won't be downloaded from the Hub.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
+ allowed by Git.
+ subfolder (`str`, *optional*, defaults to `""`):
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False):
+ Whether unused keyword arguments of the config are returned.
+ return_commit_hash (`bool`, *optional*, defaults to `False):
+ Whether the `commit_hash` of the loaded configuration are returned.
+
+ Returns:
+ `dict`:
+ A dictionary of all the parameters stored in a JSON configuration file.
+
+ """
+ cache_dir = kwargs.pop("cache_dir", None)
+ local_dir = kwargs.pop("local_dir", None)
+ local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
+ force_download = kwargs.pop("force_download", False)
+ proxies = kwargs.pop("proxies", None)
+ token = kwargs.pop("token", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ revision = kwargs.pop("revision", None)
+ _ = kwargs.pop("mirror", None)
+ subfolder = kwargs.pop("subfolder", None)
+ user_agent = kwargs.pop("user_agent", {})
+ dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
+
+ user_agent = {**user_agent, "file_type": "config"}
+ user_agent = http_user_agent(user_agent)
+
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+
+ if cls.config_name is None:
+ raise ValueError(
+ "`self.config_name` is not defined. Note that one should not load a config from "
+ "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
+ )
+ # Custom path for now
+ if dduf_entries:
+ if subfolder is not None:
+ raise ValueError(
+ "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). "
+ "Please check the DDUF structure"
+ )
+ config_file = cls._get_config_file_from_dduf(pretrained_model_name_or_path, dduf_entries)
+ elif os.path.isfile(pretrained_model_name_or_path):
+ config_file = pretrained_model_name_or_path
+ elif os.path.isdir(pretrained_model_name_or_path):
+ if subfolder is not None and os.path.isfile(
+ os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
+ ):
+ config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
+ elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
+ # Load from a PyTorch checkpoint
+ config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
+ else:
+ raise EnvironmentError(
+ f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
+ )
+ else:
+ try:
+ # Load from URL or cache if already cached
+ config_file = hf_hub_download(
+ pretrained_model_name_or_path,
+ filename=cls.config_name,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ subfolder=subfolder,
+ revision=revision,
+ local_dir=local_dir,
+ local_dir_use_symlinks=local_dir_use_symlinks,
+ )
+ except RepositoryNotFoundError:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
+ " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
+ " token having permission to this repo with `token` or log in with `hf auth login`."
+ )
+ except RevisionNotFoundError:
+ raise EnvironmentError(
+ f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
+ " this model name. Check the model page at"
+ f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
+ )
+ except EntryNotFoundError:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
+ )
+ except HfHubHTTPError as err:
+ raise EnvironmentError(
+ "There was a specific connection error when trying to load"
+ f" {pretrained_model_name_or_path}:\n{err}"
+ )
+ except ValueError:
+ raise EnvironmentError(
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
+ f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
+ f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
+ " run the library in offline mode at"
+ " 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
+ )
+ except EnvironmentError:
+ raise EnvironmentError(
+ f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
+ "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
+ f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
+ f"containing a {cls.config_name} file"
+ )
+ try:
+ config_dict = cls._dict_from_json_file(config_file, dduf_entries=dduf_entries)
+
+ commit_hash = extract_commit_hash(config_file)
+ except (json.JSONDecodeError, UnicodeDecodeError):
+ raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.")
+
+ if not (return_unused_kwargs or return_commit_hash):
+ return config_dict
+
+ outputs = (config_dict,)
+
+ if return_unused_kwargs:
+ outputs += (kwargs,)
+
+ if return_commit_hash:
+ outputs += (commit_hash,)
+
+ return outputs
+
+ @staticmethod
+ def _get_init_keys(input_class):
+ return set(dict(inspect.signature(input_class.__init__).parameters).keys())
+
+ @classmethod
+ def extract_init_dict(cls, config_dict, **kwargs):
+ # Skip keys that were not present in the original config, so default __init__ values were used
+ used_defaults = config_dict.get("_use_default_values", [])
+ config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"}
+
+ # 0. Copy origin config dict
+ original_dict = dict(config_dict.items())
+
+ # 1. Retrieve expected config attributes from __init__ signature
+ expected_keys = cls._get_init_keys(cls)
+ expected_keys.remove("self")
+ # remove general kwargs if present in dict
+ if "kwargs" in expected_keys:
+ expected_keys.remove("kwargs")
+
+ # 2. Remove attributes that cannot be expected from expected config attributes
+ # remove keys to be ignored
+ if len(cls.ignore_for_config) > 0:
+ expected_keys = expected_keys - set(cls.ignore_for_config)
+
+ # load diffusers library to import compatible and original scheduler
+ diffusers_library = importlib.import_module(__name__.split(".")[0])
+
+ if cls.has_compatibles:
+ compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]
+ else:
+ compatible_classes = []
+
+ expected_keys_comp_cls = set()
+ for c in compatible_classes:
+ expected_keys_c = cls._get_init_keys(c)
+ expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c)
+ expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls)
+ config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls}
+
+ # remove attributes from orig class that cannot be expected
+ orig_cls_name = config_dict.pop("_class_name", cls.__name__)
+ if (
+ isinstance(orig_cls_name, str)
+ and orig_cls_name != cls.__name__
+ and hasattr(diffusers_library, orig_cls_name)
+ ):
+ orig_cls = getattr(diffusers_library, orig_cls_name)
+ unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys
+ config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig}
+ elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)):
+ raise ValueError(
+ "Make sure that the `_class_name` is of type string or list of string (for custom pipelines)."
+ )
+
+ # remove private attributes
+ config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")}
+
+ # remove quantization_config
+ config_dict = {k: v for k, v in config_dict.items() if k != "quantization_config"}
+
+ # 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments
+ init_dict = {}
+ for key in expected_keys:
+ # if config param is passed to kwarg and is present in config dict
+ # it should overwrite existing config dict key
+ if key in kwargs and key in config_dict:
+ config_dict[key] = kwargs.pop(key)
+
+ if key in kwargs:
+ # overwrite key
+ init_dict[key] = kwargs.pop(key)
+ elif key in config_dict:
+ # use value from config dict
+ init_dict[key] = config_dict.pop(key)
+
+ # 4. Give nice warning if unexpected values have been passed
+ if len(config_dict) > 0:
+ logger.warning(
+ f"The config attributes {config_dict} were passed to {cls.__name__}, "
+ "but are not expected and will be ignored. Please verify your "
+ f"{cls.config_name} configuration file."
+ )
+
+ # 5. Give nice info if config attributes are initialized to default because they have not been passed
+ passed_keys = set(init_dict.keys())
+ if len(expected_keys - passed_keys) > 0:
+ logger.info(
+ f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
+ )
+
+ # 6. Define unused keyword arguments
+ unused_kwargs = {**config_dict, **kwargs}
+
+ # 7. Define "hidden" config parameters that were saved for compatible classes
+ hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict}
+
+ return init_dict, unused_kwargs, hidden_config_dict
+
+ @classmethod
+ def _dict_from_json_file(cls, json_file: str | os.PathLike, dduf_entries: dict[str, DDUFEntry] | None = None):
+ if dduf_entries:
+ text = dduf_entries[json_file].read_text()
+ else:
+ with open(json_file, "r", encoding="utf-8") as reader:
+ text = reader.read()
+ return json.loads(text)
+
+ def __repr__(self):
+ return f"{self.__class__.__name__} {self.to_json_string()}"
+
+ @property
+ def config(self) -> dict[str, Any]:
+ """
+ Returns the config of the class as a frozen dictionary
+
+ Returns:
+ `dict[str, Any]`: Config of the class.
+ """
+ return self._internal_dict
+
+ def to_json_string(self) -> str:
+ """
+ Serializes the configuration instance to a JSON string.
+
+ Returns:
+ `str`:
+ String containing all the attributes that make up the configuration instance in JSON format.
+ """
+ config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
+ config_dict["_class_name"] = self.__class__.__name__
+ config_dict["_diffusers_version"] = __version__
+
+ def to_json_saveable(value):
+ if isinstance(value, np.ndarray):
+ value = value.tolist()
+ elif isinstance(value, Path):
+ value = value.as_posix()
+ elif hasattr(value, "to_dict") and callable(value.to_dict):
+ value = value.to_dict()
+ elif isinstance(value, list):
+ value = [to_json_saveable(v) for v in value]
+ return value
+
+ if "quantization_config" in config_dict:
+ config_dict["quantization_config"] = (
+ config_dict.quantization_config.to_dict()
+ if not isinstance(config_dict.quantization_config, dict)
+ else config_dict.quantization_config
+ )
+
+ config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
+ # Don't save "_ignore_files" or "_use_default_values"
+ config_dict.pop("_ignore_files", None)
+ config_dict.pop("_use_default_values", None)
+ # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable.
+ _ = config_dict.pop("_pre_quantization_dtype", None)
+
+ if getattr(self, "_auto_class", None) is not None:
+ module = self.__class__.__module__.split(".")[-1]
+ auto_map = config_dict.get("auto_map", {})
+ auto_map[self._auto_class] = f"{module}.{self.__class__.__name__}"
+ config_dict["auto_map"] = auto_map
+
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+ def to_json_file(self, json_file_path: str | os.PathLike):
+ """
+ Save the configuration instance's parameters to a JSON file.
+
+ Args:
+ json_file_path (`str` or `os.PathLike`):
+ Path to the JSON file to save a configuration instance's parameters.
+ """
+ with open(json_file_path, "w", encoding="utf-8") as writer:
+ writer.write(self.to_json_string())
+
+ @classmethod
+ def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: dict[str, DDUFEntry]):
+ # paths inside a DDUF file must always be "/"
+ config_file = (
+ cls.config_name
+ if pretrained_model_name_or_path == ""
+ else "/".join([pretrained_model_name_or_path, cls.config_name])
+ )
+ if config_file not in dduf_entries:
+ raise ValueError(
+ f"We did not manage to find the file {config_file} in the dduf file. We only have the following files {dduf_entries.keys()}"
+ )
+ return config_file
+
+
+def register_to_config(init):
+ r"""
+ Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are
+ automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that
+ shouldn't be registered in the config, use the `ignore_for_config` class variable
+
+ Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init!
+ """
+
+ @functools.wraps(init)
+ def inner_init(self, *args, **kwargs):
+ # Ignore private kwargs in the init.
+ init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
+ config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")}
+ if not isinstance(self, ConfigMixin):
+ raise RuntimeError(
+ f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
+ "not inherit from `ConfigMixin`."
+ )
+
+ ignore = getattr(self, "ignore_for_config", [])
+ # Get positional arguments aligned with kwargs
+ new_kwargs = {}
+ signature = inspect.signature(init)
+ parameters = {
+ name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore
+ }
+ for arg, name in zip(args, parameters.keys()):
+ new_kwargs[name] = arg
+
+ # Then add all kwargs
+ new_kwargs.update(
+ {
+ k: init_kwargs.get(k, default)
+ for k, default in parameters.items()
+ if k not in ignore and k not in new_kwargs
+ }
+ )
+
+ # Take note of the parameters that were not present in the loaded config
+ if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
+ new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
+
+ new_kwargs = {**config_init_kwargs, **new_kwargs}
+ getattr(self, "register_to_config")(**new_kwargs)
+ init(self, *args, **init_kwargs)
+
+ return inner_init
+
+
+class LegacyConfigMixin(ConfigMixin):
+ r"""
+ A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
+ pipeline-specific classes (like `DiTTransformer2DModel`).
+ """
+
+ @classmethod
+ def from_config(cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs):
+ # To prevent dependency import problem.
+ from .models.model_loading_utils import _fetch_remapped_cls_from_config
+
+ # resolve remapping
+ remapped_class = _fetch_remapped_cls_from_config(config, cls)
+
+ if remapped_class is cls:
+ return super(LegacyConfigMixin, remapped_class).from_config(config, return_unused_kwargs, **kwargs)
+ else:
+ return remapped_class.from_config(config, return_unused_kwargs, **kwargs)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py
new file mode 100644
index 0000000000000000000000000000000000000000..65364f9db25d14ec81b40247805fbd49df7a370c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py
@@ -0,0 +1,34 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .dependency_versions_table import deps
+from .utils.versions import require_version, require_version_core
+
+
+# define which module versions we always want to check at run time
+# (usually the ones defined in `install_requires` in setup.py)
+#
+# order specific notes:
+# - tqdm must be checked before tokenizers
+
+pkgs_to_check_at_runtime = "python requests filelock numpy".split()
+for pkg in pkgs_to_check_at_runtime:
+ if pkg in deps:
+ require_version_core(deps[pkg])
+ else:
+ raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
+
+
+def dep_version_check(pkg, hint=None):
+ require_version(deps[pkg], hint)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4f9522ede317a0cb6eb7e4e208f7c50be7b2cb9
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py
@@ -0,0 +1,57 @@
+# THIS FILE HAS BEEN AUTOGENERATED. To update:
+# 1. modify the `_deps` dict in setup.py
+# 2. run `make deps_table_update`
+deps = {
+ "Pillow": "Pillow",
+ "accelerate": "accelerate>=0.31.0",
+ "datasets": "datasets",
+ "filelock": "filelock",
+ "ftfy": "ftfy",
+ "hf-doc-builder": "hf-doc-builder>=0.3.0",
+ "httpx": "httpx<1.0.0",
+ "huggingface-hub": "huggingface-hub>=1.23.0,<2.0",
+ "requests-mock": "requests-mock==1.10.0",
+ "importlib_metadata": "importlib_metadata",
+ "invisible-watermark": "invisible-watermark>=0.2.0",
+ "isort": "isort>=5.5.4",
+ "Jinja2": "Jinja2",
+ "torchsde": "torchsde",
+ "note_seq": "note_seq",
+ "librosa": "librosa",
+ "llvmlite": "llvmlite>=0.40.0",
+ "numba": "numba>=0.57.0",
+ "numpy": "numpy",
+ "parameterized": "parameterized",
+ "peft": "peft>=0.17.0",
+ "protobuf": "protobuf>=3.20.3,<4",
+ "pytest": "pytest",
+ "pytest-timeout": "pytest-timeout",
+ "pytest-xdist": "pytest-xdist",
+ "python": "python>=3.10.0",
+ "ruff": "ruff==0.9.10",
+ "safetensors": "safetensors>=0.8.0",
+ "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
+ "GitPython": "GitPython<3.1.19",
+ "scipy": "scipy",
+ "onnx": "onnx",
+ "optimum_quanto": "optimum_quanto>=0.2.6",
+ "gguf": "gguf>=0.10.0",
+ "auto-round": "auto-round>=0.13.0",
+ "torchao": "torchao>=0.7.0",
+ "bitsandbytes": "bitsandbytes>=0.43.3",
+ "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1",
+ "sdnq": "sdnq>=0.2.2",
+ "regex": "regex!=2019.12.17",
+ "requests": "requests",
+ "tensorboard": "tensorboard",
+ "tiktoken": "tiktoken>=0.7.0",
+ "torch": "torch>=2.6",
+ "torchvision": "torchvision",
+ "transformers": "transformers>=4.41.2",
+ "urllib3": "urllib3<=2.0.0",
+ "black": "black",
+ "phonemizer": "phonemizer",
+ "opencv-python": "opencv-python",
+ "timm": "timm",
+ "flashpack": "flashpack",
+}
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..f917b172df5c2e3226afbfcd67ae22d0d213b18a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py
@@ -0,0 +1,1468 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import warnings
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageFilter, ImageOps
+
+from .configuration_utils import ConfigMixin, register_to_config
+from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate
+
+
+PipelineImageInput = (
+ PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image] | list[np.ndarray] | list[torch.Tensor]
+)
+
+PipelineDepthInput = PipelineImageInput
+
+
+def is_valid_image(image) -> bool:
+ r"""
+ Checks if the input is a valid image.
+
+ A valid image can be:
+ - A `PIL.Image.Image`.
+ - A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image).
+
+ Args:
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
+ The image to validate. It can be a PIL image, a NumPy array, or a torch tensor.
+
+ Returns:
+ `bool`:
+ `True` if the input is a valid image, `False` otherwise.
+ """
+ return isinstance(image, PIL.Image.Image) or isinstance(image, (np.ndarray, torch.Tensor)) and image.ndim in (2, 3)
+
+
+def is_valid_image_imagelist(images):
+ r"""
+ Checks if the input is a valid image or list of images.
+
+ The input can be one of the following formats:
+ - A 4D tensor or numpy array (batch of images).
+ - A valid single image: `PIL.Image.Image`, 2D `np.ndarray` or `torch.Tensor` (grayscale image), 3D `np.ndarray` or
+ `torch.Tensor`.
+ - A list of valid images.
+
+ Args:
+ images (`np.ndarray | torch.Tensor | PIL.Image.Image | list`):
+ The image(s) to check. Can be a batch of images (4D tensor/array), a single image, or a list of valid
+ images.
+
+ Returns:
+ `bool`:
+ `True` if the input is valid, `False` otherwise.
+ """
+ if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
+ return True
+ elif is_valid_image(images):
+ return True
+ elif isinstance(images, list):
+ return all(is_valid_image(image) for image in images)
+ return False
+
+
+class VaeImageProcessor(ConfigMixin):
+ """
+ Image processor for VAE.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
+ `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
+ resample (`str`, *optional*, defaults to `lanczos`):
+ Resampling filter to use when resizing the image.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether to normalize the image to [-1,1].
+ do_binarize (`bool`, *optional*, defaults to `False`):
+ Whether to binarize the image to 0/1.
+ do_convert_rgb (`bool`, *optional*, defaults to be `False`):
+ Whether to convert the images to RGB format.
+ do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
+ Whether to convert the images to grayscale format.
+ """
+
+ config_name = CONFIG_NAME
+
+ @register_to_config
+ def __init__(
+ self,
+ do_resize: bool = True,
+ vae_scale_factor: int = 8,
+ vae_latent_channels: int = 4,
+ resample: str = "lanczos",
+ reducing_gap: int | None = None,
+ do_normalize: bool = True,
+ do_binarize: bool = False,
+ do_convert_rgb: bool = False,
+ do_convert_grayscale: bool = False,
+ ):
+ super().__init__()
+ if do_convert_rgb and do_convert_grayscale:
+ raise ValueError(
+ "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`,"
+ " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.",
+ " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`",
+ )
+
+ @staticmethod
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
+ r"""
+ Convert a numpy image or a batch of images to a PIL image.
+
+ Args:
+ images (`np.ndarray`):
+ The image array to convert to PIL format.
+
+ Returns:
+ `list[PIL.Image.Image]`:
+ A list of PIL images.
+ """
+ if images.ndim == 3:
+ images = images[None, ...]
+ images = (images * 255).round().astype("uint8")
+ if images.shape[-1] == 1:
+ # special case for grayscale (single channel) images
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+ else:
+ pil_images = [Image.fromarray(image) for image in images]
+
+ return pil_images
+
+ @staticmethod
+ def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
+ r"""
+ Convert a PIL image or a list of PIL images to NumPy arrays.
+
+ Args:
+ images (`PIL.Image.Image` or `list[PIL.Image.Image]`):
+ The PIL image or list of images to convert to NumPy format.
+
+ Returns:
+ `np.ndarray`:
+ A NumPy array representation of the images.
+ """
+ if not isinstance(images, list):
+ images = [images]
+ images = [np.array(image).astype(np.float32) / 255.0 for image in images]
+ images = np.stack(images, axis=0)
+
+ return images
+
+ @staticmethod
+ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
+ r"""
+ Convert a NumPy image to a PyTorch tensor.
+
+ Args:
+ images (`np.ndarray`):
+ The NumPy image array to convert to PyTorch format.
+
+ Returns:
+ `torch.Tensor`:
+ A PyTorch tensor representation of the images.
+ """
+ if images.ndim == 3:
+ images = images[..., None]
+
+ images = torch.from_numpy(images.transpose(0, 3, 1, 2))
+ return images
+
+ @staticmethod
+ def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
+ r"""
+ Convert a PyTorch tensor to a NumPy image.
+
+ Args:
+ images (`torch.Tensor`):
+ The PyTorch tensor to convert to NumPy format.
+
+ Returns:
+ `np.ndarray`:
+ A NumPy array representation of the images.
+ """
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+ return images
+
+ @staticmethod
+ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
+ r"""
+ Normalize an image array to [-1,1].
+
+ Args:
+ images (`np.ndarray` or `torch.Tensor`):
+ The image array to normalize.
+
+ Returns:
+ `np.ndarray` or `torch.Tensor`:
+ The normalized image array.
+ """
+ return 2.0 * images - 1.0
+
+ @staticmethod
+ def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
+ r"""
+ Denormalize an image array to [0,1].
+
+ Args:
+ images (`np.ndarray` or `torch.Tensor`):
+ The image array to denormalize.
+
+ Returns:
+ `np.ndarray` or `torch.Tensor`:
+ The denormalized image array.
+ """
+ return (images * 0.5 + 0.5).clamp(0, 1)
+
+ @staticmethod
+ def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:
+ r"""
+ Converts a PIL image to RGB format.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The PIL image to convert to RGB.
+
+ Returns:
+ `PIL.Image.Image`:
+ The RGB-converted PIL image.
+ """
+ image = image.convert("RGB")
+
+ return image
+
+ @staticmethod
+ def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image:
+ r"""
+ Converts a given PIL image to grayscale.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The input image to convert.
+
+ Returns:
+ `PIL.Image.Image`:
+ The image converted to grayscale.
+ """
+ image = image.convert("L")
+
+ return image
+
+ @staticmethod
+ def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image:
+ r"""
+ Applies Gaussian blur to an image.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The PIL image to convert to grayscale.
+
+ Returns:
+ `PIL.Image.Image`:
+ The grayscale-converted PIL image.
+ """
+ image = image.filter(ImageFilter.GaussianBlur(blur_factor))
+
+ return image
+
+ @staticmethod
+ def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0):
+ r"""
+ Finds a rectangular region that contains all masked ares in an image, and expands region to match the aspect
+ ratio of the original image; for example, if user drew mask in a 128x32 region, and the dimensions for
+ processing are 512x512, the region will be expanded to 128x128.
+
+ Args:
+ mask_image (PIL.Image.Image): Mask image.
+ width (int): Width of the image to be processed.
+ height (int): Height of the image to be processed.
+ pad (int, optional): Padding to be added to the crop region. Defaults to 0.
+
+ Returns:
+ tuple: (x1, y1, x2, y2) represent a rectangular region that contains all masked ares in an image and
+ matches the original aspect ratio.
+ """
+
+ mask_image = mask_image.convert("L")
+ mask = np.array(mask_image)
+
+ # 1. find a rectangular region that contains all masked ares in an image
+ h, w = mask.shape
+ crop_left = 0
+ for i in range(w):
+ if not (mask[:, i] == 0).all():
+ break
+ crop_left += 1
+
+ crop_right = 0
+ for i in reversed(range(w)):
+ if not (mask[:, i] == 0).all():
+ break
+ crop_right += 1
+
+ crop_top = 0
+ for i in range(h):
+ if not (mask[i] == 0).all():
+ break
+ crop_top += 1
+
+ crop_bottom = 0
+ for i in reversed(range(h)):
+ if not (mask[i] == 0).all():
+ break
+ crop_bottom += 1
+
+ # 2. add padding to the crop region
+ x1, y1, x2, y2 = (
+ int(max(crop_left - pad, 0)),
+ int(max(crop_top - pad, 0)),
+ int(min(w - crop_right + pad, w)),
+ int(min(h - crop_bottom + pad, h)),
+ )
+
+ # 3. expands crop region to match the aspect ratio of the image to be processed
+ ratio_crop_region = (x2 - x1) / (y2 - y1)
+ ratio_processing = width / height
+
+ if ratio_crop_region > ratio_processing:
+ desired_height = (x2 - x1) / ratio_processing
+ desired_height_diff = int(desired_height - (y2 - y1))
+ y1 -= desired_height_diff // 2
+ y2 += desired_height_diff - desired_height_diff // 2
+ if y2 >= mask_image.height:
+ diff = y2 - mask_image.height
+ y2 -= diff
+ y1 -= diff
+ if y1 < 0:
+ y2 -= y1
+ y1 -= y1
+ if y2 >= mask_image.height:
+ y2 = mask_image.height
+ else:
+ desired_width = (y2 - y1) * ratio_processing
+ desired_width_diff = int(desired_width - (x2 - x1))
+ x1 -= desired_width_diff // 2
+ x2 += desired_width_diff - desired_width_diff // 2
+ if x2 >= mask_image.width:
+ diff = x2 - mask_image.width
+ x2 -= diff
+ x1 -= diff
+ if x1 < 0:
+ x2 -= x1
+ x1 -= x1
+ if x2 >= mask_image.width:
+ x2 = mask_image.width
+
+ return x1, y1, x2, y2
+
+ def _resize_and_fill(
+ self,
+ image: PIL.Image.Image,
+ width: int,
+ height: int,
+ ) -> PIL.Image.Image:
+ r"""
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
+ the image within the dimensions, filling empty with data from image.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The image to resize and fill.
+ width (`int`):
+ The width to resize the image to.
+ height (`int`):
+ The height to resize the image to.
+
+ Returns:
+ `PIL.Image.Image`:
+ The resized and filled image.
+ """
+
+ ratio = width / height
+ src_ratio = image.width / image.height
+
+ src_w = width if ratio < src_ratio else image.width * height // image.height
+ src_h = height if ratio >= src_ratio else image.height * width // image.width
+
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
+ res = Image.new("RGB", (width, height))
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
+
+ if ratio < src_ratio:
+ fill_height = height // 2 - src_h // 2
+ if fill_height > 0:
+ res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
+ res.paste(
+ resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)),
+ box=(0, fill_height + src_h),
+ )
+ elif ratio > src_ratio:
+ fill_width = width // 2 - src_w // 2
+ if fill_width > 0:
+ res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
+ res.paste(
+ resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)),
+ box=(fill_width + src_w, 0),
+ )
+
+ return res
+
+ def _resize_and_crop(
+ self,
+ image: PIL.Image.Image,
+ width: int,
+ height: int,
+ ) -> PIL.Image.Image:
+ r"""
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
+ the image within the dimensions, cropping the excess.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The image to resize and crop.
+ width (`int`):
+ The width to resize the image to.
+ height (`int`):
+ The height to resize the image to.
+
+ Returns:
+ `PIL.Image.Image`:
+ The resized and cropped image.
+ """
+ ratio = width / height
+ src_ratio = image.width / image.height
+
+ src_w = width if ratio > src_ratio else image.width * height // image.height
+ src_h = height if ratio <= src_ratio else image.height * width // image.width
+
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
+ res = Image.new("RGB", (width, height))
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
+ return res
+
+ def resize(
+ self,
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ height: int,
+ width: int,
+ resize_mode: str = "default", # "default", "fill", "crop"
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
+ """
+ Resize image.
+
+ Args:
+ image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
+ The image input, can be a PIL image, numpy array or pytorch tensor.
+ height (`int`):
+ The height to resize to.
+ width (`int`):
+ The width to resize to.
+ resize_mode (`str`, *optional*, defaults to `default`):
+ The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit
+ within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`,
+ will resize the image to fit within the specified width and height, maintaining the aspect ratio, and
+ then center the image within the dimensions, filling empty with data from image. If `crop`, will resize
+ the image to fit within the specified width and height, maintaining the aspect ratio, and then center
+ the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
+ supported for PIL image input.
+
+ Returns:
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
+ The resized image.
+ """
+ if resize_mode != "default" and not isinstance(image, PIL.Image.Image):
+ raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}")
+ if isinstance(image, PIL.Image.Image):
+ if resize_mode == "default":
+ image = image.resize(
+ (width, height),
+ resample=PIL_INTERPOLATION[self.config.resample],
+ reducing_gap=self.config.reducing_gap,
+ )
+ elif resize_mode == "fill":
+ image = self._resize_and_fill(image, width, height)
+ elif resize_mode == "crop":
+ image = self._resize_and_crop(image, width, height)
+ else:
+ raise ValueError(f"resize_mode {resize_mode} is not supported")
+
+ elif isinstance(image, torch.Tensor):
+ image = torch.nn.functional.interpolate(
+ image,
+ size=(height, width),
+ )
+ elif isinstance(image, np.ndarray):
+ image = self.numpy_to_pt(image)
+ image = torch.nn.functional.interpolate(
+ image,
+ size=(height, width),
+ )
+ image = self.pt_to_numpy(image)
+
+ return image
+
+ def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:
+ """
+ Create a mask.
+
+ Args:
+ image (`PIL.Image.Image`):
+ The image input, should be a PIL image.
+
+ Returns:
+ `PIL.Image.Image`:
+ The binarized image. Values less than 0.5 are set to 0, values greater than 0.5 are set to 1.
+ """
+ image[image < 0.5] = 0
+ image[image >= 0.5] = 1
+
+ return image
+
+ def _denormalize_conditionally(
+ self, images: torch.Tensor, do_denormalize: list[bool] | None = None
+ ) -> torch.Tensor:
+ r"""
+ Denormalize a batch of images based on a condition list.
+
+ Args:
+ images (`torch.Tensor`):
+ The input image tensor.
+ do_denormalize (`Optional[list[bool]`, *optional*, defaults to `None`):
+ A list of booleans indicating whether to denormalize each image in the batch. If `None`, will use the
+ value of `do_normalize` in the `VaeImageProcessor` config.
+ """
+ if do_denormalize is None:
+ return self.denormalize(images) if self.config.do_normalize else images
+
+ return torch.stack(
+ [self.denormalize(images[i]) if do_denormalize[i] else images[i] for i in range(images.shape[0])]
+ )
+
+ def get_default_height_width(
+ self,
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
+ height: int | None = None,
+ width: int | None = None,
+ ) -> tuple[int, int]:
+ r"""
+ Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
+
+ Args:
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
+ The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
+ should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
+ tensor, it should have shape `[batch, channels, height, width]`.
+ height (`int | None`, *optional*, defaults to `None`):
+ The height of the preprocessed image. If `None`, the height of the `image` input will be used.
+ width (`int | None`, *optional*, defaults to `None`):
+ The width of the preprocessed image. If `None`, the width of the `image` input will be used.
+
+ Returns:
+ `tuple[int, int]`:
+ A tuple containing the height and width, both resized to the nearest integer multiple of
+ `vae_scale_factor`.
+ """
+
+ if height is None:
+ if isinstance(image, PIL.Image.Image):
+ height = image.height
+ elif isinstance(image, torch.Tensor):
+ height = image.shape[2]
+ else:
+ height = image.shape[1]
+
+ if width is None:
+ if isinstance(image, PIL.Image.Image):
+ width = image.width
+ elif isinstance(image, torch.Tensor):
+ width = image.shape[3]
+ else:
+ width = image.shape[2]
+
+ width, height = (
+ x - x % self.config.vae_scale_factor for x in (width, height)
+ ) # resize to integer multiple of vae_scale_factor
+
+ return height, width
+
+ def preprocess(
+ self,
+ image: PipelineImageInput,
+ height: int | None = None,
+ width: int | None = None,
+ resize_mode: str = "default", # "default", "fill", "crop"
+ crops_coords: tuple[int, int, int, int] | None = None,
+ ) -> torch.Tensor:
+ """
+ Preprocess the image input.
+
+ Args:
+ image (`PipelineImageInput`):
+ The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of
+ supported formats.
+ height (`int`, *optional*):
+ The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default
+ height.
+ width (`int`, *optional*):
+ The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width.
+ resize_mode (`str`, *optional*, defaults to `default`):
+ The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within
+ the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will
+ resize the image to fit within the specified width and height, maintaining the aspect ratio, and then
+ center the image within the dimensions, filling empty with data from image. If `crop`, will resize the
+ image to fit within the specified width and height, maintaining the aspect ratio, and then center the
+ image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
+ supported for PIL image input.
+ crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`):
+ The crop coordinates for each image in the batch. If `None`, will not crop the image.
+
+ Returns:
+ `torch.Tensor`:
+ The preprocessed image.
+ """
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
+
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
+ if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
+ if isinstance(image, torch.Tensor):
+ # if image is a pytorch tensor could have 2 possible shapes:
+ # 1. batch x height x width: we should insert the channel dimension at position 1
+ # 2. channel x height x width: we should insert batch dimension at position 0,
+ # however, since both channel and batch dimension has same size 1, it is same to insert at position 1
+ # for simplicity, we insert a dimension of size 1 at position 1 for both cases
+ image = image.unsqueeze(1)
+ else:
+ # if it is a numpy array, it could have 2 possible shapes:
+ # 1. batch x height x width: insert channel dimension on last position
+ # 2. height x width x channel: insert batch dimension on first position
+ if image.shape[-1] == 1:
+ image = np.expand_dims(image, axis=0)
+ else:
+ image = np.expand_dims(image, axis=-1)
+
+ if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4:
+ warnings.warn(
+ "Passing `image` as a list of 4d np.ndarray is deprecated."
+ "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
+ FutureWarning,
+ )
+ image = np.concatenate(image, axis=0)
+ if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
+ warnings.warn(
+ "Passing `image` as a list of 4d torch.Tensor is deprecated."
+ "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
+ FutureWarning,
+ )
+ image = torch.cat(image, axis=0)
+
+ if not is_valid_image_imagelist(image):
+ raise ValueError(
+ f"Input is in incorrect format. Currently, we only support {', '.join(str(x) for x in supported_formats)}"
+ )
+ if not isinstance(image, list):
+ image = [image]
+
+ if isinstance(image[0], PIL.Image.Image):
+ if crops_coords is not None:
+ image = [i.crop(crops_coords) for i in image]
+ if self.config.do_resize:
+ height, width = self.get_default_height_width(image[0], height, width)
+ image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
+ if self.config.do_convert_rgb:
+ image = [self.convert_to_rgb(i) for i in image]
+ elif self.config.do_convert_grayscale:
+ image = [self.convert_to_grayscale(i) for i in image]
+ image = self.pil_to_numpy(image) # to np
+ image = self.numpy_to_pt(image) # to pt
+
+ elif isinstance(image[0], np.ndarray):
+ image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
+
+ image = self.numpy_to_pt(image)
+
+ height, width = self.get_default_height_width(image, height, width)
+ if self.config.do_resize:
+ image = self.resize(image, height, width)
+
+ elif isinstance(image[0], torch.Tensor):
+ image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
+
+ if self.config.do_convert_grayscale and image.ndim == 3:
+ image = image.unsqueeze(1)
+
+ channel = image.shape[1]
+ # don't need any preprocess if the image is latents
+ if channel == self.config.vae_latent_channels:
+ return image
+
+ height, width = self.get_default_height_width(image, height, width)
+ if self.config.do_resize:
+ image = self.resize(image, height, width)
+
+ # expected range [0,1], normalize to [-1,1]
+ do_normalize = self.config.do_normalize
+ if do_normalize and image.min() < 0:
+ warnings.warn(
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
+ FutureWarning,
+ )
+ do_normalize = False
+ if do_normalize:
+ image = self.normalize(image)
+
+ if self.config.do_binarize:
+ image = self.binarize(image)
+
+ return image
+
+ def postprocess(
+ self,
+ image: torch.Tensor,
+ output_type: str = "pil",
+ do_denormalize: list[bool] | None = None,
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
+ """
+ Postprocess the image output from tensor to `output_type`.
+
+ Args:
+ image (`torch.Tensor`):
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
+ output_type (`str`, *optional*, defaults to `pil`):
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
+ `VaeImageProcessor` config.
+
+ Returns:
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
+ The postprocessed image.
+ """
+ if not isinstance(image, torch.Tensor):
+ raise ValueError(
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
+ )
+ if output_type not in ["latent", "pt", "np", "pil"]:
+ deprecation_message = (
+ f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
+ "`pil`, `np`, `pt`, `latent`"
+ )
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
+ output_type = "np"
+
+ if output_type == "latent":
+ return image
+
+ image = self._denormalize_conditionally(image, do_denormalize)
+
+ if output_type == "pt":
+ return image
+
+ image = self.pt_to_numpy(image)
+
+ if output_type == "np":
+ return image
+
+ if output_type == "pil":
+ return self.numpy_to_pil(image)
+
+ def apply_overlay(
+ self,
+ mask: PIL.Image.Image,
+ init_image: PIL.Image.Image,
+ image: PIL.Image.Image,
+ crop_coords: tuple[int, int, int, int] | None = None,
+ ) -> PIL.Image.Image:
+ r"""
+ Applies an overlay of the mask and the inpainted image on the original image.
+
+ Args:
+ mask (`PIL.Image.Image`):
+ The mask image that highlights regions to overlay.
+ init_image (`PIL.Image.Image`):
+ The original image to which the overlay is applied.
+ image (`PIL.Image.Image`):
+ The image to overlay onto the original.
+ crop_coords (`tuple[int, int, int, int]`, *optional*):
+ Coordinates to crop the image. If provided, the image will be cropped accordingly.
+
+ Returns:
+ `PIL.Image.Image`:
+ The final image with the overlay applied.
+ """
+
+ width, height = init_image.width, init_image.height
+
+ init_image_masked = PIL.Image.new("RGBa", (width, height))
+ init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L")))
+
+ init_image_masked = init_image_masked.convert("RGBA")
+
+ if crop_coords is not None:
+ x, y, x2, y2 = crop_coords
+ w = x2 - x
+ h = y2 - y
+ base_image = PIL.Image.new("RGBA", (width, height))
+ image = self.resize(image, height=h, width=w, resize_mode="crop")
+ base_image.paste(image, (x, y))
+ image = base_image.convert("RGB")
+
+ image = image.convert("RGBA")
+ image.alpha_composite(init_image_masked)
+ image = image.convert("RGB")
+
+ return image
+
+
+class InpaintProcessor(ConfigMixin):
+ """
+ Image processor for inpainting image and mask.
+ """
+
+ config_name = CONFIG_NAME
+
+ @register_to_config
+ def __init__(
+ self,
+ do_resize: bool = True,
+ vae_scale_factor: int = 8,
+ vae_latent_channels: int = 4,
+ resample: str = "lanczos",
+ reducing_gap: int | None = None,
+ do_normalize: bool = True,
+ do_binarize: bool = False,
+ do_convert_grayscale: bool = False,
+ mask_do_normalize: bool = False,
+ mask_do_binarize: bool = True,
+ mask_do_convert_grayscale: bool = True,
+ ):
+ super().__init__()
+
+ self._image_processor = VaeImageProcessor(
+ do_resize=do_resize,
+ vae_scale_factor=vae_scale_factor,
+ vae_latent_channels=vae_latent_channels,
+ resample=resample,
+ reducing_gap=reducing_gap,
+ do_normalize=do_normalize,
+ do_binarize=do_binarize,
+ do_convert_grayscale=do_convert_grayscale,
+ )
+ self._mask_processor = VaeImageProcessor(
+ do_resize=do_resize,
+ vae_scale_factor=vae_scale_factor,
+ vae_latent_channels=vae_latent_channels,
+ resample=resample,
+ reducing_gap=reducing_gap,
+ do_normalize=mask_do_normalize,
+ do_binarize=mask_do_binarize,
+ do_convert_grayscale=mask_do_convert_grayscale,
+ )
+
+ def preprocess(
+ self,
+ image: PIL.Image.Image,
+ mask: PIL.Image.Image | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ padding_mask_crop: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Preprocess the image and mask.
+ """
+ if mask is None and padding_mask_crop is not None:
+ raise ValueError("mask must be provided if padding_mask_crop is provided")
+
+ # if mask is None, same behavior as regular image processor
+ if mask is None:
+ return self._image_processor.preprocess(image, height=height, width=width)
+
+ if padding_mask_crop is not None:
+ crops_coords = self._image_processor.get_crop_region(mask, width, height, pad=padding_mask_crop)
+ resize_mode = "fill"
+ else:
+ crops_coords = None
+ resize_mode = "default"
+
+ processed_image = self._image_processor.preprocess(
+ image,
+ height=height,
+ width=width,
+ crops_coords=crops_coords,
+ resize_mode=resize_mode,
+ )
+
+ processed_mask = self._mask_processor.preprocess(
+ mask,
+ height=height,
+ width=width,
+ resize_mode=resize_mode,
+ crops_coords=crops_coords,
+ )
+
+ if crops_coords is not None:
+ postprocessing_kwargs = {
+ "crops_coords": crops_coords,
+ "original_image": image,
+ "original_mask": mask,
+ }
+ else:
+ postprocessing_kwargs = {
+ "crops_coords": None,
+ "original_image": None,
+ "original_mask": None,
+ }
+
+ return processed_image, processed_mask, postprocessing_kwargs
+
+ def postprocess(
+ self,
+ image: torch.Tensor,
+ output_type: str = "pil",
+ original_image: PIL.Image.Image | None = None,
+ original_mask: PIL.Image.Image | None = None,
+ crops_coords: tuple[int, int, int, int] | None = None,
+ ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
+ """
+ Postprocess the image, optionally apply mask overlay
+ """
+ image = self._image_processor.postprocess(
+ image,
+ output_type=output_type,
+ )
+ # optionally apply the mask overlay
+ if crops_coords is not None and (original_image is None or original_mask is None):
+ raise ValueError("original_image and original_mask must be provided if crops_coords is provided")
+
+ elif crops_coords is not None and output_type != "pil":
+ raise ValueError("output_type must be 'pil' if crops_coords is provided")
+
+ elif crops_coords is not None:
+ image = [
+ self._image_processor.apply_overlay(original_mask, original_image, i, crops_coords) for i in image
+ ]
+
+ return image
+
+
+class VaeImageProcessorLDM3D(VaeImageProcessor):
+ """
+ Image processor for VAE LDM3D.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
+ resample (`str`, *optional*, defaults to `lanczos`):
+ Resampling filter to use when resizing the image.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether to normalize the image to [-1,1].
+ """
+
+ config_name = CONFIG_NAME
+
+ @register_to_config
+ def __init__(
+ self,
+ do_resize: bool = True,
+ vae_scale_factor: int = 8,
+ resample: str = "lanczos",
+ do_normalize: bool = True,
+ ):
+ super().__init__()
+
+ @staticmethod
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
+ r"""
+ Convert a NumPy image or a batch of images to a list of PIL images.
+
+ Args:
+ images (`np.ndarray`):
+ The input NumPy array of images, which can be a single image or a batch.
+
+ Returns:
+ `list[PIL.Image.Image]`:
+ A list of PIL images converted from the input NumPy array.
+ """
+ if images.ndim == 3:
+ images = images[None, ...]
+ images = (images * 255).round().astype("uint8")
+ if images.shape[-1] == 1:
+ # special case for grayscale (single channel) images
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+ else:
+ pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
+
+ return pil_images
+
+ @staticmethod
+ def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
+ r"""
+ Convert a PIL image or a list of PIL images to NumPy arrays.
+
+ Args:
+ images (`list[PIL.Image.Image, PIL.Image.Image]`):
+ The input image or list of images to be converted.
+
+ Returns:
+ `np.ndarray`:
+ A NumPy array of the converted images.
+ """
+ if not isinstance(images, list):
+ images = [images]
+
+ images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]
+ images = np.stack(images, axis=0)
+ return images
+
+ @staticmethod
+ def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
+ r"""
+ Convert an RGB-like depth image to a depth map.
+ """
+ # 1. Cast the tensor to a larger integer type (e.g., int32)
+ # to safely perform the multiplication by 256.
+ # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte.
+ # 3. Cast the final result to the desired depth map type (uint16) if needed
+ # before returning, though leaving it as int32/int64 is often safer
+ # for return value from a library function.
+
+ if isinstance(image, torch.Tensor):
+ # Cast to a safe dtype (e.g., int32 or int64) for the calculation
+ original_dtype = image.dtype
+ image_safe = image.to(torch.int32)
+
+ # Calculate the depth map
+ depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
+
+ # You may want to cast the final result to uint16, but casting to a
+ # larger int type (like int32) is sufficient to fix the overflow.
+ # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required
+ return depth_map.to(original_dtype)
+
+ elif isinstance(image, np.ndarray):
+ # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
+ original_dtype = image.dtype
+ image_safe = image.astype(np.int32)
+
+ # Calculate the depth map
+ depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
+
+ # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required
+ return depth_map.astype(original_dtype)
+ else:
+ raise TypeError("Input image must be a torch.Tensor or np.ndarray")
+
+ def numpy_to_depth(self, images: np.ndarray) -> list[PIL.Image.Image]:
+ r"""
+ Convert a NumPy depth image or a batch of images to a list of PIL images.
+
+ Args:
+ images (`np.ndarray`):
+ The input NumPy array of depth images, which can be a single image or a batch.
+
+ Returns:
+ `list[PIL.Image.Image]`:
+ A list of PIL images converted from the input NumPy depth images.
+ """
+ if images.ndim == 3:
+ images = images[None, ...]
+ images_depth = images[:, :, :, 3:]
+ if images.shape[-1] == 6:
+ images_depth = (images_depth * 255).round().astype("uint8")
+ pil_images = [
+ Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
+ ]
+ elif images.shape[-1] == 4:
+ images_depth = (images_depth * 65535.0).astype(np.uint16)
+ pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
+ else:
+ raise Exception("Not supported")
+
+ return pil_images
+
+ def postprocess(
+ self,
+ image: torch.Tensor,
+ output_type: str = "pil",
+ do_denormalize: list[bool] | None = None,
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
+ """
+ Postprocess the image output from tensor to `output_type`.
+
+ Args:
+ image (`torch.Tensor`):
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
+ output_type (`str`, *optional*, defaults to `pil`):
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
+ `VaeImageProcessor` config.
+
+ Returns:
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
+ The postprocessed image.
+ """
+ if not isinstance(image, torch.Tensor):
+ raise ValueError(
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
+ )
+ if output_type not in ["latent", "pt", "np", "pil"]:
+ deprecation_message = (
+ f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
+ "`pil`, `np`, `pt`, `latent`"
+ )
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
+ output_type = "np"
+
+ image = self._denormalize_conditionally(image, do_denormalize)
+
+ image = self.pt_to_numpy(image)
+
+ if output_type == "np":
+ if image.shape[-1] == 6:
+ image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
+ else:
+ image_depth = image[:, :, :, 3:]
+ return image[:, :, :, :3], image_depth
+
+ if output_type == "pil":
+ return self.numpy_to_pil(image), self.numpy_to_depth(image)
+ else:
+ raise Exception(f"This type {output_type} is not supported")
+
+ def preprocess(
+ self,
+ rgb: torch.Tensor | PIL.Image.Image | np.ndarray,
+ depth: torch.Tensor | PIL.Image.Image | np.ndarray,
+ height: int | None = None,
+ width: int | None = None,
+ target_res: int | None = None,
+ ) -> torch.Tensor:
+ r"""
+ Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors.
+
+ Args:
+ rgb (`torch.Tensor | PIL.Image.Image | np.ndarray`):
+ The RGB input image, which can be a single image or a batch.
+ depth (`torch.Tensor | PIL.Image.Image | np.ndarray`):
+ The depth input image, which can be a single image or a batch.
+ height (`int | None`, *optional*, defaults to `None`):
+ The desired height of the processed image. If `None`, defaults to the height of the input image.
+ width (`int | None`, *optional*, defaults to `None`):
+ The desired width of the processed image. If `None`, defaults to the width of the input image.
+ target_res (`int | None`, *optional*, defaults to `None`):
+ Target resolution for resizing the images. If specified, overrides height and width.
+
+ Returns:
+ `tuple[torch.Tensor, torch.Tensor]`:
+ A tuple containing the processed RGB and depth images as PyTorch tensors.
+ """
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
+
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
+ if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:
+ raise Exception("This is not yet supported")
+
+ if isinstance(rgb, supported_formats):
+ rgb = [rgb]
+ depth = [depth]
+ elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):
+ raise ValueError(
+ f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"
+ )
+
+ if isinstance(rgb[0], PIL.Image.Image):
+ if self.config.do_convert_rgb:
+ raise Exception("This is not yet supported")
+ # rgb = [self.convert_to_rgb(i) for i in rgb]
+ # depth = [self.convert_to_depth(i) for i in depth] #TODO define convert_to_depth
+ if self.config.do_resize or target_res:
+ height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res
+ rgb = [self.resize(i, height, width) for i in rgb]
+ depth = [self.resize(i, height, width) for i in depth]
+ rgb = self.pil_to_numpy(rgb) # to np
+ rgb = self.numpy_to_pt(rgb) # to pt
+
+ depth = self.depth_pil_to_numpy(depth) # to np
+ depth = self.numpy_to_pt(depth) # to pt
+
+ elif isinstance(rgb[0], np.ndarray):
+ rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)
+ rgb = self.numpy_to_pt(rgb)
+ height, width = self.get_default_height_width(rgb, height, width)
+ if self.config.do_resize:
+ rgb = self.resize(rgb, height, width)
+
+ depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0)
+ depth = self.numpy_to_pt(depth)
+ height, width = self.get_default_height_width(depth, height, width)
+ if self.config.do_resize:
+ depth = self.resize(depth, height, width)
+
+ elif isinstance(rgb[0], torch.Tensor):
+ raise Exception("This is not yet supported")
+ # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)
+
+ # if self.config.do_convert_grayscale and rgb.ndim == 3:
+ # rgb = rgb.unsqueeze(1)
+
+ # channel = rgb.shape[1]
+
+ # height, width = self.get_default_height_width(rgb, height, width)
+ # if self.config.do_resize:
+ # rgb = self.resize(rgb, height, width)
+
+ # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)
+
+ # if self.config.do_convert_grayscale and depth.ndim == 3:
+ # depth = depth.unsqueeze(1)
+
+ # channel = depth.shape[1]
+ # # don't need any preprocess if the image is latents
+ # if depth == 4:
+ # return rgb, depth
+
+ # height, width = self.get_default_height_width(depth, height, width)
+ # if self.config.do_resize:
+ # depth = self.resize(depth, height, width)
+ # expected range [0,1], normalize to [-1,1]
+ do_normalize = self.config.do_normalize
+ if rgb.min() < 0 and do_normalize:
+ warnings.warn(
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",
+ FutureWarning,
+ )
+ do_normalize = False
+
+ if do_normalize:
+ rgb = self.normalize(rgb)
+ depth = self.normalize(depth)
+
+ if self.config.do_binarize:
+ rgb = self.binarize(rgb)
+ depth = self.binarize(depth)
+
+ return rgb, depth
+
+
+class IPAdapterMaskProcessor(VaeImageProcessor):
+ """
+ Image processor for IP Adapter image masks.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
+ resample (`str`, *optional*, defaults to `lanczos`):
+ Resampling filter to use when resizing the image.
+ do_normalize (`bool`, *optional*, defaults to `False`):
+ Whether to normalize the image to [-1,1].
+ do_binarize (`bool`, *optional*, defaults to `True`):
+ Whether to binarize the image to 0/1.
+ do_convert_grayscale (`bool`, *optional*, defaults to be `True`):
+ Whether to convert the images to grayscale format.
+
+ """
+
+ config_name = CONFIG_NAME
+
+ @register_to_config
+ def __init__(
+ self,
+ do_resize: bool = True,
+ vae_scale_factor: int = 8,
+ resample: str = "lanczos",
+ do_normalize: bool = False,
+ do_binarize: bool = True,
+ do_convert_grayscale: bool = True,
+ ):
+ super().__init__(
+ do_resize=do_resize,
+ vae_scale_factor=vae_scale_factor,
+ resample=resample,
+ do_normalize=do_normalize,
+ do_binarize=do_binarize,
+ do_convert_grayscale=do_convert_grayscale,
+ )
+
+ @staticmethod
+ def downsample(mask: torch.Tensor, batch_size: int, num_queries: int, value_embed_dim: int):
+ """
+ Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the
+ aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.
+
+ Args:
+ mask (`torch.Tensor`):
+ The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
+ batch_size (`int`):
+ The batch size.
+ num_queries (`int`):
+ The number of queries.
+ value_embed_dim (`int`):
+ The dimensionality of the value embeddings.
+
+ Returns:
+ `torch.Tensor`:
+ The downsampled mask tensor.
+
+ """
+ o_h = mask.shape[1]
+ o_w = mask.shape[2]
+ ratio = o_w / o_h
+ mask_h = int(math.sqrt(num_queries / ratio))
+ mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
+ mask_w = num_queries // mask_h
+
+ mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)
+
+ # Repeat batch_size times
+ if mask_downsample.shape[0] < batch_size:
+ mask_downsample = mask_downsample.repeat(batch_size, 1, 1)
+
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)
+
+ downsampled_area = mask_h * mask_w
+ # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
+ # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
+ if downsampled_area < num_queries:
+ warnings.warn(
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
+ "Please update your masks or adjust the output size for optimal performance.",
+ UserWarning,
+ )
+ mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
+ # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
+ if downsampled_area > num_queries:
+ warnings.warn(
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
+ "Please update your masks or adjust the output size for optimal performance.",
+ UserWarning,
+ )
+ mask_downsample = mask_downsample[:, :num_queries]
+
+ # Repeat last dimension to match SDPA output shape
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
+ 1, 1, value_embed_dim
+ )
+
+ return mask_downsample
+
+
+class PixArtImageProcessor(VaeImageProcessor):
+ """
+ Image processor for PixArt image resize and crop.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
+ `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
+ resample (`str`, *optional*, defaults to `lanczos`):
+ Resampling filter to use when resizing the image.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether to normalize the image to [-1,1].
+ do_binarize (`bool`, *optional*, defaults to `False`):
+ Whether to binarize the image to 0/1.
+ do_convert_rgb (`bool`, *optional*, defaults to be `False`):
+ Whether to convert the images to RGB format.
+ do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
+ Whether to convert the images to grayscale format.
+ """
+
+ @register_to_config
+ def __init__(
+ self,
+ do_resize: bool = True,
+ vae_scale_factor: int = 8,
+ resample: str = "lanczos",
+ do_normalize: bool = True,
+ do_binarize: bool = False,
+ do_convert_grayscale: bool = False,
+ ):
+ super().__init__(
+ do_resize=do_resize,
+ vae_scale_factor=vae_scale_factor,
+ resample=resample,
+ do_normalize=do_normalize,
+ do_binarize=do_binarize,
+ do_convert_grayscale=do_convert_grayscale,
+ )
+
+ @staticmethod
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
+ r"""
+ Returns the binned height and width based on the aspect ratio.
+
+ Args:
+ height (`int`): The height of the image.
+ width (`int`): The width of the image.
+ ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
+
+ Returns:
+ `tuple[int, int]`: The closest binned height and width.
+ """
+ ar = float(height / width)
+ closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
+ default_hw = ratios[closest_ratio]
+ return int(default_hw[0]), int(default_hw[1])
+
+ @staticmethod
+ def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
+ r"""
+ Resizes and crops a tensor of images to the specified dimensions.
+
+ Args:
+ samples (`torch.Tensor`):
+ A tensor of shape (N, C, H, W) where N is the batch size, C is the number of channels, H is the height,
+ and W is the width.
+ new_width (`int`): The desired width of the output images.
+ new_height (`int`): The desired height of the output images.
+
+ Returns:
+ `torch.Tensor`: A tensor containing the resized and cropped images.
+ """
+ orig_height, orig_width = samples.shape[2], samples.shape[3]
+
+ # Check if resizing is needed
+ if orig_height != new_height or orig_width != new_width:
+ ratio = max(new_height / orig_height, new_width / orig_width)
+ resized_width = int(orig_width * ratio)
+ resized_height = int(orig_height * ratio)
+
+ # Resize
+ samples = F.interpolate(
+ samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
+ )
+
+ # Center Crop
+ start_x = (resized_width - new_width) // 2
+ end_x = start_x + new_width
+ start_y = (resized_height - new_height) // 2
+ end_y = start_y + new_height
+ samples = samples[:, :, start_y:end_y, start_x:end_x]
+
+ return samples
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3bca7741d1685fefe7ed26db48486a5af690a7be
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py
@@ -0,0 +1,230 @@
+from typing import TYPE_CHECKING
+
+from ..utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+ logging,
+)
+
+
+logger = logging.get_logger(__name__)
+logger.warning(
+ "Modular Diffusers is currently an experimental feature under active development. The API is subject to breaking changes in future releases."
+)
+
+# These modules contain pipelines from multiple libraries/frameworks
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ..utils import dummy_pt_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_pt_objects))
+else:
+ _import_structure["modular_pipeline"] = [
+ "ModularPipelineBlocks",
+ "ModularPipeline",
+ "AutoPipelineBlocks",
+ "SequentialPipelineBlocks",
+ "ConditionalPipelineBlocks",
+ "LoopSequentialPipelineBlocks",
+ "PipelineState",
+ "BlockState",
+ ]
+ _import_structure["modular_pipeline_utils"] = [
+ "ComponentSpec",
+ "ConfigSpec",
+ "InputParam",
+ "OutputParam",
+ "InsertableDict",
+ ]
+ _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"]
+ _import_structure["stable_diffusion_3"] = ["StableDiffusion3AutoBlocks", "StableDiffusion3ModularPipeline"]
+ _import_structure["wan"] = [
+ "WanBlocks",
+ "Wan22Blocks",
+ "WanImage2VideoAutoBlocks",
+ "Wan22Image2VideoBlocks",
+ "WanModularPipeline",
+ "Wan22ModularPipeline",
+ "WanImage2VideoModularPipeline",
+ "Wan22Image2VideoModularPipeline",
+ ]
+ _import_structure["helios"] = [
+ "HeliosAutoBlocks",
+ "HeliosModularPipeline",
+ "HeliosPyramidAutoBlocks",
+ "HeliosPyramidDistilledAutoBlocks",
+ "HeliosPyramidDistilledModularPipeline",
+ "HeliosPyramidModularPipeline",
+ ]
+ _import_structure["flux"] = [
+ "FluxAutoBlocks",
+ "FluxModularPipeline",
+ "FluxKontextAutoBlocks",
+ "FluxKontextModularPipeline",
+ ]
+ _import_structure["flux2"] = [
+ "Flux2AutoBlocks",
+ "Flux2KleinAutoBlocks",
+ "Flux2KleinBaseAutoBlocks",
+ "Flux2ModularPipeline",
+ "Flux2KleinModularPipeline",
+ "Flux2KleinBaseModularPipeline",
+ ]
+ _import_structure["ideogram4"] = [
+ "Ideogram4AutoBlocks",
+ "Ideogram4ModularPipeline",
+ ]
+ _import_structure["krea2"] = [
+ "Krea2AutoBlocks",
+ "Krea2ModularPipeline",
+ "Krea2TurboAutoBlocks",
+ "Krea2TurboModularPipeline",
+ ]
+ _import_structure["qwenimage"] = [
+ "QwenImageAutoBlocks",
+ "QwenImageModularPipeline",
+ "QwenImageEditModularPipeline",
+ "QwenImageEditAutoBlocks",
+ "QwenImageEditPlusModularPipeline",
+ "QwenImageEditPlusAutoBlocks",
+ "QwenImageLayeredModularPipeline",
+ "QwenImageLayeredAutoBlocks",
+ ]
+ _import_structure["anima"] = [
+ "AnimaAutoBlocks",
+ "AnimaModularPipeline",
+ ]
+ _import_structure["cosmos"] = [
+ "Cosmos3DistilledBlocks",
+ "Cosmos3DistilledModularPipeline",
+ "Cosmos3OmniBlocks",
+ "Cosmos3OmniModularPipeline",
+ ]
+ _import_structure["ernie_image"] = [
+ "ErnieImageAutoBlocks",
+ "ErnieImageModularPipeline",
+ ]
+ _import_structure["hunyuan_video1_5"] = [
+ "HunyuanVideo15AutoBlocks",
+ "HunyuanVideo15ModularPipeline",
+ ]
+ _import_structure["ltx"] = [
+ "LTXAutoBlocks",
+ "LTXModularPipeline",
+ ]
+ _import_structure["minimax_h3"] = [
+ "MiniMaxH3Blocks",
+ "MiniMaxH3ModularPipeline",
+ ]
+ _import_structure["z_image"] = [
+ "ZImageAutoBlocks",
+ "ZImageModularPipeline",
+ ]
+ _import_structure["components_manager"] = ["ComponentsManager"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not is_torch_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ..utils.dummy_pt_objects import * # noqa F403
+ else:
+ from .anima import AnimaAutoBlocks, AnimaModularPipeline
+ from .components_manager import ComponentsManager
+ from .cosmos import (
+ Cosmos3DistilledBlocks,
+ Cosmos3DistilledModularPipeline,
+ Cosmos3OmniBlocks,
+ Cosmos3OmniModularPipeline,
+ )
+ from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline
+ from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
+ from .flux2 import (
+ Flux2AutoBlocks,
+ Flux2KleinAutoBlocks,
+ Flux2KleinBaseAutoBlocks,
+ Flux2KleinBaseModularPipeline,
+ Flux2KleinModularPipeline,
+ Flux2ModularPipeline,
+ )
+ from .helios import (
+ HeliosAutoBlocks,
+ HeliosModularPipeline,
+ HeliosPyramidAutoBlocks,
+ HeliosPyramidDistilledAutoBlocks,
+ HeliosPyramidDistilledModularPipeline,
+ HeliosPyramidModularPipeline,
+ )
+ from .hunyuan_video1_5 import (
+ HunyuanVideo15AutoBlocks,
+ HunyuanVideo15ModularPipeline,
+ )
+ from .ideogram4 import (
+ Ideogram4AutoBlocks,
+ Ideogram4ModularPipeline,
+ )
+ from .krea2 import (
+ Krea2AutoBlocks,
+ Krea2ModularPipeline,
+ Krea2TurboAutoBlocks,
+ Krea2TurboModularPipeline,
+ )
+ from .ltx import LTXAutoBlocks, LTXModularPipeline
+ from .minimax_h3 import (
+ MiniMaxH3Blocks,
+ MiniMaxH3ModularPipeline,
+ )
+ from .modular_pipeline import (
+ AutoPipelineBlocks,
+ BlockState,
+ ConditionalPipelineBlocks,
+ LoopSequentialPipelineBlocks,
+ ModularPipeline,
+ ModularPipelineBlocks,
+ PipelineState,
+ SequentialPipelineBlocks,
+ )
+ from .modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam
+ from .qwenimage import (
+ QwenImageAutoBlocks,
+ QwenImageEditAutoBlocks,
+ QwenImageEditModularPipeline,
+ QwenImageEditPlusAutoBlocks,
+ QwenImageEditPlusModularPipeline,
+ QwenImageLayeredAutoBlocks,
+ QwenImageLayeredModularPipeline,
+ QwenImageModularPipeline,
+ )
+ from .stable_diffusion_3 import StableDiffusion3AutoBlocks, StableDiffusion3ModularPipeline
+ from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline
+ from .wan import (
+ Wan22Blocks,
+ Wan22Image2VideoBlocks,
+ Wan22Image2VideoModularPipeline,
+ Wan22ModularPipeline,
+ WanBlocks,
+ WanImage2VideoAutoBlocks,
+ WanImage2VideoModularPipeline,
+ WanModularPipeline,
+ )
+ from .z_image import ZImageAutoBlocks, ZImageModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..09b2b96369ae0a0222c16ababdbded135cdd7a20
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_anima"] = ["AnimaAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["AnimaModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_anima import AnimaAutoBlocks
+ from .modular_pipeline import AnimaModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..16f7ee609009f2ac84f04ed047d1494eff5e5d70
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py
@@ -0,0 +1,714 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...models import AnimaTextConditioner, CosmosTransformer3DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import AnimaModularPipeline
+
+
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# Copied from diffusers.modular_pipelines.z_image.before_denoise.repeat_tensor_to_batch_size
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_images_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_images_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_images_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_images_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_images_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+class AnimaTextConditioningStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Map Qwen text encoder states and T5 token ids to Cosmos text conditioning for Anima."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_conditioner", AnimaTextConditioner),
+ ComponentSpec("transformer", CosmosTransformer3DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "qwen_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Qwen prompt embeddings generated by the text encoder step.",
+ ),
+ InputParam(
+ "qwen_attention_mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Qwen prompt attention mask generated by the text encoder step.",
+ ),
+ InputParam(
+ "t5_input_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="T5 prompt token ids generated by the text encoder step.",
+ ),
+ InputParam(
+ "t5_attention_mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="T5 prompt attention mask generated by the text encoder step.",
+ ),
+ InputParam(
+ "negative_qwen_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Negative Qwen prompt embeddings generated by the text encoder step.",
+ ),
+ InputParam(
+ "negative_qwen_attention_mask",
+ type_hint=torch.Tensor,
+ description="Negative Qwen prompt attention mask generated by the text encoder step.",
+ ),
+ InputParam(
+ "negative_t5_input_ids",
+ type_hint=torch.Tensor,
+ description="Negative T5 prompt token ids generated by the text encoder step.",
+ ),
+ InputParam(
+ "negative_t5_attention_mask",
+ type_hint=torch.Tensor,
+ description="Negative T5 prompt attention mask generated by the text encoder step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Conditioned prompt embeddings generated by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Conditioned negative prompt embeddings generated by the Anima text conditioner.",
+ ),
+ ]
+
+ @staticmethod
+ def _condition_prompt_embeds(
+ components: AnimaModularPipeline,
+ qwen_prompt_embeds: torch.Tensor,
+ qwen_attention_mask: torch.Tensor,
+ t5_input_ids: torch.Tensor,
+ t5_attention_mask: torch.Tensor,
+ device: torch.device,
+ conditioning_dtype: torch.dtype,
+ output_dtype: torch.dtype,
+ ) -> torch.Tensor:
+ prompt_embeds = components.text_conditioner(
+ source_hidden_states=qwen_prompt_embeds.to(device=device, dtype=conditioning_dtype),
+ target_input_ids=t5_input_ids.to(device),
+ target_attention_mask=t5_attention_mask.to(device),
+ source_attention_mask=qwen_attention_mask.to(device),
+ )
+ return prompt_embeds.to(dtype=output_dtype, device=device)
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ conditioning_dtype = components.text_conditioner.dtype
+ output_dtype = components.transformer.dtype
+
+ block_state.prompt_embeds = self._condition_prompt_embeds(
+ components,
+ qwen_prompt_embeds=block_state.qwen_prompt_embeds,
+ qwen_attention_mask=block_state.qwen_attention_mask,
+ t5_input_ids=block_state.t5_input_ids,
+ t5_attention_mask=block_state.t5_attention_mask,
+ device=device,
+ conditioning_dtype=conditioning_dtype,
+ output_dtype=output_dtype,
+ )
+
+ block_state.negative_prompt_embeds = None
+ if block_state.negative_qwen_prompt_embeds is not None:
+ block_state.negative_prompt_embeds = self._condition_prompt_embeds(
+ components,
+ qwen_prompt_embeds=block_state.negative_qwen_prompt_embeds,
+ qwen_attention_mask=block_state.negative_qwen_attention_mask,
+ t5_input_ids=block_state.negative_t5_input_ids,
+ t5_attention_mask=block_state.negative_t5_attention_mask,
+ device=device,
+ conditioning_dtype=conditioning_dtype,
+ output_dtype=output_dtype,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaTextInputStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Input processing step that expands Anima prompt embeddings for the requested image batch."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", CosmosTransformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt"),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditioned prompt embeddings generated by the Anima text conditioner.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Conditioned negative prompt embeddings generated by the Anima text conditioner.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Prompt embeddings expanded to the final denoising batch.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Negative prompt embeddings expanded to the final denoising batch.",
+ ),
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
+ ),
+ OutputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = components.transformer.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaImageInputStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that expands Anima image latents to the final denoising batch "
+ "and derives height/width from the latents when not provided."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image_latents"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
+ ),
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Image latents expanded to the final denoising batch.",
+ ),
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latent_height, latent_width = block_state.image_latents.shape[-2:]
+ block_state.height = block_state.height or latent_height * components.vae_scale_factor
+ block_state.width = block_state.width or latent_width * components.vae_scale_factor
+
+ block_state.image_latents = repeat_tensor_to_batch_size(
+ input_name="image_latents",
+ input_tensor=block_state.image_latents,
+ batch_size=block_state.batch_size,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Prepare noisy image latents and padding mask for Anima denoising."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", CosmosTransformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("latents"),
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy latents for the denoising process."),
+ OutputParam("padding_mask", type_hint=torch.Tensor, description="Cosmos padding mask for image latents."),
+ ]
+
+ def check_inputs(self, components: AnimaModularPipeline, block_state):
+ divisor = components.vae_scale_factor * 2
+ if block_state.height % divisor != 0 or block_state.width % divisor != 0:
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {divisor} but are {block_state.height} and"
+ f" {block_state.width}."
+ )
+
+ @staticmethod
+ def prepare_latents(
+ batch_size: int,
+ num_channels_latents: int,
+ height: int,
+ width: int,
+ vae_scale_factor: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ generator: torch.Generator | list[torch.Generator] | None,
+ latents: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ if latents is not None:
+ return latents.to(device=device, dtype=dtype)
+
+ latent_height = height // vae_scale_factor
+ latent_width = width // vae_scale_factor
+ shape = (batch_size, num_channels_latents, 1, latent_height, latent_width)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+ self.check_inputs(components, block_state)
+
+ device = components._execution_device
+ block_state.latents = self.prepare_latents(
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
+ num_channels_latents=components.num_channels_latents,
+ height=block_state.height,
+ width=block_state.width,
+ vae_scale_factor=components.vae_scale_factor,
+ dtype=torch.float32,
+ device=device,
+ generator=block_state.generator,
+ latents=block_state.latents,
+ )
+ block_state.padding_mask = block_state.latents.new_zeros(
+ 1, 1, block_state.height, block_state.width, dtype=block_state.dtype
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.before_denoise.get_timesteps
+def get_timesteps(scheduler, num_inference_steps, strength):
+ # get the original timestep using init_timestep
+ init_timestep = min(num_inference_steps * strength, num_inference_steps)
+
+ t_start = int(max(num_inference_steps - init_timestep, 0))
+ timesteps = scheduler.timesteps[t_start * scheduler.order :]
+ if hasattr(scheduler, "set_begin_index"):
+ scheduler.set_begin_index(t_start * scheduler.order)
+
+ return timesteps, num_inference_steps - t_start
+
+
+class AnimaSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Set the scheduler timesteps for Anima inference."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for the denoising loop."),
+ OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ sigmas = (
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+ if block_state.sigmas is None
+ else block_state.sigmas
+ )
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ device=device,
+ sigmas=sigmas,
+ )
+ components.scheduler.set_begin_index(0)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaImg2ImgSetTimestepsStep(ModularPipelineBlocks):
+ """Set the scheduler timesteps for Anima image-to-image inference.
+
+ This step computes the full timestep schedule, then slices it based on ``strength`` via ``get_timesteps()``, which
+ also sets the scheduler's begin index.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ How much to transform the reference image.
+
+ Outputs:
+ timesteps (`Tensor`):
+ Timestep schedule sliced by ``strength``.
+ num_inference_steps (`int`):
+ Number of denoising steps after strength-based slicing.
+ """
+
+ model_name = "anima"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Set the scheduler timesteps for Anima image-to-image inference, sliced by strength."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ InputParam.template("strength"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "timesteps",
+ type_hint=torch.Tensor,
+ description="Timestep schedule sliced by strength.",
+ ),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="Number of denoising steps after strength-based slicing.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ sigmas = (
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+ if block_state.sigmas is None
+ else block_state.sigmas
+ )
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ device=device,
+ sigmas=sigmas,
+ )
+ block_state.timesteps, block_state.num_inference_steps = get_timesteps(
+ components.scheduler, block_state.num_inference_steps, block_state.strength
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaImg2ImgPrepareLatentsStep(ModularPipelineBlocks):
+ """Prepares noisy latents for Anima image-to-image generation.
+
+ Generates noise and mixes it with the image latents via ``scheduler.scale_noise()`` at the first sliced timestep.
+ The image latents are expected to already be expanded to the final batch size by ``AnimaImageInputStep``.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ image_latents (`Tensor`):
+ Encoded image latents, expanded to the final denoising batch.
+ timesteps (`Tensor`):
+ Timestep schedule sliced by ``strength`` from ``AnimaImg2ImgSetTimestepsStep``.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`, *optional*):
+ Pre-computed noise tensor. Generated randomly if ``None``.
+ dtype (`torch.dtype`):
+ Dtype used by the Anima denoiser.
+ height (`int`):
+ Image height.
+ width (`int`):
+ Image width.
+
+ Outputs:
+ latents (`Tensor`):
+ Noisy image latents for the denoising loop.
+ padding_mask (`Tensor`):
+ Cosmos padding mask for the image latents.
+ """
+
+ model_name = "anima"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepares noisy image-to-image latents for Anima by adding noise to the encoded "
+ "image latents via scheduler.scale_noise()."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image_latents"),
+ InputParam.template("timesteps", required=True),
+ InputParam.template("generator"),
+ InputParam.template("latents"),
+ InputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy latents for the denoising loop."),
+ OutputParam("padding_mask", type_hint=torch.Tensor, description="Cosmos padding mask for image latents."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ image_latents = block_state.image_latents.to(device=device, dtype=torch.float32)
+
+ if block_state.latents is None:
+ noise = randn_tensor(
+ image_latents.shape,
+ generator=block_state.generator,
+ device=device,
+ dtype=torch.float32,
+ )
+ else:
+ noise = block_state.latents.to(device=device, dtype=torch.float32)
+
+ latent_timestep = block_state.timesteps[:1].repeat(image_latents.shape[0])
+ block_state.latents = components.scheduler.scale_noise(image_latents, latent_timestep, noise)
+
+ block_state.padding_mask = block_state.latents.new_zeros(
+ 1, 1, block_state.height, block_state.width, dtype=block_state.dtype
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ba0371d45af0acd908650fe41367e4d27a10893
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py
@@ -0,0 +1,120 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKLQwenImage
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import AnimaModularPipeline
+
+
+class AnimaVaeDecoderStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes Anima latents into image tensors."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKLQwenImage)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latents", required=True, type_hint=torch.Tensor, description="Denoised Anima latents."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images", note="tensor output of the VAE decoder")]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latents = block_state.latents.to(components.vae.dtype)
+ latents_mean = (
+ torch.tensor(components.vae.config.latents_mean)
+ .view(1, components.vae.config.z_dim, 1, 1, 1)
+ .to(latents.device, latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+ 1, components.vae.config.z_dim, 1, 1, 1
+ ).to(latents.device, latents.dtype)
+ latents = latents / latents_std + latents_mean
+
+ block_state.images = components.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaProcessImagesOutputStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Postprocess decoded Anima image tensors."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("images", required=True, type_hint=torch.Tensor, description="Decoded Anima image tensors."),
+ InputParam.template("output_type"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | np.ndarray | torch.Tensor,
+ description="Generated images.",
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(output_type):
+ if output_type not in ["pil", "np", "pt"]:
+ raise ValueError(f"Invalid output_type: {output_type}")
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state.output_type)
+
+ block_state.images = components.image_processor.postprocess(
+ image=block_state.images,
+ output_type=block_state.output_type,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..51fdb8d7a707d5bb50e86e5939ec3cc4736c9ce2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py
@@ -0,0 +1,211 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import CosmosTransformer3DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam
+from .modular_pipeline import AnimaModularPipeline
+
+
+class AnimaLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that prepares Anima latent and timestep inputs."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latents", required=True, type_hint=torch.Tensor, description="Current Anima latents."),
+ InputParam("dtype", required=True, type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = block_state.latents.to(block_state.dtype)
+
+ timestep = t.expand(block_state.latents.shape[0]).to(block_state.dtype)
+ block_state.timestep = timestep / components.scheduler.config.num_train_timesteps
+ return components, block_state
+
+
+class AnimaLoopDenoiser(ModularPipelineBlocks):
+ model_name = "anima"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] | None = None,
+ ):
+ if guider_input_fields is None:
+ guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"`guider_input_fields` must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", CosmosTransformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that predicts Anima noise with guidance."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="Number of denoising steps.",
+ ),
+ InputParam(
+ "padding_mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Cosmos padding mask for image latents.",
+ ),
+ InputParam(
+ kwargs_type="denoiser_input_fields",
+ description="The conditional model inputs for the Anima denoiser.",
+ ),
+ ]
+
+ guider_input_names = []
+ uncond_guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.append(value[0])
+ uncond_guider_input_names.append(value[1])
+ else:
+ guider_input_names.append(value)
+
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True))
+ for name in uncond_guider_input_names:
+ inputs.append(InputParam(name=name))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {
+ key: getattr(guider_state_batch, key).to(block_state.dtype) for key in self._guider_input_fields.keys()
+ }
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=block_state.timestep,
+ padding_mask=block_state.padding_mask,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state)[0]
+ return components, block_state
+
+
+class AnimaLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that updates Anima latents."
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred, t, block_state.latents, return_dict=False
+ )[0]
+ if block_state.latents.dtype != latents_dtype and torch.backends.mps.is_available():
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class AnimaDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Pipeline block that iteratively denoises Anima latents over scheduler timesteps."
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam("timesteps", required=True, type_hint=torch.Tensor, description="Timesteps to denoise over."),
+ InputParam("num_inference_steps", required=True, type_hint=int, description="Number of denoising steps."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ num_warmup_steps = len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class AnimaDenoiseStep(AnimaDenoiseLoopWrapper):
+ block_classes = [
+ AnimaLoopBeforeDenoiser,
+ AnimaLoopDenoiser(guider_input_fields={"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}),
+ AnimaLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return "Denoise step that iteratively denoises image latents for Anima."
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d79ddc640a587de9674eecf1735d652eff92518
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py
@@ -0,0 +1,404 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from transformers import Qwen2Tokenizer, Qwen3Model, T5TokenizerFast
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKLQwenImage
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import AnimaModularPipeline
+
+
+class AnimaTextEncoderStep(ModularPipelineBlocks):
+ model_name = "anima"
+
+ @property
+ def description(self) -> str:
+ return "Text encoder step that encodes Anima prompts into Qwen states and T5 token ids."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3Model),
+ ComponentSpec("tokenizer", Qwen2Tokenizer),
+ ComponentSpec("t5_tokenizer", T5TokenizerFast),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam.template("max_sequence_length"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "qwen_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Qwen prompt embeddings to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "qwen_attention_mask",
+ type_hint=torch.Tensor,
+ description="Qwen prompt attention mask to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "t5_input_ids",
+ type_hint=torch.Tensor,
+ description="T5 prompt token ids to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "t5_attention_mask",
+ type_hint=torch.Tensor,
+ description="T5 prompt attention mask to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "negative_qwen_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Negative Qwen prompt embeddings to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "negative_qwen_attention_mask",
+ type_hint=torch.Tensor,
+ description="Negative Qwen prompt attention mask to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "negative_t5_input_ids",
+ type_hint=torch.Tensor,
+ description="Negative T5 prompt token ids to be consumed by the Anima text conditioner.",
+ ),
+ OutputParam(
+ "negative_t5_attention_mask",
+ type_hint=torch.Tensor,
+ description="Negative T5 prompt attention mask to be consumed by the Anima text conditioner.",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ if not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+ if block_state.max_sequence_length is not None and block_state.max_sequence_length > 4096:
+ raise ValueError(
+ f"`max_sequence_length` cannot be greater than 4096 but is {block_state.max_sequence_length}"
+ )
+
+ @staticmethod
+ def _get_qwen_prompt_embeds(
+ components: AnimaModularPipeline,
+ prompt: str | list[str],
+ max_sequence_length: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ text_inputs = components.tokenizer(
+ prompt,
+ padding="longest",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids.to(device)
+ prompt_attention_mask = text_inputs.attention_mask.to(device)
+ if text_input_ids.shape[-1] == 0:
+ text_input_ids = text_input_ids.new_zeros((text_input_ids.shape[0], 1))
+ prompt_attention_mask = prompt_attention_mask.new_zeros((prompt_attention_mask.shape[0], 1))
+
+ prompt_embeds = components.text_encoder(
+ input_ids=text_input_ids,
+ attention_mask=prompt_attention_mask,
+ output_hidden_states=False,
+ ).last_hidden_state
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+ prompt_embeds = prompt_embeds * prompt_attention_mask.to(prompt_embeds).unsqueeze(-1)
+
+ return prompt_embeds, prompt_attention_mask
+
+ @staticmethod
+ def _get_t5_prompt_ids(
+ components: AnimaModularPipeline,
+ prompt: str | list[str],
+ max_sequence_length: int,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ text_inputs = components.t5_tokenizer(
+ prompt,
+ padding="longest",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+ return text_inputs.input_ids.to(device), text_inputs.attention_mask.to(device)
+
+ @classmethod
+ def encode_prompt(
+ cls,
+ components: AnimaModularPipeline,
+ prompt: str | list[str],
+ negative_prompt: str | list[str] | None = None,
+ prepare_unconditional_embeds: bool = True,
+ max_sequence_length: int = 512,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ) -> dict[str, torch.Tensor | None]:
+ device = device or components._execution_device
+ dtype = dtype or components.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ prompt_embeds, prompt_attention_mask = cls._get_qwen_prompt_embeds(
+ components=components,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+ t5_input_ids, t5_attention_mask = cls._get_t5_prompt_ids(
+ components=components,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ negative_prompt_embeds = None
+ negative_prompt_attention_mask = None
+ negative_t5_input_ids = None
+ negative_t5_attention_mask = None
+ if prepare_unconditional_embeds:
+ negative_prompt = negative_prompt if negative_prompt is not None else ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ if batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds, negative_prompt_attention_mask = cls._get_qwen_prompt_embeds(
+ components=components,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+ negative_t5_input_ids, negative_t5_attention_mask = cls._get_t5_prompt_ids(
+ components=components,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ return {
+ "qwen_prompt_embeds": prompt_embeds,
+ "qwen_attention_mask": prompt_attention_mask,
+ "t5_input_ids": t5_input_ids,
+ "t5_attention_mask": t5_attention_mask,
+ "negative_qwen_prompt_embeds": negative_prompt_embeds,
+ "negative_qwen_attention_mask": negative_prompt_attention_mask,
+ "negative_t5_input_ids": negative_t5_input_ids,
+ "negative_t5_attention_mask": negative_t5_attention_mask,
+ }
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ prompt_outputs = self.encode_prompt(
+ components=components,
+ prompt=block_state.prompt,
+ negative_prompt=block_state.negative_prompt,
+ prepare_unconditional_embeds=components.guider.num_conditions > 1,
+ max_sequence_length=block_state.max_sequence_length,
+ device=components._execution_device,
+ dtype=components.text_encoder.dtype,
+ )
+ for name, value in prompt_outputs.items():
+ setattr(block_state, name, value)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.encoders.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.encoders.encode_vae_image
+def encode_vae_image(
+ image: torch.Tensor,
+ vae: AutoencoderKLQwenImage,
+ generator: torch.Generator,
+ device: torch.device,
+ dtype: torch.dtype,
+ latent_channels: int = 16,
+ sample_mode: str = "argmax",
+):
+ if not isinstance(image, torch.Tensor):
+ raise ValueError(f"Expected image to be a tensor, got {type(image)}.")
+
+ # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width
+ if image.dim() == 4:
+ image = image.unsqueeze(2)
+ elif image.dim() != 5:
+ raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.")
+
+ image = image.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode)
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean)
+ .view(1, latent_channels, 1, 1, 1)
+ .to(image_latents.device, image_latents.dtype)
+ )
+ latents_std = (
+ torch.tensor(vae.config.latents_std)
+ .view(1, latent_channels, 1, 1, 1)
+ .to(image_latents.device, image_latents.dtype)
+ )
+ image_latents = (image_latents - latents_mean) / latents_std
+
+ return image_latents
+
+
+class AnimaImg2ImgVaeEncoderStep(ModularPipelineBlocks):
+ """VAE Encoder step for Anima image-to-image generation.
+
+ Preprocesses the input image and encodes it with the VAE, producing ``image_latents``. Timestep slicing is handled
+ downstream by ``AnimaImg2ImgSetTimestepsStep`` and noise addition by ``AnimaImg2ImgPrepareLatentsStep``.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`PIL.Image.Image`):
+ Input image to encode.
+ height (`int`, *optional*):
+ Height of the output image. Defaults to pipeline default.
+ width (`int`, *optional*):
+ Width of the output image. Defaults to pipeline default.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ image_latents (`Tensor`):
+ Encoded image latents.
+ height (`int`):
+ Output image height.
+ width (`int`):
+ Output image width.
+ """
+
+ model_name = "anima"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLQwenImage),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "VAE Encoder step for Anima image-to-image generation. Encodes the input image to produce image_latents."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("image_latents", type_hint=torch.Tensor, description="Encoded image latents."),
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ processed_image = components.image_processor.preprocess(
+ image=block_state.image, height=block_state.height, width=block_state.width
+ )
+
+ block_state.image_latents = encode_vae_image(
+ image=processed_image,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=components.vae.dtype,
+ latent_channels=components.num_channels_latents,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e90c5959bef49f0ee4405c7489dd42dff1fa286
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py
@@ -0,0 +1,381 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ AnimaImageInputStep,
+ AnimaImg2ImgPrepareLatentsStep,
+ AnimaImg2ImgSetTimestepsStep,
+ AnimaPrepareLatentsStep,
+ AnimaSetTimestepsStep,
+ AnimaTextConditioningStep,
+ AnimaTextInputStep,
+)
+from .decoders import AnimaProcessImagesOutputStep, AnimaVaeDecoderStep
+from .denoise import AnimaDenoiseStep
+from .encoders import AnimaImg2ImgVaeEncoderStep, AnimaTextEncoderStep
+
+
+# auto_docstring
+class AnimaCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block that takes encoded Anima text inputs and runs the denoising process.
+
+ Components:
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ qwen_prompt_embeds (`Tensor`):
+ Qwen prompt embeddings generated by the text encoder step.
+ qwen_attention_mask (`Tensor`):
+ Qwen prompt attention mask generated by the text encoder step.
+ t5_input_ids (`Tensor`):
+ T5 prompt token ids generated by the text encoder step.
+ t5_attention_mask (`Tensor`):
+ T5 prompt attention mask generated by the text encoder step.
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
+ Negative Qwen prompt embeddings generated by the text encoder step.
+ negative_qwen_attention_mask (`Tensor`, *optional*):
+ Negative Qwen prompt attention mask generated by the text encoder step.
+ negative_t5_input_ids (`Tensor`, *optional*):
+ Negative T5 prompt token ids generated by the text encoder step.
+ negative_t5_attention_mask (`Tensor`, *optional*):
+ Negative T5 prompt attention mask generated by the text encoder step.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the Anima denoiser.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ AnimaTextConditioningStep,
+ AnimaTextInputStep,
+ AnimaPrepareLatentsStep,
+ AnimaSetTimestepsStep,
+ AnimaDenoiseStep,
+ ]
+ block_names = ["text_conditioning", "input", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def description(self) -> str:
+ return "Denoise block that takes encoded Anima text inputs and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class AnimaDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode Anima latents into generated images.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ Denoised Anima latents.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ block_classes = [AnimaVaeDecoderStep, AnimaProcessImagesOutputStep]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self) -> str:
+ return "Decode Anima latents into generated images."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
+
+
+# auto_docstring
+class AnimaImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block for Anima image-to-image generation. Uses image_latents already in state from
+ AnimaImg2ImgVaeEncoderStep.
+
+ Components:
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ qwen_prompt_embeds (`Tensor`):
+ Qwen prompt embeddings generated by the text encoder step.
+ qwen_attention_mask (`Tensor`):
+ Qwen prompt attention mask generated by the text encoder step.
+ t5_input_ids (`Tensor`):
+ T5 prompt token ids generated by the text encoder step.
+ t5_attention_mask (`Tensor`):
+ T5 prompt attention mask generated by the text encoder step.
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
+ Negative Qwen prompt embeddings generated by the text encoder step.
+ negative_qwen_attention_mask (`Tensor`, *optional*):
+ Negative Qwen prompt attention mask generated by the text encoder step.
+ negative_t5_input_ids (`Tensor`, *optional*):
+ Negative T5 prompt token ids generated by the text encoder step.
+ negative_t5_attention_mask (`Tensor`, *optional*):
+ Negative T5 prompt attention mask generated by the text encoder step.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the Anima denoiser.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ AnimaTextConditioningStep,
+ AnimaTextInputStep,
+ AnimaImageInputStep,
+ AnimaImg2ImgSetTimestepsStep,
+ AnimaImg2ImgPrepareLatentsStep,
+ AnimaDenoiseStep,
+ ]
+ block_names = ["text_conditioning", "input", "image_input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise block for Anima image-to-image generation. "
+ "Uses image_latents already in state from AnimaImg2ImgVaeEncoderStep."
+ )
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class AnimaAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Denoise step that selects between text-to-image and image-to-image denoising based on whether image_latents is
+ present in state. - `AnimaCoreDenoiseStep` (text2image) is used when no image_latents are present. -
+ `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when image_latents are present.
+
+ Components:
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ qwen_prompt_embeds (`Tensor`):
+ Qwen prompt embeddings generated by the text encoder step.
+ qwen_attention_mask (`Tensor`):
+ Qwen prompt attention mask generated by the text encoder step.
+ t5_input_ids (`Tensor`):
+ T5 prompt token ids generated by the text encoder step.
+ t5_attention_mask (`Tensor`):
+ T5 prompt attention mask generated by the text encoder step.
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
+ Negative Qwen prompt embeddings generated by the text encoder step.
+ negative_qwen_attention_mask (`Tensor`, *optional*):
+ Negative Qwen prompt attention mask generated by the text encoder step.
+ negative_t5_input_ids (`Tensor`, *optional*):
+ Negative T5 prompt token ids generated by the text encoder step.
+ negative_t5_attention_mask (`Tensor`, *optional*):
+ Negative T5 prompt attention mask generated by the text encoder step.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the Anima denoiser.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [AnimaImg2ImgCoreDenoiseStep, AnimaCoreDenoiseStep]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that selects between text-to-image and image-to-image denoising based on whether "
+ "image_latents is present in state."
+ " - `AnimaCoreDenoiseStep` (text2image) is used when no image_latents are present."
+ " - `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when image_latents are present."
+ )
+
+
+# auto_docstring
+class AnimaAutoVaeImageEncoderStep(AutoPipelineBlocks):
+ """
+ VAE Image Encoder step that encodes the input image to produce image_latents. Skipped when no image is provided
+ (text-to-image workflow).
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ image_latents (`Tensor`):
+ Encoded image latents.
+ height (`int`):
+ Image height used for generation.
+ width (`int`):
+ Image width used for generation.
+ """
+
+ block_classes = [AnimaImg2ImgVaeEncoderStep]
+ block_names = ["vae_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "VAE Image Encoder step that encodes the input image to produce image_latents. "
+ "Skipped when no image is provided (text-to-image workflow)."
+ )
+
+
+# auto_docstring
+class AnimaAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-to-image generation using Anima.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `img2img`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) t5_tokenizer (`T5Tokenizer`) guider
+ (`ClassifierFreeGuidance`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the Anima denoiser.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ block_classes = [
+ AnimaTextEncoderStep,
+ AnimaAutoVaeImageEncoderStep,
+ AnimaAutoCoreDenoiseStep,
+ AnimaDecodeStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "img2img": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return "Auto Modular pipeline for text-to-image and image-to-image generation using Anima."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c34d1548e9c25f30773e1bf119e3b93f60c7c59
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py
@@ -0,0 +1,52 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...loaders import AnimaLoraLoaderMixin
+from ..modular_pipeline import ModularPipeline
+
+
+class AnimaModularPipeline(ModularPipeline, AnimaLoraLoaderMixin):
+ """
+ A ModularPipeline for Anima.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "AnimaAutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ return 128
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if self.vae is not None:
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels
+ return num_channels_latents
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..89fb1129cf86e5504b8b3aabed130bbc29e6753f
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py
@@ -0,0 +1,1109 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+import time
+from collections import OrderedDict
+from itertools import combinations
+from typing import Any
+
+import torch
+
+from ..hooks import ModelHook
+from ..utils import (
+ is_accelerate_available,
+ logging,
+)
+from ..utils.torch_utils import get_device
+
+
+if is_accelerate_available():
+ from accelerate.hooks import add_hook_to_module, remove_hook_from_module
+ from accelerate.state import PartialState
+ from accelerate.utils import send_to_device
+ from accelerate.utils.memory import clear_device_cache
+ from accelerate.utils.modeling import convert_file_size_to_int
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class CustomOffloadHook(ModelHook):
+ """
+ A hook that offloads a model on the CPU until its forward pass is called. It ensures the model and its inputs are
+ on the given device. Optionally offloads other models to the CPU before the forward pass is called.
+
+ Args:
+ execution_device(`str`, `int` or `torch.device`, *optional*):
+ The device on which the model should be executed. Will default to the MPS device if it's available, then
+ GPU 0 if there is a GPU, and finally to the CPU.
+ """
+
+ no_grad = False
+
+ def __init__(
+ self,
+ execution_device: str | int | torch.device | None = None,
+ other_hooks: list["UserCustomOffloadHook"] | None = None,
+ offload_strategy: "AutoOffloadStrategy" | None = None,
+ ):
+ self.execution_device = execution_device if execution_device is not None else PartialState().default_device
+ self.other_hooks = other_hooks
+ self.offload_strategy = offload_strategy
+ self.model_id = None
+
+ def set_strategy(self, offload_strategy: "AutoOffloadStrategy"):
+ self.offload_strategy = offload_strategy
+
+ def add_other_hook(self, hook: "UserCustomOffloadHook"):
+ """
+ Add a hook to the list of hooks to consider for offloading.
+ """
+ if self.other_hooks is None:
+ self.other_hooks = []
+ self.other_hooks.append(hook)
+
+ def init_hook(self, module):
+ return module.to("cpu")
+
+ def pre_forward(self, module, *args, **kwargs):
+ if module.device != self.execution_device:
+ if self.other_hooks is not None:
+ hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device]
+ # offload all other hooks
+ start_time = time.perf_counter()
+ if self.offload_strategy is not None:
+ hooks_to_offload = self.offload_strategy(
+ hooks=hooks_to_offload,
+ model_id=self.model_id,
+ model=module,
+ execution_device=self.execution_device,
+ )
+ end_time = time.perf_counter()
+ logger.info(
+ f" time taken to apply offload strategy for {self.model_id}: {(end_time - start_time):.2f} seconds"
+ )
+
+ for hook in hooks_to_offload:
+ logger.info(
+ f"moving {self.model_id} to {self.execution_device}, offloading {hook.model_id} to cpu"
+ )
+ hook.offload()
+
+ if hooks_to_offload:
+ clear_device_cache()
+ module.to(self.execution_device)
+ return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
+
+
+class UserCustomOffloadHook:
+ """
+ A simple hook grouping a model and a `CustomOffloadHook`, which provides easy APIs for to call the init method of
+ the hook or remove it entirely.
+ """
+
+ def __init__(self, model_id, model, hook):
+ self.model_id = model_id
+ self.model = model
+ self.hook = hook
+
+ def offload(self):
+ self.hook.init_hook(self.model)
+
+ def attach(self):
+ add_hook_to_module(self.model, self.hook)
+ self.hook.model_id = self.model_id
+
+ def remove(self):
+ remove_hook_from_module(self.model)
+ self.hook.model_id = None
+
+ def add_other_hook(self, hook: "UserCustomOffloadHook"):
+ self.hook.add_other_hook(hook)
+
+
+def custom_offload_with_hook(
+ model_id: str,
+ model: torch.nn.Module,
+ execution_device: str | int | torch.device = None,
+ offload_strategy: "AutoOffloadStrategy" | None = None,
+):
+ hook = CustomOffloadHook(execution_device=execution_device, offload_strategy=offload_strategy)
+ user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook)
+ user_hook.attach()
+ return user_hook
+
+
+# this is the class that user can customize to implement their own offload strategy
+class AutoOffloadStrategy:
+ """
+ Offload strategy that should be used with `CustomOffloadHook` to automatically offload models to the CPU based on
+ the available memory on the device.
+ """
+
+ # YiYi TODO: instead of memory_reserve_margin, we should let user set the maximum_total_models_size to keep on device
+ # the actual memory usage would be higher. But it's simpler this way, and can be tested
+ def __init__(self, memory_reserve_margin="3GB"):
+ self.memory_reserve_margin = convert_file_size_to_int(memory_reserve_margin)
+
+ def __call__(self, hooks, model_id, model, execution_device):
+ if len(hooks) == 0:
+ return []
+
+ try:
+ current_module_size = model.get_memory_footprint()
+ except AttributeError:
+ raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.")
+
+ device_type = execution_device.type
+ device_module = getattr(torch, device_type, torch.cuda)
+ try:
+ mem_on_device = device_module.mem_get_info(execution_device.index)[0]
+ except AttributeError:
+ raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.")
+
+ mem_on_device = mem_on_device - self.memory_reserve_margin
+ if current_module_size < mem_on_device:
+ return []
+
+ min_memory_offload = current_module_size - mem_on_device
+ logger.info(f" search for models to offload in order to free up {min_memory_offload / 1024**3:.2f} GB memory")
+
+ # exlucde models that's not currently loaded on the device
+ module_sizes = dict(
+ sorted(
+ {hook.model_id: hook.model.get_memory_footprint() for hook in hooks}.items(),
+ key=lambda x: x[1],
+ reverse=True,
+ )
+ )
+
+ # YiYi/Dhruv TODO: sort smallest to largest, and offload in that order we would tend to keep the larger models on GPU more often
+ def search_best_candidate(module_sizes, min_memory_offload):
+ """
+ search the optimal combination of models to offload to cpu, given a dictionary of module sizes and a
+ minimum memory offload size. the combination of models should add up to the smallest modulesize that is
+ larger than `min_memory_offload`
+ """
+ model_ids = list(module_sizes.keys())
+ best_candidate = None
+ best_size = float("inf")
+ for r in range(1, len(model_ids) + 1):
+ for candidate_model_ids in combinations(model_ids, r):
+ candidate_size = sum(
+ module_sizes[candidate_model_id] for candidate_model_id in candidate_model_ids
+ )
+ if candidate_size < min_memory_offload:
+ continue
+ else:
+ if best_candidate is None or candidate_size < best_size:
+ best_candidate = candidate_model_ids
+ best_size = candidate_size
+
+ return best_candidate
+
+ best_offload_model_ids = search_best_candidate(module_sizes, min_memory_offload)
+
+ if best_offload_model_ids is None:
+ # if no combination is found, meaning that we cannot meet the memory requirement, offload all models
+ logger.warning("no combination of models to offload to cpu is found, offloading all models")
+ hooks_to_offload = hooks
+ else:
+ hooks_to_offload = [hook for hook in hooks if hook.model_id in best_offload_model_ids]
+
+ return hooks_to_offload
+
+
+# utils for display component info in a readable format
+# TODO: move to a different file
+def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]:
+ """Summarizes a dictionary by finding common prefixes that share the same value.
+
+ For a dictionary with dot-separated keys like: {
+ 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor': [0.6],
+ 'down_blocks.1.attentions.1.transformer_blocks.1.attn2.processor': [0.6],
+ 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor': [0.3],
+ }
+
+ Returns a dictionary where keys are the shortest common prefixes and values are their shared values: {
+ 'down_blocks': [0.6], 'up_blocks': [0.3]
+ }
+ """
+ # First group by values - convert lists to tuples to make them hashable
+ value_to_keys = {}
+ for key, value in d.items():
+ value_tuple = tuple(value) if isinstance(value, list) else value
+ if value_tuple not in value_to_keys:
+ value_to_keys[value_tuple] = []
+ value_to_keys[value_tuple].append(key)
+
+ def find_common_prefix(keys: list[str]) -> str:
+ """Find the shortest common prefix among a list of dot-separated keys."""
+ if not keys:
+ return ""
+ if len(keys) == 1:
+ return keys[0]
+
+ # Split all keys into parts
+ key_parts = [k.split(".") for k in keys]
+
+ # Find how many initial parts are common
+ common_length = 0
+ for parts in zip(*key_parts):
+ if len(set(parts)) == 1: # All parts at this position are the same
+ common_length += 1
+ else:
+ break
+
+ if common_length == 0:
+ return ""
+
+ # Return the common prefix
+ return ".".join(key_parts[0][:common_length])
+
+ # Create summary by finding common prefixes for each value group
+ summary = {}
+ for value_tuple, keys in value_to_keys.items():
+ prefix = find_common_prefix(keys)
+ if prefix: # Only add if we found a common prefix
+ # Convert tuple back to list if it was originally a list
+ value = list(value_tuple) if isinstance(d[keys[0]], list) else value_tuple
+ summary[prefix] = value
+ else:
+ summary[""] = value # Use empty string if no common prefix
+
+ return summary
+
+
+class ComponentsManager:
+ """
+ A central registry and management system for model components across multiple pipelines.
+
+ [`ComponentsManager`] provides a unified way to register, track, and reuse model components (like UNet, VAE, text
+ encoders, etc.) across different modular pipelines. It includes features for duplicate detection, memory
+ management, and component organization.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+
+ Example:
+ ```python
+ from diffusers import ComponentsManager
+
+ # Create a components manager
+ cm = ComponentsManager()
+
+ # Add components
+ cm.add("unet", unet_model, collection="sdxl")
+ cm.add("vae", vae_model, collection="sdxl")
+
+ # Enable auto offloading
+ cm.enable_auto_cpu_offload()
+
+ # Retrieve components
+ unet = cm.get_one(name="unet", collection="sdxl")
+ ```
+ """
+
+ _available_info_fields = [
+ "model_id",
+ "added_time",
+ "collection",
+ "class_name",
+ "size_gb",
+ "adapters",
+ "has_hook",
+ "execution_device",
+ "ip_adapter",
+ "quantization",
+ ]
+
+ def __init__(self):
+ self.components = OrderedDict()
+ # YiYi TODO: can remove once confirm we don't need this in mellon
+ self.added_time = OrderedDict() # Store when components were added
+ self.collections = OrderedDict() # collection_name -> set of component_names
+ self.model_hooks = None
+ self._auto_offload_enabled = False
+
+ def _lookup_ids(
+ self,
+ name: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
+ components: OrderedDict | None = None,
+ ):
+ """
+ Lookup component_ids by name, collection, or load_id. Does not support pattern matching. Returns a set of
+ component_ids
+ """
+ if components is None:
+ components = self.components
+
+ if name:
+ ids_by_name = set()
+ for component_id, component in components.items():
+ comp_name = self._id_to_name(component_id)
+ if comp_name == name:
+ ids_by_name.add(component_id)
+ else:
+ ids_by_name = set(components.keys())
+ if collection and collection not in self.collections:
+ return set()
+ elif collection and collection in self.collections:
+ ids_by_collection = set()
+ for component_id, component in components.items():
+ if component_id in self.collections[collection]:
+ ids_by_collection.add(component_id)
+ else:
+ ids_by_collection = set(components.keys())
+ if load_id:
+ ids_by_load_id = set()
+ for name, component in components.items():
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id == load_id:
+ ids_by_load_id.add(name)
+ else:
+ ids_by_load_id = set(components.keys())
+
+ ids = ids_by_name.intersection(ids_by_collection).intersection(ids_by_load_id)
+ return ids
+
+ @staticmethod
+ def _id_to_name(component_id: str):
+ return "_".join(component_id.split("_")[:-1])
+
+ def add(self, name: str, component: Any, collection: str | None = None):
+ """
+ Add a component to the ComponentsManager.
+
+ Args:
+ name (str): The name of the component
+ component (Any): The component to add
+ collection (str | None): The collection to add the component to
+
+ Returns:
+ str: The unique component ID, which is generated as "{name}_{id(component)}" where
+ id(component) is Python's built-in unique identifier for the object
+ """
+ component_id = f"{name}_{id(component)}"
+ is_new_component = True
+
+ # check for duplicated components
+ for comp_id, comp in self.components.items():
+ if comp == component:
+ comp_name = self._id_to_name(comp_id)
+ if comp_name == name:
+ logger.warning(f"ComponentsManager: component '{name}' already exists as '{comp_id}'")
+ component_id = comp_id
+ is_new_component = False
+ break
+ else:
+ logger.warning(
+ f"ComponentsManager: adding component '{name}' as '{component_id}', but it is duplicate of '{comp_id}'"
+ f"To remove a duplicate, call `components_manager.remove('')`."
+ )
+
+ # check for duplicated load_id and warn (we do not delete for you)
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id != "null":
+ components_with_same_load_id = self._lookup_ids(load_id=component._diffusers_load_id)
+ components_with_same_load_id = [id for id in components_with_same_load_id if id != component_id]
+
+ if components_with_same_load_id:
+ existing = ", ".join(components_with_same_load_id)
+ logger.warning(
+ f"ComponentsManager: adding component '{component_id}', but it has duplicate load_id '{component._diffusers_load_id}' with existing components: {existing}. "
+ f"To remove a duplicate, call `components_manager.remove('')`."
+ )
+
+ # add component to components manager
+ self.components[component_id] = component
+ if is_new_component:
+ self.added_time[component_id] = time.time()
+
+ if collection:
+ if collection not in self.collections:
+ self.collections[collection] = set()
+ if component_id not in self.collections[collection]:
+ comp_ids_in_collection = self._lookup_ids(name=name, collection=collection)
+ for comp_id in comp_ids_in_collection:
+ logger.warning(
+ f"ComponentsManager: removing existing {name} from collection '{collection}': {comp_id}"
+ )
+ # remove existing component from this collection (if it is not in any other collection, will be removed from ComponentsManager)
+ self.remove_from_collection(comp_id, collection)
+
+ self.collections[collection].add(component_id)
+ logger.info(
+ f"ComponentsManager: added component '{name}' in collection '{collection}': {component_id}"
+ )
+ else:
+ logger.info(f"ComponentsManager: added component '{name}' as '{component_id}'")
+
+ if self._auto_offload_enabled and is_new_component:
+ self.enable_auto_cpu_offload(self._auto_offload_device)
+
+ return component_id
+
+ def remove_from_collection(self, component_id: str, collection: str):
+ """
+ Remove a component from a collection.
+ """
+ if collection not in self.collections:
+ logger.warning(f"Collection '{collection}' not found in ComponentsManager")
+ return
+ if component_id not in self.collections[collection]:
+ logger.warning(f"Component '{component_id}' not found in collection '{collection}'")
+ return
+ # remove from the collection
+ self.collections[collection].remove(component_id)
+ # check if this component is in any other collection
+ comp_colls = [coll for coll, comps in self.collections.items() if component_id in comps]
+ if not comp_colls: # only if no other collection contains this component, remove it
+ logger.warning(f"ComponentsManager: removing component '{component_id}' from ComponentsManager")
+ self.remove(component_id)
+
+ def remove(self, component_id: str = None):
+ """
+ Remove a component from the ComponentsManager.
+
+ Args:
+ component_id (str): The ID of the component to remove
+ """
+ if component_id not in self.components:
+ logger.warning(f"Component '{component_id}' not found in ComponentsManager")
+ return
+
+ component = self.components.pop(component_id)
+ self.added_time.pop(component_id)
+
+ for collection in self.collections:
+ if component_id in self.collections[collection]:
+ self.collections[collection].remove(component_id)
+
+ if self._auto_offload_enabled:
+ self.enable_auto_cpu_offload(self._auto_offload_device)
+ else:
+ if isinstance(component, torch.nn.Module):
+ component.to("cpu")
+ del component
+ import gc
+
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ if torch.xpu.is_available():
+ torch.xpu.empty_cache()
+
+ # YiYi TODO: rename to search_components for now, may remove this method
+ def search_components(
+ self,
+ names: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
+ return_dict_with_names: bool = True,
+ ):
+ """
+ Search components by name with simple pattern matching. Optionally filter by collection or load_id.
+
+ Args:
+ names: Component name(s) or pattern(s)
+ Patterns:
+ - "unet" : match any component with base name "unet" (e.g., unet_123abc)
+ - "!unet" : everything except components with base name "unet"
+ - "unet*" : anything with base name starting with "unet"
+ - "!unet*" : anything with base name NOT starting with "unet"
+ - "*unet*" : anything with base name containing "unet"
+ - "!*unet*" : anything with base name NOT containing "unet"
+ - "refiner|vae|unet" : anything with base name exactly matching "refiner", "vae", or "unet"
+ - "!refiner|vae|unet" : anything with base name NOT exactly matching "refiner", "vae", or "unet"
+ - "unet*|vae*" : anything with base name starting with "unet" OR starting with "vae"
+ collection: Optional collection to filter by
+ load_id: Optional load_id to filter by
+ return_dict_with_names:
+ If True, returns a dictionary with component names as keys, throw an error if
+ multiple components with the same name are found If False, returns a dictionary
+ with component IDs as keys
+
+ Returns:
+ Dictionary mapping component names to components if return_dict_with_names=True, or a dictionary mapping
+ component IDs to components if return_dict_with_names=False
+ """
+
+ # select components based on collection and load_id filters
+ selected_ids = self._lookup_ids(collection=collection, load_id=load_id)
+ components = {k: self.components[k] for k in selected_ids}
+
+ def get_return_dict(components, return_dict_with_names):
+ """
+ Create a dictionary mapping component names to components if return_dict_with_names=True, or a dictionary
+ mapping component IDs to components if return_dict_with_names=False, throw an error if duplicate component
+ names are found when return_dict_with_names=True
+ """
+ if return_dict_with_names:
+ dict_to_return = {}
+ for comp_id, comp in components.items():
+ comp_name = self._id_to_name(comp_id)
+ if comp_name in dict_to_return:
+ raise ValueError(
+ f"Duplicate component names found in the search results: {comp_name}, please set `return_dict_with_names=False` to return a dictionary with component IDs as keys"
+ )
+ dict_to_return[comp_name] = comp
+ return dict_to_return
+ else:
+ return components
+
+ # if no names are provided, return the filtered components as it is
+ if names is None:
+ return get_return_dict(components, return_dict_with_names)
+
+ # if names is not a string, raise an error
+ elif not isinstance(names, str):
+ raise ValueError(f"Invalid type for `names: {type(names)}, only support string")
+
+ # Create mapping from component_id to base_name for components to be used for pattern matching
+ base_names = {comp_id: self._id_to_name(comp_id) for comp_id in components.keys()}
+
+ # Helper function to check if a component matches a pattern based on its base name
+ def matches_pattern(component_id, pattern, exact_match=False):
+ """
+ Helper function to check if a component matches a pattern based on its base name.
+
+ Args:
+ component_id: The component ID to check
+ pattern: The pattern to match against
+ exact_match: If True, only exact matches to base_name are considered
+ """
+ base_name = base_names[component_id]
+
+ # Exact match with base name
+ if exact_match:
+ return pattern == base_name
+
+ # Prefix match (ends with *)
+ elif pattern.endswith("*"):
+ prefix = pattern[:-1]
+ return base_name.startswith(prefix)
+
+ # Contains match (starts with *)
+ elif pattern.startswith("*"):
+ search = pattern[1:-1] if pattern.endswith("*") else pattern[1:]
+ return search in base_name
+
+ # Exact match (no wildcards)
+ else:
+ return pattern == base_name
+
+ # Check if this is a "not" pattern
+ is_not_pattern = names.startswith("!")
+ if is_not_pattern:
+ names = names[1:] # Remove the ! prefix
+
+ # Handle OR patterns (containing |)
+ if "|" in names:
+ terms = names.split("|")
+ matches = {}
+
+ for comp_id, comp in components.items():
+ # For OR patterns with exact names (no wildcards), we do exact matching on base names
+ exact_match = all(not (term.startswith("*") or term.endswith("*")) for term in terms)
+
+ # Check if any of the terms match this component
+ should_include = any(matches_pattern(comp_id, term, exact_match) for term in terms)
+
+ # Flip the decision if this is a NOT pattern
+ if is_not_pattern:
+ should_include = not should_include
+
+ if should_include:
+ matches[comp_id] = comp
+
+ log_msg = "NOT " if is_not_pattern else ""
+ match_type = "exactly matching" if exact_match else "matching any of patterns"
+ logger.info(f"Getting components {log_msg}{match_type} {terms}: {list(matches.keys())}")
+
+ # Try exact match with a base name
+ elif any(names == base_name for base_name in base_names.values()):
+ # Find all components with this base name
+ matches = {
+ comp_id: comp
+ for comp_id, comp in components.items()
+ if (base_names[comp_id] == names) != is_not_pattern
+ }
+
+ if is_not_pattern:
+ logger.info(f"Getting all components except those with base name '{names}': {list(matches.keys())}")
+ else:
+ logger.info(f"Getting components with base name '{names}': {list(matches.keys())}")
+
+ # Prefix match (ends with *)
+ elif names.endswith("*"):
+ prefix = names[:-1]
+ matches = {
+ comp_id: comp
+ for comp_id, comp in components.items()
+ if base_names[comp_id].startswith(prefix) != is_not_pattern
+ }
+ if is_not_pattern:
+ logger.info(f"Getting components NOT starting with '{prefix}': {list(matches.keys())}")
+ else:
+ logger.info(f"Getting components starting with '{prefix}': {list(matches.keys())}")
+
+ # Contains match (starts with *)
+ elif names.startswith("*"):
+ search = names[1:-1] if names.endswith("*") else names[1:]
+ matches = {
+ comp_id: comp
+ for comp_id, comp in components.items()
+ if (search in base_names[comp_id]) != is_not_pattern
+ }
+ if is_not_pattern:
+ logger.info(f"Getting components NOT containing '{search}': {list(matches.keys())}")
+ else:
+ logger.info(f"Getting components containing '{search}': {list(matches.keys())}")
+
+ # Substring match (no wildcards, but not an exact component name)
+ elif any(names in base_name for base_name in base_names.values()):
+ matches = {
+ comp_id: comp
+ for comp_id, comp in components.items()
+ if (names in base_names[comp_id]) != is_not_pattern
+ }
+ if is_not_pattern:
+ logger.info(f"Getting components NOT containing '{names}': {list(matches.keys())}")
+ else:
+ logger.info(f"Getting components containing '{names}': {list(matches.keys())}")
+
+ else:
+ raise ValueError(f"Component or pattern '{names}' not found in ComponentsManager")
+
+ if not matches:
+ raise ValueError(f"No components found matching pattern '{names}'")
+
+ return get_return_dict(matches, return_dict_with_names)
+
+ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"):
+ """
+ Enable automatic CPU offloading for all components.
+
+ The algorithm works as follows:
+ 1. All models start on CPU by default
+ 2. When a model's forward pass is called, it's moved to the execution device
+ 3. If there's insufficient memory, other models on the device are moved back to CPU
+ 4. The system tries to offload the smallest combination of models that frees enough memory
+ 5. Models stay on the execution device until another model needs memory and forces them off
+
+ Args:
+ device (str | int | torch.device): The execution device where models are moved for forward passes
+ memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of
+ memory to keep free on the device to avoid running out of memory during model
+ execution (e.g., for intermediate activations, gradients, etc.)
+ """
+ if not is_accelerate_available():
+ raise ImportError("Make sure to install accelerate to use auto_cpu_offload")
+
+ if device is None:
+ device = get_device()
+ if not isinstance(device, torch.device):
+ device = torch.device(device)
+
+ device_type = device.type
+ device_module = getattr(torch, device_type, torch.cuda)
+ if not hasattr(device_module, "mem_get_info"):
+ raise NotImplementedError(
+ f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}."
+ )
+
+ if device.index is None:
+ device = torch.device(f"{device.type}:{0}")
+
+ for name, component in self.components.items():
+ if isinstance(component, torch.nn.Module) and hasattr(component, "_hf_hook"):
+ remove_hook_from_module(component, recurse=True)
+
+ self.disable_auto_cpu_offload()
+ offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin)
+
+ all_hooks = []
+ for name, component in self.components.items():
+ if isinstance(component, torch.nn.Module):
+ hook = custom_offload_with_hook(name, component, device, offload_strategy=offload_strategy)
+ all_hooks.append(hook)
+
+ for hook in all_hooks:
+ other_hooks = [h for h in all_hooks if h is not hook]
+ for other_hook in other_hooks:
+ if other_hook.hook.execution_device == hook.hook.execution_device:
+ hook.add_other_hook(other_hook)
+
+ self.model_hooks = all_hooks
+ self._auto_offload_enabled = True
+ self._auto_offload_device = device
+
+ def disable_auto_cpu_offload(self):
+ """
+ Disable automatic CPU offloading for all components.
+ """
+ if self.model_hooks is None:
+ self._auto_offload_enabled = False
+ return
+
+ for hook in self.model_hooks:
+ hook.offload()
+ hook.remove()
+ if self.model_hooks:
+ clear_device_cache()
+ self.model_hooks = None
+ self._auto_offload_enabled = False
+
+ def get_model_info(
+ self,
+ component_id: str,
+ fields: str | list[str] | None = None,
+ ) -> dict[str, Any] | None:
+ """Get comprehensive information about a component.
+
+ Args:
+ component_id (str): Name of the component to get info for
+ fields (str | list[str] | None):
+ Field(s) to return. Can be a string for single field or list of fields. If None, uses the
+ available_info_fields setting.
+
+ Returns:
+ Dictionary containing requested component metadata. If fields is specified, returns only those fields.
+ Otherwise, returns all fields.
+ """
+ if component_id not in self.components:
+ raise ValueError(f"Component '{component_id}' not found in ComponentsManager")
+
+ component = self.components[component_id]
+
+ # Validate fields if specified
+ if fields is not None:
+ if isinstance(fields, str):
+ fields = [fields]
+ for field in fields:
+ if field not in self._available_info_fields:
+ raise ValueError(f"Field '{field}' not found in available_info_fields")
+
+ # Build complete info dict first
+ info = {
+ "model_id": component_id,
+ "added_time": self.added_time[component_id],
+ "collection": ", ".join([coll for coll, comps in self.collections.items() if component_id in comps])
+ or None,
+ }
+
+ # Additional info for torch.nn.Module components
+ if isinstance(component, torch.nn.Module):
+ # Check for hook information
+ has_hook = hasattr(component, "_hf_hook")
+ execution_device = None
+ if has_hook and hasattr(component._hf_hook, "execution_device"):
+ execution_device = component._hf_hook.execution_device
+
+ info.update(
+ {
+ "class_name": component.__class__.__name__,
+ "size_gb": component.get_memory_footprint() / (1024**3),
+ "adapters": None, # Default to None
+ "has_hook": has_hook,
+ "execution_device": execution_device,
+ }
+ )
+
+ # Get adapters if applicable
+ if hasattr(component, "peft_config"):
+ info["adapters"] = list(component.peft_config.keys())
+
+ # Check for IP-Adapter scales
+ if hasattr(component, "_load_ip_adapter_weights") and hasattr(component, "attn_processors"):
+ processors = copy.deepcopy(component.attn_processors)
+ # First check if any processor is an IP-Adapter
+ processor_types = [v.__class__.__name__ for v in processors.values()]
+ if any("IPAdapter" in ptype for ptype in processor_types):
+ # Then get scales only from IP-Adapter processors
+ scales = {
+ k: v.scale
+ for k, v in processors.items()
+ if hasattr(v, "scale") and "IPAdapter" in v.__class__.__name__
+ }
+ if scales:
+ info["ip_adapter"] = summarize_dict_by_value_and_parts(scales)
+
+ # Check for quantization
+ hf_quantizer = getattr(component, "hf_quantizer", None)
+ if hf_quantizer is not None:
+ quant_config = hf_quantizer.quantization_config
+ if hasattr(quant_config, "to_diff_dict"):
+ info["quantization"] = quant_config.to_diff_dict()
+ else:
+ info["quantization"] = quant_config.to_dict()
+ else:
+ info["quantization"] = None
+
+ # If fields specified, filter info
+ if fields is not None:
+ return {k: v for k, v in info.items() if k in fields}
+ else:
+ return info
+
+ # YiYi TODO: (1) add display fields, allow user to set which fields to display in the comnponents table
+ def __repr__(self):
+ # Handle empty components case
+ if not self.components:
+ return "Components:\n" + "=" * 50 + "\nNo components registered.\n" + "=" * 50
+
+ # Extract load_id if available
+ def get_load_id(component):
+ if hasattr(component, "_diffusers_load_id"):
+ return component._diffusers_load_id
+ return "N/A"
+
+ # Format device info compactly
+ def format_device(component, info):
+ if not info["has_hook"]:
+ return str(getattr(component, "device", "N/A"))
+ else:
+ device = str(getattr(component, "device", "N/A"))
+ exec_device = str(info["execution_device"] or "N/A")
+ return f"{device}({exec_device})"
+
+ # Get max length of load_ids for models
+ load_ids = [
+ get_load_id(component)
+ for component in self.components.values()
+ if isinstance(component, torch.nn.Module) and hasattr(component, "_diffusers_load_id")
+ ]
+ max_load_id_len = max([15] + [len(str(lid)) for lid in load_ids]) if load_ids else 15
+
+ # Get all collections for each component
+ component_collections = {}
+ for name in self.components.keys():
+ component_collections[name] = []
+ for coll, comps in self.collections.items():
+ if name in comps:
+ component_collections[name].append(coll)
+ if not component_collections[name]:
+ component_collections[name] = ["N/A"]
+
+ # Find the maximum collection name length
+ all_collections = [coll for colls in component_collections.values() for coll in colls]
+ max_collection_len = max(10, max(len(str(c)) for c in all_collections)) if all_collections else 10
+
+ col_widths = {
+ "id": max(15, max(len(name) for name in self.components.keys())),
+ "class": max(25, max(len(component.__class__.__name__) for component in self.components.values())),
+ "device": 20,
+ "dtype": 15,
+ "size": 10,
+ "load_id": max_load_id_len,
+ "collection": max_collection_len,
+ }
+
+ # Create the header lines
+ sep_line = "=" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n"
+ dash_line = "-" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n"
+
+ output = "Components:\n" + sep_line
+
+ # Separate components into models and others
+ models = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)}
+ others = {k: v for k, v in self.components.items() if not isinstance(v, torch.nn.Module)}
+
+ # Models section
+ if models:
+ output += "Models:\n" + dash_line
+ # Column headers
+ output += f"{'Name_ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | "
+ output += f"{'Device: act(exec)':<{col_widths['device']}} | {'Dtype':<{col_widths['dtype']}} | "
+ output += f"{'Size (GB)':<{col_widths['size']}} | {'Load ID':<{col_widths['load_id']}} | Collection\n"
+ output += dash_line
+
+ # Model entries
+ for name, component in models.items():
+ info = self.get_model_info(name)
+ device_str = format_device(component, info)
+ dtype = str(component.dtype) if hasattr(component, "dtype") else "N/A"
+ load_id = get_load_id(component)
+
+ # Print first collection on the main line
+ first_collection = component_collections[name][0] if component_collections[name] else "N/A"
+
+ output += f"{name:<{col_widths['id']}} | {info['class_name']:<{col_widths['class']}} | "
+ output += f"{device_str:<{col_widths['device']}} | {dtype:<{col_widths['dtype']}} | "
+ output += f"{info['size_gb']:<{col_widths['size']}.2f} | {load_id:<{col_widths['load_id']}} | {first_collection}\n"
+
+ # Print additional collections on separate lines if they exist
+ for i in range(1, len(component_collections[name])):
+ collection = component_collections[name][i]
+ output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | "
+ output += f"{'':<{col_widths['device']}} | {'':<{col_widths['dtype']}} | "
+ output += f"{'':<{col_widths['size']}} | {'':<{col_widths['load_id']}} | {collection}\n"
+
+ output += dash_line
+
+ # Other components section
+ if others:
+ if models: # Add extra newline if we had models section
+ output += "\n"
+ output += "Other Components:\n" + dash_line
+ # Column headers for other components
+ output += f"{'ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | Collection\n"
+ output += dash_line
+
+ # Other component entries
+ for name, component in others.items():
+ info = self.get_model_info(name)
+
+ # Print first collection on the main line
+ first_collection = component_collections[name][0] if component_collections[name] else "N/A"
+
+ output += f"{name:<{col_widths['id']}} | {component.__class__.__name__:<{col_widths['class']}} | {first_collection}\n"
+
+ # Print additional collections on separate lines if they exist
+ for i in range(1, len(component_collections[name])):
+ collection = component_collections[name][i]
+ output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | {collection}\n"
+
+ output += dash_line
+
+ # Add additional component info
+ output += "\nAdditional Component Info:\n" + "=" * 50 + "\n"
+ for name in self.components:
+ info = self.get_model_info(name)
+ if info is not None and (
+ info.get("adapters") is not None or info.get("ip_adapter") or info.get("quantization")
+ ):
+ output += f"\n{name}:\n"
+ if info.get("adapters") is not None:
+ output += f" Adapters: {info['adapters']}\n"
+ if info.get("ip_adapter"):
+ output += " IP-Adapter: Enabled\n"
+ if info.get("quantization"):
+ output += f" Quantization: {info['quantization']}\n"
+
+ return output
+
+ def get_one(
+ self,
+ component_id: str | None = None,
+ name: str | None = None,
+ collection: str | None = None,
+ load_id: str | None = None,
+ ) -> Any:
+ """
+ Get a single component by either:
+ - searching name (pattern matching), collection, or load_id.
+ - passing in a component_id
+ Raises an error if multiple components match or none are found.
+
+ Args:
+ component_id (str | None): Optional component ID to get
+ name (str | None): Component name or pattern
+ collection (str | None): Optional collection to filter by
+ load_id (str | None): Optional load_id to filter by
+
+ Returns:
+ A single component
+
+ Raises:
+ ValueError: If no components match or multiple components match
+ """
+
+ if component_id is not None and (name is not None or collection is not None or load_id is not None):
+ raise ValueError("If searching by component_id, do not pass name, collection, or load_id")
+
+ # search by component_id
+ if component_id is not None:
+ if component_id not in self.components:
+ raise ValueError(f"Component '{component_id}' not found in ComponentsManager")
+ return self.components[component_id]
+ # search with name/collection/load_id
+ results = self.search_components(name, collection, load_id)
+
+ if not results:
+ raise ValueError(f"No components found matching '{name}'")
+
+ if len(results) > 1:
+ raise ValueError(f"Multiple components found matching '{name}': {list(results.keys())}")
+
+ return next(iter(results.values()))
+
+ def get_ids(self, names: str | list[str] = None, collection: str | None = None):
+ """
+ Get component IDs by a list of names, optionally filtered by collection.
+
+ Args:
+ names (str | list[str]): list of component names
+ collection (str | None): Optional collection to filter by
+
+ Returns:
+ list[str]: list of component IDs
+ """
+ ids = set()
+ if not isinstance(names, list):
+ names = [names]
+ for name in names:
+ ids.update(self._lookup_ids(name=name, collection=collection))
+ return list(ids)
+
+ def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | None = True):
+ """
+ Get components by a list of IDs.
+
+ Args:
+ ids (list[str]):
+ list of component IDs
+ return_dict_with_names (bool | None):
+ Whether to return a dictionary with component names as keys:
+
+ Returns:
+ dict[str, Any]: Dictionary of components.
+ - If return_dict_with_names=True, keys are component names.
+ - If return_dict_with_names=False, keys are component IDs.
+
+ Raises:
+ ValueError: If duplicate component names are found in the search results when return_dict_with_names=True
+ """
+ components = {id: self.components[id] for id in ids}
+
+ if return_dict_with_names:
+ dict_to_return = {}
+ for comp_id, comp in components.items():
+ comp_name = self._id_to_name(comp_id)
+ if comp_name in dict_to_return:
+ raise ValueError(
+ f"Duplicate component names found in the search results: {comp_name}, please set `return_dict_with_names=False` to return a dictionary with component IDs as keys"
+ )
+ dict_to_return[comp_name] = comp
+ return dict_to_return
+ else:
+ return components
+
+ def get_components_by_names(self, names: list[str], collection: str | None = None):
+ """
+ Get components by a list of names, optionally filtered by collection.
+
+ Args:
+ names (list[str]): list of component names
+ collection (str | None): Optional collection to filter by
+
+ Returns:
+ dict[str, Any]: Dictionary of components with component names as keys
+
+ Raises:
+ ValueError: If duplicate component names are found in the search results
+ """
+ ids = self.get_ids(names, collection)
+ return self.get_components_by_ids(ids)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5645c741198621af3b7a6559a81cb171c8736d4
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py
@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"]
+ _import_structure["modular_blocks_cosmos3_distilled"] = ["Cosmos3DistilledBlocks"]
+ _import_structure["modular_pipeline"] = ["Cosmos3DistilledModularPipeline", "Cosmos3OmniModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_cosmos3 import Cosmos3OmniBlocks
+ from .modular_blocks_cosmos3_distilled import Cosmos3DistilledBlocks
+ from .modular_pipeline import Cosmos3DistilledModularPipeline, Cosmos3OmniModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py
new file mode 100644
index 0000000000000000000000000000000000000000..d72bc9a6dde2daea671ecaf1ceb00b6e392a7550
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py
@@ -0,0 +1,113 @@
+import torch
+
+from ...utils import encode_video, export_to_video
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+class Cosmos3ActionOutputStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Post-processes action latents into action outputs."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="action_latents",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Denoised action latents.",
+ ),
+ InputParam(
+ name="action_mode", type_hint=str, default=None, description="Requested action-generation mode."
+ ),
+ InputParam(
+ name="raw_action_dim_resolved",
+ type_hint=int,
+ default=None,
+ description="Unpadded action-vector dimension.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ action_output = None
+ if block_state.action_mode in {"inverse_dynamics", "policy"} and block_state.action_latents is not None:
+ action_output = block_state.action_latents
+ if block_state.raw_action_dim_resolved is not None:
+ action_output = action_output[:, : block_state.raw_action_dim_resolved]
+ action_output = [action_output.detach().cpu()]
+ block_state.action = action_output
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ExportStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Optional export block that writes decoded outputs to disk. Writes `videos` to `output_path` via "
+ "`export_to_video`, or muxes `videos` with `sound` via `encode_video` when a waveform is present. "
+ "Not wired into the default blocks; add it explicitly when you want the pipeline to produce a file."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="videos", required=True, description="Generated video frames to export."),
+ InputParam(
+ name="output_path",
+ type_hint=str,
+ required=True,
+ description="Destination path for the exported video.",
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the exported video."),
+ InputParam(
+ name="sound",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Generated waveform to mux into the video.",
+ ),
+ InputParam(
+ name="sampling_rate",
+ type_hint=int,
+ default=None,
+ description="Sample rate of the generated waveform in Hz.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("output_path", type_hint=str, description="Path of the exported video file.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ output_path = str(block_state.output_path)
+ fps = int(round(block_state.fps))
+ if block_state.sound is not None:
+ if block_state.sampling_rate is None:
+ raise ValueError("`sampling_rate` is required to export a video with sound.")
+ encode_video(
+ block_state.videos,
+ fps=fps,
+ audio=block_state.sound,
+ audio_sample_rate=int(block_state.sampling_rate),
+ output_path=output_path,
+ )
+ else:
+ export_to_video(block_state.videos, output_path, fps=fps, macro_block_size=1)
+ block_state.output_path = output_path
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebcfb10980d45704ce9f5ce131d9840615d01240
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py
@@ -0,0 +1,1331 @@
+import copy
+
+import numpy as np
+import torch
+
+from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Builds cond/uncond text segments before denoising."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="cond_input_ids", required=True, description="Token IDs for the conditional prompt."),
+ InputParam(name="uncond_input_ids", required=True, description="Token IDs for the unconditional prompt."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_text_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional text segment for the denoiser.",
+ ),
+ OutputParam(
+ "uncond_text_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional text segment for the denoiser.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ block_state.cond_text_segment = components._prepare_text_segment(block_state.cond_input_ids, device=device)
+ block_state.uncond_text_segment = components._prepare_text_segment(block_state.uncond_input_ids, device=device)
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3VisionPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares noisy vision latents and the vision conditioning mask."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="x0_tokens_vision",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Vision latents encoded from the conditioning image or video.",
+ ),
+ InputParam(
+ name="vision_condition_frames",
+ type_hint=list[int],
+ default=None,
+ description="Latent-frame indexes fixed by visual conditioning.",
+ ),
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
+ InputParam(
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
+ ),
+ InputParam(
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Pre-generated noisy vision latents.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy vision latents for denoising."),
+ OutputParam("fps_vision", type_hint=float, description="Frame rate used to pack vision latents."),
+ OutputParam(
+ "vision_condition_mask",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Mask marking conditioned vision latent frames.",
+ ),
+ OutputParam(
+ "vision_condition_indexes_for_pack",
+ type_hint=list[int],
+ description="Indexes of conditioned vision latent frames.",
+ ),
+ OutputParam(
+ "vision_conditioning_latents",
+ type_hint=torch.Tensor,
+ description="Clean encoded vision latents used to re-anchor image conditioning each step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ x0_tokens_vision = block_state.x0_tokens_vision
+ if x0_tokens_vision is None:
+ if block_state.num_frames < 1:
+ raise ValueError(f"num_frames must be >= 1, got {block_state.num_frames}.")
+ sf_spatial = components.vae_scale_factor_spatial
+ if block_state.height % sf_spatial != 0 or block_state.width % sf_spatial != 0:
+ raise ValueError(
+ f"height and width must be multiples of {sf_spatial}, got ({block_state.height}, {block_state.width})."
+ )
+ latent_shape = (
+ 1,
+ components.num_channels_latents,
+ (block_state.num_frames - 1) // components.vae_scale_factor_temporal + 1,
+ block_state.height // sf_spatial,
+ block_state.width // sf_spatial,
+ )
+ x0_tokens_vision = torch.zeros(latent_shape, device=device, dtype=torch.float32)
+ else:
+ x0_tokens_vision = x0_tokens_vision.to(device=device, dtype=torch.float32)
+
+ block_state.fps_vision = float(block_state.fps)
+ condition_frames = block_state.vision_condition_frames or []
+ block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype)
+ for frame_idx in condition_frames:
+ if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]:
+ block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0
+
+ if block_state.latents is None:
+ pure_noise = randn_tensor(
+ tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype
+ )
+ block_state.latents = (
+ block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype)
+ + (1.0 - block_state.vision_condition_mask) * pure_noise
+ )
+ else:
+ block_state.latents = block_state.latents.to(device=device, dtype=dtype)
+
+ vision_condition_indexes = torch.nonzero(
+ block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False
+ ).flatten()
+ block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes]
+ block_state.vision_conditioning_latents = x0_tokens_vision
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3SoundPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares noisy sound latents and the sound conditioning mask."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ InputParam(
+ name="sound_latents",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Pre-generated noisy sound latents.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Noisy sound latents for denoising."),
+ OutputParam("fps_sound", type_hint=float, description="Frame rate of the sound latent sequence."),
+ OutputParam(
+ "sound_condition_mask",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Mask marking conditioned sound latent frames.",
+ ),
+ OutputParam("sound_scheduler", description="Scheduler used to update sound latents."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ if not components.transformer.config.sound_gen:
+ raise ValueError("Sound generation requires a transformer trained with sound_gen=True.")
+
+ sound_dim = components.transformer.config.sound_dim
+ block_state.fps_sound = float(components.transformer.config.sound_latent_fps)
+ n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate)
+ hop_size = components.sound_hop_size
+ t_sound = (n_audio_samples + hop_size - 1) // hop_size
+ x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype)
+ block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype)
+
+ if block_state.sound_latents is None:
+ pure_noise = randn_tensor(
+ tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype
+ )
+ block_state.sound_latents = (
+ block_state.sound_condition_mask.T * x0_tokens_sound
+ + (1.0 - block_state.sound_condition_mask.T) * pure_noise
+ )
+ else:
+ block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype)
+
+ block_state.sound_scheduler = copy.deepcopy(components.scheduler)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ActionPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares noisy action latents and the action conditioning mask."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="action",
+ type_hint=CosmosActionCondition,
+ required=True,
+ description="Action-conditioning metadata.",
+ ),
+ InputParam(
+ name="action_condition_frame_indexes",
+ type_hint=list[int],
+ default=None,
+ description="Action-frame indexes fixed by action conditioning.",
+ ),
+ InputParam(
+ name="action_latents",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Pre-generated noisy action latents.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Noisy action latents for denoising."),
+ OutputParam(
+ "action_condition_mask",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Mask marking conditioned action latent frames.",
+ ),
+ OutputParam(
+ "action_domain_ids",
+ type_hint=list[torch.Tensor],
+ kwargs_type="denoiser_input_fields",
+ description="Embodiment domain IDs for action conditioning.",
+ ),
+ OutputParam(
+ "raw_action_dim_resolved",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Unpadded action-vector dimension.",
+ ),
+ OutputParam("action_scheduler", description="Scheduler used to update action latents."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+ action = block_state.action
+
+ if not components.transformer.config.action_gen:
+ raise ValueError("action requires a transformer trained with action_gen=True.")
+
+ block_state.raw_action_dim_resolved = int(action.raw_action_dim) if action.raw_action_dim is not None else None
+ if (
+ block_state.raw_action_dim_resolved is not None
+ and block_state.raw_action_dim_resolved > components.transformer.config.action_dim
+ ):
+ raise ValueError(
+ f"raw_action_dim={block_state.raw_action_dim_resolved} exceeds the model action_dim="
+ f"{components.transformer.config.action_dim}."
+ )
+
+ action_chunk_size = action.chunk_size
+ action_dim = components.transformer.action_dim
+ if action.mode == "forward_dynamics":
+ raw_actions = action.raw_actions
+ if raw_actions is None:
+ raise ValueError("action_mode='forward_dynamics' requires an action tensor.")
+ raw_actions = raw_actions.to(device=device, dtype=dtype)
+ if raw_actions.shape[-1] > action_dim:
+ raise ValueError(
+ f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}."
+ )
+ if raw_actions.shape[0] < action_chunk_size:
+ raw_actions = torch.cat(
+ [raw_actions, raw_actions[-1:].expand(action_chunk_size - raw_actions.shape[0], -1)],
+ dim=0,
+ )
+ raw_actions = raw_actions[:action_chunk_size]
+ if raw_actions.shape[-1] < action_dim:
+ action_padding = torch.zeros(
+ raw_actions.shape[0],
+ action_dim - raw_actions.shape[-1],
+ dtype=raw_actions.dtype,
+ device=raw_actions.device,
+ )
+ raw_actions = torch.cat([raw_actions, action_padding], dim=-1)
+ x0_tokens_action = raw_actions
+ else:
+ x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype)
+
+ if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID:
+ raise ValueError(
+ f"Unknown Cosmos3 action domain_name={action.domain_name!r}; expected one of {sorted(_EMBODIMENT_TO_DOMAIN_ID)}."
+ )
+ block_state.action_domain_ids = [
+ torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device)
+ ]
+ condition_frames = block_state.action_condition_frame_indexes or []
+ block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype)
+ for frame_idx in condition_frames:
+ if 0 <= frame_idx < block_state.action_condition_mask.shape[0]:
+ block_state.action_condition_mask[frame_idx, 0] = 1.0
+
+ if block_state.action_latents is None:
+ pure_noise = randn_tensor(
+ tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype
+ )
+ block_state.action_latents = (
+ block_state.action_condition_mask * x0_tokens_action
+ + (1.0 - block_state.action_condition_mask) * pure_noise
+ )
+ if block_state.raw_action_dim_resolved is not None:
+ block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0
+ else:
+ block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype)
+
+ block_state.action_scheduler = copy.deepcopy(components.scheduler)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3VisionPackSequenceStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Builds separate cond/uncond vision sequence segments."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+ ),
+ InputParam(
+ name="uncond_text_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional text segment.",
+ ),
+ InputParam(
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy vision latents to pack."
+ ),
+ InputParam(
+ name="fps_vision",
+ type_hint=float,
+ required=True,
+ description="Frame rate used to pack vision latents.",
+ ),
+ InputParam(
+ name="vision_condition_indexes_for_pack",
+ type_hint=list[int],
+ required=True,
+ description="Indexes of conditioned vision latent frames.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_vision_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional vision segment for the denoiser.",
+ ),
+ OutputParam(
+ "uncond_vision_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional vision segment for the denoiser.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ has_image_condition = bool(block_state.vision_condition_indexes_for_pack)
+
+ block_state.cond_vision_segment = components._prepare_vision_segment(
+ input_vision_tokens=block_state.latents,
+ has_image_condition=has_image_condition,
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
+ vision_fps=block_state.fps_vision,
+ curr=block_state.cond_text_segment["und_len"],
+ device=device,
+ condition_frame_indexes=block_state.vision_condition_indexes_for_pack,
+ )
+ block_state.uncond_vision_segment = components._prepare_vision_segment(
+ input_vision_tokens=block_state.latents,
+ has_image_condition=has_image_condition,
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
+ vision_fps=block_state.fps_vision,
+ curr=block_state.uncond_text_segment["und_len"],
+ device=device,
+ condition_frame_indexes=block_state.vision_condition_indexes_for_pack,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3SoundPackSequenceStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Builds separate cond/uncond sound sequence segments."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+ ),
+ InputParam(
+ name="uncond_text_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional text segment.",
+ ),
+ InputParam(
+ name="cond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Conditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="uncond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Unconditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="sound_latents", type_hint=torch.Tensor, required=True, description="Noisy sound latents to pack."
+ ),
+ InputParam(
+ name="fps_sound",
+ type_hint=float,
+ required=True,
+ description="Frame rate of the sound latent sequence.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_sound_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional sound segment for the denoiser.",
+ ),
+ OutputParam(
+ "uncond_sound_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional sound segment for the denoiser.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ block_state.cond_sound_segment = components._prepare_sound_segment(
+ input_sound_tokens=block_state.sound_latents,
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
+ sound_fps=block_state.fps_sound,
+ curr=block_state.cond_sequence_length,
+ device=device,
+ )
+ block_state.uncond_sound_segment = components._prepare_sound_segment(
+ input_sound_tokens=block_state.sound_latents,
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
+ sound_fps=block_state.fps_sound,
+ curr=block_state.uncond_sequence_length,
+ device=device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ActionPackSequenceStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Builds separate cond/uncond action sequence segments."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+ ),
+ InputParam(
+ name="uncond_text_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional text segment.",
+ ),
+ InputParam(
+ name="cond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Conditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="uncond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Unconditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="action_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy action latents to pack.",
+ ),
+ InputParam(
+ name="action_condition_frame_indexes",
+ type_hint=list[int],
+ default=None,
+ description="Action-frame indexes fixed by action conditioning.",
+ ),
+ InputParam(
+ name="fps_vision",
+ type_hint=float,
+ required=True,
+ description="Frame rate used to pack vision latents.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_action_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional action segment for the denoiser.",
+ ),
+ OutputParam(
+ "uncond_action_segment",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional action segment for the denoiser.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ block_state.cond_action_segment = components._prepare_action_segment(
+ input_action_tokens=block_state.action_latents,
+ condition_frame_indexes=block_state.action_condition_frame_indexes,
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
+ action_fps=block_state.fps_vision,
+ curr=block_state.cond_sequence_length,
+ device=device,
+ )
+ block_state.uncond_action_segment = components._prepare_action_segment(
+ input_action_tokens=block_state.action_latents,
+ condition_frame_indexes=block_state.action_condition_frame_indexes,
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
+ action_fps=block_state.fps_vision,
+ curr=block_state.uncond_sequence_length,
+ device=device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3VisionDenoiseInputStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Assembles text and vision sequence metadata for the denoising loop."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+ ),
+ InputParam(
+ name="uncond_text_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional text segment.",
+ ),
+ InputParam(
+ name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment."
+ ),
+ InputParam(
+ name="uncond_vision_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional vision segment.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "uncond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "cond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal sequence length.",
+ ),
+ OutputParam(
+ "uncond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal sequence length.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.cond_position_ids = torch.cat(
+ [
+ block_state.cond_text_segment["text_mrope_ids"],
+ block_state.cond_vision_segment["vision_mrope_ids"],
+ ],
+ dim=1,
+ )
+ block_state.uncond_position_ids = torch.cat(
+ [
+ block_state.uncond_text_segment["text_mrope_ids"],
+ block_state.uncond_vision_segment["vision_mrope_ids"],
+ ],
+ dim=1,
+ )
+ block_state.cond_sequence_length = (
+ block_state.cond_text_segment["und_len"] + block_state.cond_vision_segment["num_vision_tokens"]
+ )
+ block_state.uncond_sequence_length = (
+ block_state.uncond_text_segment["und_len"] + block_state.uncond_vision_segment["num_vision_tokens"]
+ )
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3SoundDenoiseInputStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Appends sound sequence metadata to the denoising-loop inputs."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_position_ids",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Conditional multimodal RoPE position IDs.",
+ ),
+ InputParam(
+ name="uncond_position_ids",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Unconditional multimodal RoPE position IDs.",
+ ),
+ InputParam(
+ name="cond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Conditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="uncond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Unconditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment."
+ ),
+ InputParam(
+ name="uncond_sound_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional sound segment.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "uncond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "cond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal sequence length.",
+ ),
+ OutputParam(
+ "uncond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal sequence length.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.cond_position_ids = torch.cat(
+ [block_state.cond_position_ids, block_state.cond_sound_segment["sound_mrope_ids"]], dim=1
+ )
+ block_state.uncond_position_ids = torch.cat(
+ [block_state.uncond_position_ids, block_state.uncond_sound_segment["sound_mrope_ids"]], dim=1
+ )
+ block_state.cond_sequence_length += block_state.cond_sound_segment["sound_len"]
+ block_state.uncond_sequence_length += block_state.uncond_sound_segment["sound_len"]
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ActionDenoiseInputStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Appends action sequence metadata to the denoising-loop inputs."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_position_ids",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Conditional multimodal RoPE position IDs.",
+ ),
+ InputParam(
+ name="uncond_position_ids",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Unconditional multimodal RoPE position IDs.",
+ ),
+ InputParam(
+ name="cond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Conditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="uncond_sequence_length",
+ type_hint=int,
+ required=True,
+ description="Unconditional multimodal sequence length.",
+ ),
+ InputParam(
+ name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment."
+ ),
+ InputParam(
+ name="uncond_action_segment",
+ type_hint=dict,
+ required=True,
+ description="Unconditional action segment.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "uncond_position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal RoPE position IDs.",
+ ),
+ OutputParam(
+ "cond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional multimodal sequence length.",
+ ),
+ OutputParam(
+ "uncond_sequence_length",
+ type_hint=int,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional multimodal sequence length.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.cond_position_ids = torch.cat(
+ [block_state.cond_position_ids, block_state.cond_action_segment["action_mrope_ids"]], dim=1
+ )
+ block_state.uncond_position_ids = torch.cat(
+ [block_state.uncond_position_ids, block_state.uncond_action_segment["action_mrope_ids"]], dim=1
+ )
+ block_state.cond_sequence_length += block_state.cond_action_segment["action_len"]
+ block_state.uncond_sequence_length += block_state.uncond_action_segment["action_len"]
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Initializes scheduler timesteps."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [ConfigSpec(name="use_native_flow_schedule", default=False)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."),
+ OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ if components.config.use_native_flow_schedule:
+ sigmas = np.linspace(
+ 1.0 - 1.0 / components.scheduler.config.num_train_timesteps,
+ 0.0,
+ block_state.num_inference_steps + 1,
+ )[:-1]
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device, sigmas=sigmas)
+ else:
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.num_warmup_steps = (
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
+ )
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Per-chunk transfer latent prep: takes the clean target latents encoded by "
+ "Cosmos3TransferChunkVaeEncoderStep and builds the noisy target latents, velocity mask, condition latents "
+ "and conditioned-frame indexes for this chunk."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="x0_tokens_vision",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Clean target vision latents encoded from the seeded target frames.",
+ ),
+ InputParam(
+ name="current_conditional_frames",
+ type_hint=int,
+ required=True,
+ description="Number of pixel frames used to seed this chunk's target.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy target latents for this chunk."),
+ OutputParam(
+ "velocity_mask",
+ type_hint=torch.Tensor,
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
+ ),
+ OutputParam(
+ "condition_latents",
+ type_hint=torch.Tensor,
+ description="Clean target latents on the conditioned frames (the autoregressive seed).",
+ ),
+ OutputParam(
+ "target_condition_indexes",
+ type_hint=list[int],
+ description="Latent-frame indexes fixed by the chunk's conditioning.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+ tcf = components.vae_scale_factor_temporal
+
+ target_x0 = block_state.x0_tokens_vision.to(device=device)
+ current_conditional_frames = block_state.current_conditional_frames
+
+ # Build the noisy target latents + conditioning mask from the clean target latents.
+ latent_t = target_x0.shape[2]
+ condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype)
+ latent_condition_frames = 0
+ if current_conditional_frames > 0:
+ latent_condition_frames = (current_conditional_frames - 1) // tcf + 1
+ condition_mask[:latent_condition_frames] = 1.0
+ noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype)
+ block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise
+ block_state.velocity_mask = 1.0 - condition_mask
+ block_state.condition_latents = condition_mask * target_x0
+ block_state.target_condition_indexes = list(range(latent_condition_frames))
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferPackSequenceStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pre-packs the three transfer CFG sequence variants: cond_full / uncond_full carry every control item, "
+ "the no-control branch drops them (only [text, target]) so the control axis can be amplified."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
+ ),
+ InputParam(
+ name="uncond_text_segment", type_hint=dict, required=True, description="Unconditional text segment."
+ ),
+ InputParam(
+ name="control_latents",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Clean control latents for this chunk, one per hint in canonical order.",
+ ),
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy target latents for this chunk.",
+ ),
+ InputParam(
+ name="target_condition_indexes",
+ type_hint=list[int],
+ required=True,
+ description="Latent-frame indexes fixed by the chunk's conditioning.",
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "cond_full_static",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional [control..., target] transfer sequence carrying every control item.",
+ ),
+ OutputParam(
+ "cond_no_control_static",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Conditional [target] transfer sequence with the control items dropped.",
+ ),
+ OutputParam(
+ "uncond_full_static",
+ type_hint=dict,
+ kwargs_type="denoiser_input_fields",
+ description="Unconditional [control..., target] transfer sequence for text CFG.",
+ ),
+ OutputParam(
+ "num_noisy_vision_tokens",
+ type_hint=int,
+ description="Number of noisy target vision tokens denoised each step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ num_hints = len(block_state.control_latents)
+
+ def _vision_pack(text_segment: dict, include_controls: bool) -> dict:
+ if include_controls:
+ vision_items = [*block_state.control_latents, block_state.latents]
+ condition_indexes = [None] * num_hints + [block_state.target_condition_indexes]
+ clean_flags = [True] * num_hints + [False]
+ else:
+ vision_items = [block_state.latents]
+ condition_indexes = [block_state.target_condition_indexes]
+ clean_flags = [False]
+
+ # Transfer packs [ctrl_1, ..., ctrl_N, target] into one vision segment
+ mrope_offset = text_segment["vision_start_temporal_offset"]
+ item_curr = text_segment["und_len"]
+ token_shapes = []
+ sequence_index_parts = []
+ mse_loss_index_parts = []
+ noisy_frame_indexes_per_item = []
+ mrope_id_parts = []
+ num_vision_tokens = 0
+ num_noisy_vision_tokens = 0
+ for item, item_condition, is_clean in zip(vision_items, condition_indexes, clean_flags):
+ latent_t = item.shape[2]
+ if is_clean:
+ frame_condition = list(range(latent_t))
+ else:
+ frame_condition = item_condition if item_condition is not None else []
+ item_segment = components._prepare_vision_segment(
+ input_vision_tokens=item,
+ has_image_condition=False,
+ mrope_offset=mrope_offset,
+ vision_fps=block_state.fps,
+ curr=item_curr,
+ device=device,
+ condition_frame_indexes=frame_condition,
+ )
+ token_shapes.extend(item_segment["vision_token_shapes"])
+ sequence_index_parts.append(item_segment["vision_sequence_indexes"])
+ mse_loss_index_parts.append(item_segment["vision_mse_loss_indexes"])
+ noisy_frame_indexes_per_item.extend(item_segment["vision_noisy_frame_indexes"])
+ mrope_id_parts.append(item_segment["vision_mrope_ids"])
+ num_vision_tokens += item_segment["num_vision_tokens"]
+ num_noisy_vision_tokens += item_segment["num_noisy_vision_tokens"]
+ item_curr += item_segment["num_vision_tokens"]
+
+ vision_segment = {
+ "vision_token_shapes": token_shapes,
+ "vision_sequence_indexes": torch.cat(sequence_index_parts, dim=0),
+ "vision_mse_loss_indexes": torch.cat(mse_loss_index_parts, dim=0),
+ "vision_noisy_frame_indexes": noisy_frame_indexes_per_item,
+ "vision_mrope_ids": torch.cat(mrope_id_parts, dim=1),
+ "num_vision_tokens": num_vision_tokens,
+ "num_noisy_vision_tokens": num_noisy_vision_tokens,
+ }
+ return {
+ **text_segment,
+ **vision_segment,
+ "position_ids": torch.cat([text_segment["text_mrope_ids"], vision_segment["vision_mrope_ids"]], dim=1),
+ "sequence_length": text_segment["und_len"] + vision_segment["num_vision_tokens"],
+ }
+
+ block_state.cond_full_static = _vision_pack(block_state.cond_text_segment, include_controls=True)
+ block_state.cond_no_control_static = _vision_pack(block_state.cond_text_segment, include_controls=False)
+ block_state.uncond_full_static = _vision_pack(block_state.uncond_text_segment, include_controls=True)
+ block_state.num_noisy_vision_tokens = block_state.cond_full_static["num_noisy_vision_tokens"]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resets the scheduler and computes timesteps for a single transfer chunk. UniPCMultistepScheduler keeps "
+ "per-step state on the instance, so it is reset per chunk (each autoregressive chunk is a full denoise)."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [InputParam.template("num_inference_steps", required=True)]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for this chunk."),
+ OutputParam(
+ "num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps for this chunk."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.num_warmup_steps = (
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
+ )
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3DistilledSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Initializes the fixed distilled sampling schedule from the pipeline's `distilled_sigmas` config."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="is_distilled", default=True),
+ ConfigSpec(name="distilled_sigmas", default=None),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", required=False, default=None),
+ InputParam(
+ name="guidance_scale",
+ type_hint=float,
+ default=None,
+ description=(
+ "Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the "
+ "scale is forced to 1.0. Passing a value other than 1.0 raises an error."
+ ),
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."),
+ OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="Resolved number of denoising steps (fixed by the distilled schedule).",
+ ),
+ OutputParam(
+ name="guidance_scale",
+ type_hint=float,
+ description="Resolved classifier-free guidance scale (always 1.0 for distilled checkpoints).",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ sigmas = components.config.distilled_sigmas
+ if not sigmas:
+ raise ValueError(
+ "Cosmos3DistilledSetTimestepsStep requires the pipeline config `distilled_sigmas` to be set "
+ "(populated from the distilled checkpoint's `modular_model_index.json`). Load a distilled Cosmos3 "
+ "checkpoint or use `Cosmos3OmniModularPipeline` for base checkpoints."
+ )
+ sigmas = [float(s) for s in sigmas]
+ distilled_steps = len(sigmas)
+
+ if block_state.num_inference_steps is not None and block_state.num_inference_steps != distilled_steps:
+ raise ValueError(
+ "This is a distilled checkpoint; the step count is fixed by the pipeline's "
+ f"`distilled_sigmas` config ({distilled_steps} steps). "
+ f"`num_inference_steps` must be {distilled_steps} or left unset (got {block_state.num_inference_steps})."
+ )
+ if block_state.guidance_scale is not None and block_state.guidance_scale != 1.0:
+ raise ValueError(
+ "This is a distilled checkpoint; classifier-free guidance is baked into the weights. "
+ f"`guidance_scale` must be 1.0 or left unset (got {block_state.guidance_scale})."
+ )
+
+ components.scheduler.set_timesteps(sigmas=sigmas, device=device)
+ block_state.num_inference_steps = distilled_steps
+ block_state.guidance_scale = 1.0
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.num_warmup_steps = len(block_state.timesteps) - distilled_steps * components.scheduler.order
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..03daff293578c480690a1a67a1b00b0e2118d399
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py
@@ -0,0 +1,166 @@
+import math
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+class Cosmos3TransferSetupStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Preprocesses the transfer control videos and resolves the autoregressive chunk geometry "
+ "(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="control_videos",
+ type_hint=dict,
+ required=True,
+ description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.",
+ ),
+ InputParam(
+ name="height", type_hint=int, default=None, description="Height of the generated video in pixels."
+ ),
+ InputParam(
+ name="width", type_hint=int, default=None, description="Width of the generated video in pixels."
+ ),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ default=None,
+ description="Optional cap on the number of output frames (defaults to the control video length).",
+ ),
+ InputParam(
+ name="num_video_frames_per_chunk",
+ type_hint=int,
+ default=None,
+ description="Number of pixel frames generated per autoregressive chunk.",
+ ),
+ InputParam(
+ name="num_conditional_frames",
+ type_hint=int,
+ default=1,
+ description="Number of frames each chunk reuses from the previous chunk's tail.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved output height in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved output width in pixels."),
+ OutputParam(
+ "control_frames",
+ type_hint=dict,
+ description="Preprocessed, time-padded control maps in canonical hint order.",
+ ),
+ OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."),
+ OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."),
+ OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."),
+ OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ if block_state.height is None:
+ block_state.height = 720
+ if block_state.width is None:
+ block_state.width = 1280
+
+ # Canonical hint order used both to validate and to order the preprocessed control maps.
+ hint_order = ["edge", "blur", "depth", "seg", "wsm"]
+ control_videos = block_state.control_videos
+ if not isinstance(control_videos, dict) or not control_videos:
+ raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.")
+ unknown = [k for k in control_videos if k not in hint_order]
+ if unknown:
+ raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.")
+ if any(v is None for v in control_videos.values()):
+ raise ValueError("`control_videos` entries must be loaded videos, not None.")
+
+ tcf = components.vae_scale_factor_temporal
+ sf = components.vae_scale_factor_spatial
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
+ raise ValueError(
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
+ )
+
+ # Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order.
+ # The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys).
+ hint_keys = [k for k in hint_order if k in control_videos]
+ control_frames = {
+ key: components.video_processor.preprocess_video(
+ control_videos[key], height=block_state.height, width=block_state.width
+ ).to(device=device, dtype=dtype)
+ for key in hint_keys
+ }
+
+ # Output frame count / chunking come from the (first) control video, optionally capped by num_frames.
+ total_frames = next(iter(control_frames.values())).shape[2]
+ if block_state.num_frames is not None:
+ total_frames = min(total_frames, block_state.num_frames)
+ total_frames = max(1, total_frames)
+
+ per_chunk = (
+ block_state.num_video_frames_per_chunk
+ if block_state.num_video_frames_per_chunk is not None
+ else total_frames
+ )
+ chunk_frames = 1 if total_frames == 1 else per_chunk
+ chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1
+
+ if total_frames <= chunk_frames:
+ num_chunks, stride = 1, chunk_frames
+ else:
+ stride = chunk_frames - block_state.num_conditional_frames
+ if stride <= 0:
+ raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.")
+ remaining = total_frames - chunk_frames
+ num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0))
+
+ # Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to
+ # keep reflecting). No truncation here; per-chunk slicing happens later.
+ padded = max(total_frames, chunk_frames)
+ control_frames_padded = {}
+ for key, frames in control_frames.items():
+ while frames.shape[2] < padded:
+ pad_len = min(frames.shape[2] - 1, padded - frames.shape[2])
+ if pad_len <= 0:
+ pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1)
+ frames = torch.cat([frames, pad_frame], dim=2)
+ break
+ frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
+ control_frames_padded[key] = frames
+ block_state.control_frames = control_frames_padded
+ block_state.total_frames = total_frames
+ block_state.chunk_frames = chunk_frames
+ block_state.num_chunks = num_chunks
+ block_state.stride = stride
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe927b32489184bc8eb61f4dec08cdce5e174349
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py
@@ -0,0 +1,259 @@
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
+from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class Cosmos3VideoDecodeStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Decodes denoised vision latents into video outputs."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True, description="Denoised vision latents to decode."),
+ InputParam.template("output_type", default="pil"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("videos")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ if block_state.output_type == "latent":
+ block_state.videos = block_state.latents
+ else:
+ in_dtype = block_state.latents.dtype
+ vae_dtype = components.vae.dtype
+ mean = components._vae_latents_mean.to(device=block_state.latents.device, dtype=vae_dtype)
+ inv_std = components._vae_latents_inv_std.to(device=block_state.latents.device, dtype=vae_dtype)
+ z_raw = block_state.latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1)
+ decoded = components.vae.decode(z_raw).sample.to(in_dtype)
+ block_state.videos = components.video_processor.postprocess_video(
+ decoded, output_type=block_state.output_type
+ )[0]
+
+ if components.requires_safety_checker and block_state.output_type != "latent":
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ block_state.videos = components._apply_video_safety_check(
+ block_state.videos, output_type=block_state.output_type, device=device
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3SoundDecodeStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Decodes sound latents into waveform output."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("sound_tokenizer", Cosmos3AVAEAudioTokenizer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="sound_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Denoised sound latents to decode.",
+ )
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."),
+ OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if components.sound_tokenizer is None:
+ raise ValueError("Sound decoding requires a sound-capable checkpoint with a sound_tokenizer.")
+ block_state.sound = components.decode_sound(block_state.sound_latents)
+ block_state.sampling_rate = int(components.sound_tokenizer.config.sampling_rate)
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferDecodeChunkStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Decodes one transfer chunk's latents to pixels (float32, clamped to [-1, 1]), records it as the "
+ "autoregressive seed for the next chunk, and appends it to output_chunks (dropping the overlap that "
+ "later chunks share with the previous chunk's conditioning frames)."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKLWan)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Denoised target latents for this chunk.",
+ ),
+ InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."),
+ InputParam(
+ name="current_conditional_frames",
+ type_hint=int,
+ required=True,
+ description="Number of pixel frames this chunk reused from the previous chunk.",
+ ),
+ InputParam(
+ name="output_chunks",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Decoded pixel chunks accumulated so far.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "previous_output",
+ type_hint=torch.Tensor,
+ description="Decoded pixels of this chunk, used to seed the next chunk.",
+ ),
+ OutputParam(
+ "output_chunks",
+ type_hint=list[torch.Tensor],
+ description="Decoded pixel chunks accumulated so far (with this chunk appended).",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ latents = block_state.latents
+ vae_dtype = components.vae.dtype
+ mean = components._vae_latents_mean.to(device=latents.device, dtype=vae_dtype)
+ inv_std = components._vae_latents_inv_std.to(device=latents.device, dtype=vae_dtype)
+ z_raw = latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1)
+ output_video = components.vae.decode(z_raw).sample.to(torch.float32).clamp(-1, 1)
+ block_state.previous_output = output_video
+ chunk = (
+ output_video if block_state.chunk_id == 0 else output_video[:, :, block_state.current_conditional_frames :]
+ )
+ block_state.output_chunks = [*block_state.output_chunks, chunk]
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferStitchStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Concatenates the decoded transfer chunks along time, truncates to total_frames, and post-processes to "
+ "the requested output type. Transfer produces no audio, so sound / sampling_rate are None."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="output_chunks",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Decoded pixel chunks to stitch together.",
+ ),
+ InputParam(
+ name="total_frames", type_hint=int, required=True, description="Total number of output frames to keep."
+ ),
+ InputParam.template("output_type", default="pil"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("videos", description="The generated transfer video."),
+ OutputParam("sound", description="Always None for transfer (no audio)."),
+ OutputParam("sampling_rate", description="Always None for transfer (no audio)."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ decoded = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames]
+ block_state.videos = components.video_processor.postprocess_video(
+ decoded, output_type=block_state.output_type
+ )[0]
+
+ if components.requires_safety_checker and block_state.output_type != "latent":
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ block_state.videos = components._apply_video_safety_check(
+ block_state.videos, output_type=block_state.output_type, device=device
+ )
+
+ block_state.sound = None
+ block_state.sampling_rate = None
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e437231c86b5a59ae1ffd12d6efd2a8f1185b92
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py
@@ -0,0 +1,889 @@
+import inspect
+
+import torch
+
+from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares vision tokens and timesteps for one denoising iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True, description="Noisy vision latents to denoise."),
+ InputParam(
+ name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "vision_tokens",
+ type_hint=list[torch.Tensor],
+ description="Vision tokens for the transformer denoiser.",
+ ),
+ OutputParam("vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the vision tokens."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ device = components._execution_device
+ block_state.vision_tokens = [block_state.latents.to(device=device, dtype=components.transformer.dtype)]
+ block_state.vision_timesteps = torch.full(
+ (block_state.cond_vision_segment["num_noisy_vision_tokens"],), t.item(), device=device
+ )
+ return components, block_state
+
+
+class Cosmos3SoundLoopPrepareStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares sound tokens and timesteps for one denoising iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="sound_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy sound latents to denoise.",
+ ),
+ InputParam(
+ name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "sound_tokens", type_hint=list[torch.Tensor], description="Sound tokens for the transformer denoiser."
+ ),
+ OutputParam("sound_timesteps", type_hint=torch.Tensor, description="Timesteps for the sound tokens."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ device = components._execution_device
+ block_state.sound_tokens = [block_state.sound_latents.to(device=device, dtype=components.transformer.dtype)]
+ block_state.sound_timesteps = torch.full(
+ (block_state.cond_sound_segment["sound_len"],), t.item(), device=device
+ )
+ return components, block_state
+
+
+class Cosmos3ActionLoopPrepareStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares action tokens and timesteps for one denoising iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="action_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy action latents to denoise.",
+ ),
+ InputParam(
+ name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "action_tokens",
+ type_hint=list[torch.Tensor],
+ description="Action tokens for the transformer denoiser.",
+ ),
+ OutputParam("action_timesteps", type_hint=torch.Tensor, description="Timesteps for the action tokens."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ device = components._execution_device
+ block_state.action_tokens = [block_state.action_latents.to(device=device, dtype=components.transformer.dtype)]
+ block_state.action_timesteps = torch.full(
+ (block_state.cond_action_segment["num_noisy_action_tokens"],), t.item(), device=device
+ )
+ return components, block_state
+
+
+class Cosmos3LoopDenoiser(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Predicts available Cosmos3 modality velocities for one denoising iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ name="guidance_scale",
+ type_hint=float,
+ default=6.0,
+ description="Scale for classifier-free guidance.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "velocity_vision", type_hint=torch.Tensor, description="Predicted velocity for vision latents."
+ ),
+ OutputParam("velocity_sound", type_hint=torch.Tensor, description="Predicted velocity for sound latents."),
+ OutputParam(
+ "velocity_action", type_hint=torch.Tensor, description="Predicted velocity for action latents."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ denoiser_input_fields = block_state.denoiser_input_fields
+ loop_input_fields = block_state.as_dict()
+ has_sound = "sound_tokens" in loop_input_fields
+ has_action = "action_tokens" in loop_input_fields
+ do_cfg = block_state.guidance_scale != 1.0
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters)
+
+ prediction_passes = ["cond"]
+ if do_cfg:
+ prediction_passes.append("uncond")
+
+ velocities = {}
+ for pass_name in prediction_passes:
+ transformer_kwargs = {}
+ for field_name, field_value in denoiser_input_fields.items():
+ if field_name.startswith(f"{pass_name}_"):
+ transformer_field_name = field_name.removeprefix(f"{pass_name}_")
+ if transformer_field_name.endswith("_segment"):
+ transformer_kwargs.update(field_value)
+ else:
+ transformer_kwargs[transformer_field_name] = field_value
+ elif field_name in transformer_args:
+ transformer_kwargs[field_name] = field_value
+ transformer_kwargs.update(
+ {
+ field_name: field_value
+ for field_name, field_value in loop_input_fields.items()
+ if field_name in transformer_args
+ }
+ )
+ transformer_kwargs = {
+ name: value for name, value in transformer_kwargs.items() if name in transformer_args
+ }
+ preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs, return_dict=False)
+ velocities[pass_name] = components._mask_velocity_predictions(
+ preds_vision,
+ preds_sound,
+ vision_condition_mask=[loop_input_fields["vision_condition_mask"]],
+ sound_condition_mask=[loop_input_fields["sound_condition_mask"]] if has_sound else None,
+ preds_action=preds_action,
+ action_condition_mask=[loop_input_fields["action_condition_mask"]] if has_action else None,
+ raw_action_dim=loop_input_fields.get("raw_action_dim_resolved"),
+ )
+
+ cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"]
+ if do_cfg:
+ uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"]
+ block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * (
+ cond_velocity_vision - uncond_velocity_vision
+ )
+ block_state.velocity_sound = (
+ uncond_velocity_sound + block_state.guidance_scale * (cond_velocity_sound - uncond_velocity_sound)
+ if has_sound
+ else None
+ )
+ block_state.velocity_action = (
+ uncond_velocity_action + block_state.guidance_scale * (cond_velocity_action - uncond_velocity_action)
+ if has_action
+ else None
+ )
+ else:
+ block_state.velocity_vision = cond_velocity_vision
+ block_state.velocity_sound = cond_velocity_sound if has_sound else None
+ block_state.velocity_action = cond_velocity_action if has_action else None
+
+ return components, block_state
+
+
+class Cosmos3VisionLoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Updates vision latents after one denoising iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True, description="Noisy vision latents to update."),
+ InputParam(
+ name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("latents")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latents = components.scheduler.step(
+ block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False
+ )[0].squeeze(0)
+ return components, block_state
+
+
+class Cosmos3DistilledVisionLoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Updates vision latents after one distilled denoising iteration, re-anchoring conditioned frames."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True, description="Noisy vision latents to update."),
+ InputParam(
+ name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity."
+ ),
+ InputParam(
+ name="vision_condition_mask",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Mask marking conditioned vision latent frames.",
+ ),
+ InputParam(
+ name="vision_conditioning_latents",
+ type_hint=torch.Tensor,
+ default=None,
+ description="Clean encoded vision latents for re-anchoring conditioned frames.",
+ ),
+ InputParam(
+ name="vision_condition_indexes_for_pack",
+ type_hint=list,
+ default=None,
+ description="Indexes of conditioned vision latent frames; non-empty for image-to-video.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("latents")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible.
+ block_state.latents = components.scheduler.step(
+ block_state.velocity_vision.unsqueeze(0),
+ t,
+ block_state.latents.unsqueeze(0),
+ generator=block_state.generator,
+ return_dict=False,
+ )[0].squeeze(0)
+
+ # Distilled checkpoints use stochastic (SDE) scheduler steps that re-noise every position.
+ # Re-anchor conditioned frames to the clean encoded reference after each step.
+ has_image_condition = bool(block_state.vision_condition_indexes_for_pack)
+ if has_image_condition and block_state.vision_conditioning_latents is not None:
+ mask = block_state.vision_condition_mask
+ reference = block_state.vision_conditioning_latents.to(block_state.latents.dtype)
+ block_state.latents = mask * reference + (1.0 - mask) * block_state.latents
+
+ return components, block_state
+
+
+class Cosmos3SoundLoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Updates sound latents after one denoising iteration."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="sound_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy sound latents to update.",
+ ),
+ InputParam(
+ name="sound_scheduler",
+ type_hint=UniPCMultistepScheduler,
+ required=True,
+ description="Scheduler used to update sound latents.",
+ ),
+ InputParam(
+ name="velocity_sound", type_hint=torch.Tensor, required=True, description="Predicted sound velocity."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("sound_latents", type_hint=torch.Tensor, description="Updated sound latents.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.sound_latents = block_state.sound_scheduler.step(
+ block_state.velocity_sound.unsqueeze(0), t, block_state.sound_latents.unsqueeze(0), return_dict=False
+ )[0].squeeze(0)
+ return components, block_state
+
+
+class Cosmos3ActionLoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Updates action latents after one denoising iteration."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="action_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Noisy action latents to update.",
+ ),
+ InputParam(
+ name="action_scheduler",
+ type_hint=UniPCMultistepScheduler,
+ required=True,
+ description="Scheduler used to update action latents.",
+ ),
+ InputParam(
+ name="velocity_action", type_hint=torch.Tensor, required=True, description="Predicted action velocity."
+ ),
+ InputParam(
+ name="action_condition_mask",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Mask marking conditioned action latent frames.",
+ ),
+ InputParam(
+ name="raw_action_dim_resolved",
+ type_hint=int,
+ default=None,
+ description="Unpadded action-vector dimension.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("action_latents", type_hint=torch.Tensor, description="Updated action latents.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ has_noisy_action = block_state.action_condition_mask.sum() < block_state.action_condition_mask.numel()
+ if has_noisy_action:
+ block_state.action_latents = block_state.action_scheduler.step(
+ block_state.velocity_action.unsqueeze(0), t, block_state.action_latents.unsqueeze(0), return_dict=False
+ )[0].squeeze(0)
+ if block_state.raw_action_dim_resolved is not None:
+ block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0
+ return components, block_state
+
+
+class Cosmos3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Iteratively denoises Cosmos3 latents over scheduler timesteps."
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("timesteps", required=True),
+ InputParam.template("num_inference_steps", required=True),
+ InputParam(
+ name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3VisionDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ block_classes = [
+ Cosmos3VisionLoopPrepareStep,
+ Cosmos3LoopDenoiser,
+ Cosmos3VisionLoopSchedulerStep,
+ ]
+ block_names = ["prepare_vision", "denoiser", "update_vision"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the vision-only Cosmos3 denoising loop."
+
+
+class Cosmos3DistilledVisionDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3VisionLoopPrepareStep,
+ Cosmos3LoopDenoiser,
+ Cosmos3DistilledVisionLoopSchedulerStep,
+ ]
+ block_names = ["prepare_vision", "denoiser", "update_vision"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the vision-only distilled Cosmos3 denoising loop."
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
+ ]
+
+
+class Cosmos3VisionSoundDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ block_classes = [
+ Cosmos3VisionLoopPrepareStep,
+ Cosmos3SoundLoopPrepareStep,
+ Cosmos3LoopDenoiser,
+ Cosmos3VisionLoopSchedulerStep,
+ Cosmos3SoundLoopSchedulerStep,
+ ]
+ block_names = ["prepare_vision", "prepare_sound", "denoiser", "update_vision", "update_sound"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the vision-and-sound Cosmos3 denoising loop."
+
+
+class Cosmos3VisionActionDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ block_classes = [
+ Cosmos3VisionLoopPrepareStep,
+ Cosmos3ActionLoopPrepareStep,
+ Cosmos3LoopDenoiser,
+ Cosmos3VisionLoopSchedulerStep,
+ Cosmos3ActionLoopSchedulerStep,
+ ]
+ block_names = ["prepare_vision", "prepare_action", "denoiser", "update_vision", "update_action"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the vision-and-action Cosmos3 denoising loop."
+
+
+class Cosmos3VisionSoundActionDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ block_classes = [
+ Cosmos3VisionLoopPrepareStep,
+ Cosmos3SoundLoopPrepareStep,
+ Cosmos3ActionLoopPrepareStep,
+ Cosmos3LoopDenoiser,
+ Cosmos3VisionLoopSchedulerStep,
+ Cosmos3SoundLoopSchedulerStep,
+ Cosmos3ActionLoopSchedulerStep,
+ ]
+ block_names = [
+ "prepare_vision",
+ "prepare_sound",
+ "prepare_action",
+ "denoiser",
+ "update_vision",
+ "update_sound",
+ "update_action",
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Runs the vision, sound, and action Cosmos3 denoising loop."
+
+
+class Cosmos3TransferLoopPrepareStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares the full [control..., target] and target-only vision token lists plus timesteps for one transfer iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="control_latents",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Clean control latents for this chunk, one per hint in canonical order.",
+ ),
+ InputParam(
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to denoise."
+ ),
+ InputParam(
+ name="num_noisy_vision_tokens",
+ type_hint=int,
+ required=True,
+ description="Number of noisy target vision tokens denoised each step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "vision_tokens_full",
+ type_hint=list[torch.Tensor],
+ description="Token list for the [control..., target] forward passes.",
+ ),
+ OutputParam(
+ "vision_tokens_target",
+ type_hint=list[torch.Tensor],
+ description="Token list for the target-only (no-control) forward pass.",
+ ),
+ OutputParam(
+ "vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the noisy target tokens."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ device = components._execution_device
+ dtype = components.transformer.dtype
+ block_state.vision_tokens_full = [c.to(device=device, dtype=dtype) for c in block_state.control_latents] + [
+ block_state.latents.to(device=device, dtype=dtype)
+ ]
+ block_state.vision_tokens_target = [block_state.latents.to(device=device, dtype=dtype)]
+ block_state.vision_timesteps = torch.full((block_state.num_noisy_vision_tokens,), t.item(), device=device)
+ return components, block_state
+
+
+class Cosmos3TransferLoopDenoiser(ModularPipelineBlocks):
+ # Dedicated (not Cosmos3LoopDenoiser): transfer runs up to 3 passes over different token sequences with nested
+ # control/text CFG and interval gating, which the generic cond/uncond denoiser cannot express.
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Predicts the transfer velocity with nested control/text CFG over [control..., target]. Each branch is "
+ "gated by its guidance interval, and the result is masked so conditioned frames get zero velocity."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ # The three pre-packed CFG sequence variants (cond_full / cond_no_control / uncond_full) flow in as
+ # denoiser_input_fields, gathered generically like the other Cosmos3 denoisers.
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ name="vision_tokens_full",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Token list for the [control..., target] forward passes.",
+ ),
+ InputParam(
+ name="vision_tokens_target",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="Token list for the target-only (no-control) forward pass.",
+ ),
+ InputParam(
+ name="vision_timesteps",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Timesteps for the noisy target tokens.",
+ ),
+ InputParam(
+ name="velocity_mask",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
+ ),
+ InputParam(
+ name="guidance_scale",
+ type_hint=float,
+ default=6.0,
+ description="Scale for text classifier-free guidance.",
+ ),
+ InputParam(
+ name="control_guidance",
+ type_hint=float,
+ default=1.0,
+ description="Scale for the control (structural) guidance axis.",
+ ),
+ InputParam(
+ name="guidance_interval",
+ type_hint=tuple,
+ default=None,
+ description="Timestep interval [lo, hi] over which text guidance is active (None = always).",
+ ),
+ InputParam(
+ name="control_guidance_interval",
+ type_hint=tuple,
+ default=None,
+ description="Timestep interval [lo, hi] over which control guidance is active (None = always).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")]
+
+ @staticmethod
+ def _forward(components, static, vision_tokens, vision_timesteps):
+ preds_vision, _, _ = components.transformer(
+ input_ids=static["input_ids"],
+ text_indexes=static["text_indexes"],
+ position_ids=static["position_ids"],
+ und_len=static["und_len"],
+ sequence_length=static["sequence_length"],
+ vision_tokens=vision_tokens,
+ vision_token_shapes=static["vision_token_shapes"],
+ vision_sequence_indexes=static["vision_sequence_indexes"],
+ vision_mse_loss_indexes=static["vision_mse_loss_indexes"],
+ vision_timesteps=vision_timesteps,
+ vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"],
+ return_dict=False,
+ )
+ return preds_vision[-1]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # active-at: a None interval is always active; otherwise the timestep must fall within [lo, hi].
+ guidance_interval = block_state.guidance_interval
+ guidance_active = guidance_interval is None or (
+ float(guidance_interval[0]) <= float(t.item()) <= float(guidance_interval[1])
+ )
+ control_interval = block_state.control_guidance_interval
+ control_active = control_interval is None or (
+ float(control_interval[0]) <= float(t.item()) <= float(control_interval[1])
+ )
+ step_guidance = block_state.guidance_scale if guidance_active else 1.0
+ step_control = block_state.control_guidance if control_active else 1.0
+ needs_text_cfg = step_guidance > 1.0
+ needs_control_cfg = step_control != 1.0
+
+ denoiser_input_fields = block_state.denoiser_input_fields
+ cond_full_static = denoiser_input_fields["cond_full_static"]
+ cond_no_control_static = denoiser_input_fields["cond_no_control_static"]
+ uncond_full_static = denoiser_input_fields["uncond_full_static"]
+
+ cond_full = self._forward(
+ components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps
+ )
+
+ cond_no_control = None
+ if needs_control_cfg:
+ cond_no_control = self._forward(
+ components,
+ cond_no_control_static,
+ block_state.vision_tokens_target,
+ block_state.vision_timesteps,
+ )
+
+ uncond_full = None
+ if needs_text_cfg:
+ uncond_full = self._forward(
+ components,
+ uncond_full_static,
+ block_state.vision_tokens_full,
+ block_state.vision_timesteps,
+ )
+
+ if needs_control_cfg and needs_text_cfg:
+ control_cond = cond_no_control + step_control * (cond_full - cond_no_control)
+ velocity = uncond_full + step_guidance * (control_cond - uncond_full)
+ elif needs_control_cfg:
+ velocity = cond_no_control + step_control * (cond_full - cond_no_control)
+ elif needs_text_cfg:
+ velocity = uncond_full + step_guidance * (cond_full - uncond_full)
+ else:
+ velocity = cond_full
+
+ block_state.velocity = velocity * block_state.velocity_mask
+ return components, block_state
+
+
+class Cosmos3TransferLoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Steps the scheduler and re-pins the conditioned frames exactly for one transfer iteration."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to update."
+ ),
+ InputParam(
+ name="velocity",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Predicted (masked) transfer velocity.",
+ ),
+ InputParam(
+ name="velocity_mask",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
+ ),
+ InputParam(
+ name="condition_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Clean target latents on the conditioned frames (the autoregressive seed).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="Updated target latents for this chunk.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latents = components.scheduler.step(
+ block_state.velocity.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False
+ )[0].squeeze(0)
+ # Re-pin conditioned frames exactly (the autoregressive seed), guarding multistep drift.
+ block_state.latents = (
+ block_state.velocity_mask * block_state.latents
+ + (1.0 - block_state.velocity_mask) * block_state.condition_latents
+ )
+ return components, block_state
+
+
+# auto_docstring
+class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
+ """
+ Runs the per-chunk transfer denoising loop over scheduler timesteps.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ Timesteps for the denoising process.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ num_warmup_steps (`int`):
+ Number of scheduler warmup steps.
+ control_latents (`list`):
+ Clean control latents for this chunk, one per hint in canonical order.
+ latents (`Tensor`):
+ Noisy target latents to denoise.
+ num_noisy_vision_tokens (`int`):
+ Number of noisy target vision tokens denoised each step.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ velocity_mask (`Tensor`):
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for text classifier-free guidance.
+ control_guidance (`float`, *optional*, defaults to 1.0):
+ Scale for the control (structural) guidance axis.
+ guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
+ control_guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
+ latents (`Tensor`):
+ Noisy target latents to update.
+ condition_latents (`Tensor`):
+ Clean target latents on the conditioned frames (the autoregressive seed).
+
+ Outputs:
+ latents (`Tensor`):
+ Updated target latents for this chunk.
+ """
+
+ block_classes = [
+ Cosmos3TransferLoopPrepareStep,
+ Cosmos3TransferLoopDenoiser,
+ Cosmos3TransferLoopSchedulerStep,
+ ]
+ block_names = ["prepare_transfer", "denoiser", "update_transfer"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the per-chunk transfer denoising loop over scheduler timesteps."
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b3774e9d0f04de75bb9b6e8012a4dfc5cc13d9b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py
@@ -0,0 +1,1056 @@
+import torch
+from transformers import AutoTokenizer
+
+from ...configuration_utils import FrozenDict
+from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
+from ...pipelines.cosmos.pipeline_cosmos3_omni import (
+ _ACTION_RESOLUTION_BINS,
+ CosmosActionCondition,
+)
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+# Transfer conditions on control signals (edge/blur/depth/seg/wsm), so it uses its own system prompt instead of the
+# plain image/video ones. Defined here (not on the task pipeline) so the transfer text block is self-contained.
+_SYSTEM_PROMPT_TRANSFER = (
+ "You are a helpful assistant that generates images or videos following the user's instructions"
+ " and control signals (edge maps, blur, depth, or segmentation)."
+)
+
+
+class Cosmos3TextEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares non-action prompt token IDs for downstream text-segment packing."
+
+ @staticmethod
+ def _check_inputs(block_state) -> None:
+ prompt = block_state.prompt
+ negative_prompt = block_state.negative_prompt
+
+ if not isinstance(prompt, str):
+ raise ValueError(
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
+ )
+ if negative_prompt is not None and not isinstance(negative_prompt, str):
+ raise ValueError(
+ "`negative_prompt` must be a str or None; batched prompts are not supported, "
+ f"got {type(negative_prompt).__name__}."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_tokenizer", AutoTokenizer),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="default_use_system_prompt", default=True),
+ ConfigSpec(name="enable_safety_checker", default=True),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
+ InputParam.template(
+ "negative_prompt", description="The negative text prompt used for classifier-free guidance."
+ ),
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
+ InputParam(
+ name="height",
+ type_hint=int,
+ default=None,
+ description="Height of the generated video or image in pixels.",
+ ),
+ InputParam(
+ name="width",
+ type_hint=int,
+ default=None,
+ description="Width of the generated video or image in pixels.",
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ InputParam(
+ name="use_system_prompt",
+ type_hint=bool | None,
+ default=None,
+ description="Whether to prepend the Cosmos3 system prompt.",
+ ),
+ InputParam(
+ name="add_resolution_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add resolution metadata to the prompt.",
+ ),
+ InputParam(
+ name="add_duration_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add duration metadata to the prompt.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
+ OutputParam(
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if block_state.num_frames is None:
+ block_state.num_frames = 189
+ if block_state.height is None:
+ block_state.height = 720
+ if block_state.width is None:
+ block_state.width = 1280
+ if block_state.use_system_prompt is None:
+ block_state.use_system_prompt = components.config.default_use_system_prompt
+
+ self._check_inputs(block_state)
+ if components.requires_safety_checker:
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ device = components._execution_device
+ components.safety_checker.to(device)
+ try:
+ if not components.safety_checker.check_text_safety(block_state.prompt):
+ raise ValueError(
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
+ )
+ finally:
+ components.safety_checker.to("cpu")
+
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
+ block_state.prompt,
+ block_state.negative_prompt,
+ num_frames=block_state.num_frames,
+ height=block_state.height,
+ width=block_state.width,
+ fps=block_state.fps,
+ use_system_prompt=block_state.use_system_prompt,
+ add_resolution_template=block_state.add_resolution_template,
+ add_duration_template=block_state.add_duration_template,
+ action_mode=None,
+ action_view_point=None,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferTextStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Tokenizes the transfer prompt with the transfer system prompt. Transfer prompts are pre-upsampled JSON "
+ "captions passed through verbatim (no resolution/duration templates), so this is self-contained and does "
+ "not reuse the standard text step."
+ )
+
+ @staticmethod
+ def _check_inputs(block_state) -> None:
+ prompt = block_state.prompt
+ negative_prompt = block_state.negative_prompt
+
+ if not isinstance(prompt, (str, list)) or (
+ isinstance(prompt, list) and not all(isinstance(p, str) for p in prompt)
+ ):
+ raise ValueError(f"`prompt` must be a str or list of str, got {type(prompt).__name__}.")
+ if negative_prompt is not None and not isinstance(negative_prompt, (str, list)):
+ raise ValueError(
+ f"`negative_prompt` must be a str, list of str, or None, got {type(negative_prompt).__name__}."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_tokenizer", AutoTokenizer),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="prompt",
+ type_hint=str,
+ required=True,
+ description="The text prompt that guides Cosmos3 generation.",
+ ),
+ InputParam(
+ name="negative_prompt",
+ type_hint=str,
+ default=None,
+ description="The negative text prompt used for classifier-free guidance.",
+ ),
+ InputParam(
+ name="use_system_prompt",
+ type_hint=bool,
+ default=True,
+ description="Whether to prepend the Cosmos3 transfer system prompt.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
+ OutputParam(
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self._check_inputs(block_state)
+
+ if isinstance(block_state.prompt, list):
+ block_state.prompt = block_state.prompt[0]
+ if isinstance(block_state.negative_prompt, list):
+ block_state.negative_prompt = block_state.negative_prompt[0]
+
+ if components.requires_safety_checker:
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ device = components._execution_device
+ components.safety_checker.to(device)
+ try:
+ if not components.safety_checker.check_text_safety(block_state.prompt):
+ raise ValueError(
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
+ )
+ finally:
+ components.safety_checker.to("cpu")
+
+ # Transfer prompts are pre-upsampled JSON captions: tokenize them verbatim (no resolution/duration templates)
+ # under the transfer system prompt. Kept self-contained here rather than adding a flag to the standard step.
+ negative_prompt = block_state.negative_prompt if block_state.negative_prompt is not None else ""
+ special_tokens = components.llm_special_tokens
+
+ def _tokenize(text: str) -> list[int]:
+ conversations = []
+ if block_state.use_system_prompt:
+ conversations.append({"role": "system", "content": _SYSTEM_PROMPT_TRANSFER})
+ conversations.append({"role": "user", "content": text})
+ encoding = components.text_tokenizer.apply_chat_template(
+ conversations,
+ tokenize=True,
+ add_generation_prompt=True,
+ add_vision_id=False,
+ return_dict=True,
+ )
+ return list(encoding.input_ids) + [
+ special_tokens["eos_token_id"],
+ special_tokens["start_of_generation"],
+ ]
+
+ block_state.cond_input_ids = _tokenize(block_state.prompt)
+ block_state.uncond_input_ids = _tokenize(negative_prompt)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3DistilledTextEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepares distilled prompt token IDs. Classifier-free guidance is baked into the weights, so "
+ "`negative_prompt` is not exposed and the unconditional branch is derived from an empty prompt."
+ )
+
+ @staticmethod
+ def _check_inputs(block_state) -> None:
+ prompt = block_state.prompt
+ if not isinstance(prompt, str):
+ raise ValueError(
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_tokenizer", AutoTokenizer),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="default_use_system_prompt", default=True),
+ ConfigSpec(name="enable_safety_checker", default=True),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
+ InputParam(
+ name="height",
+ type_hint=int,
+ default=None,
+ description="Height of the generated video or image in pixels.",
+ ),
+ InputParam(
+ name="width",
+ type_hint=int,
+ default=None,
+ description="Width of the generated video or image in pixels.",
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ InputParam(
+ name="use_system_prompt",
+ type_hint=bool,
+ default=True,
+ description="Whether to prepend the Cosmos3 system prompt.",
+ ),
+ InputParam(
+ name="add_resolution_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add resolution metadata to the prompt.",
+ ),
+ InputParam(
+ name="add_duration_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add duration metadata to the prompt.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
+ OutputParam(
+ "uncond_input_ids",
+ type_hint=torch.Tensor,
+ description="Token IDs for the unconditional prompt (empty prompt; guidance is baked in).",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if block_state.num_frames is None:
+ block_state.num_frames = 189
+ if block_state.height is None:
+ block_state.height = 720
+ if block_state.width is None:
+ block_state.width = 1280
+
+ self._check_inputs(block_state)
+ if components.requires_safety_checker:
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ device = components._execution_device
+ components.safety_checker.to(device)
+ try:
+ if not components.safety_checker.check_text_safety(block_state.prompt):
+ raise ValueError(
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
+ )
+ finally:
+ components.safety_checker.to("cpu")
+
+ # Guidance is baked into distilled weights: the unconditional branch is built from an empty prompt
+ # (negative_prompt is not a user-facing input) so the downstream text-segment packing contract still holds.
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
+ block_state.prompt,
+ None,
+ num_frames=block_state.num_frames,
+ height=block_state.height,
+ width=block_state.width,
+ fps=block_state.fps,
+ use_system_prompt=block_state.use_system_prompt,
+ add_resolution_template=block_state.add_resolution_template,
+ add_duration_template=block_state.add_duration_template,
+ action_mode=None,
+ action_view_point=None,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ActionTextStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Prepares action prompt token IDs from prompt + action metadata."
+
+ @staticmethod
+ def _check_inputs(block_state) -> None:
+ prompt = block_state.prompt
+ negative_prompt = block_state.negative_prompt
+ action = block_state.action
+ num_frames = block_state.num_frames
+ height = block_state.height
+ width = block_state.width
+ if not isinstance(prompt, str):
+ raise ValueError(
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
+ )
+ if negative_prompt is not None and not isinstance(negative_prompt, str):
+ raise ValueError(
+ "`negative_prompt` must be a str or None; batched prompts are not supported, "
+ f"got {type(negative_prompt).__name__}."
+ )
+ if action is None:
+ raise ValueError("`action` is required for Cosmos3ActionTextStep.")
+ if action.image is None and action.video is None:
+ raise ValueError("`action.image` or `action.video` must be provided for action-conditioned generation.")
+ if num_frames is not None:
+ raise ValueError("`num_frames` has to be None if action is not None.")
+ if height is not None or width is not None:
+ raise ValueError("`height` and `width` have to be None if action is not None.")
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_tokenizer", AutoTokenizer),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="default_use_system_prompt", default=True),
+ ConfigSpec(name="enable_safety_checker", default=True),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
+ InputParam.template(
+ "negative_prompt", description="The negative text prompt used for classifier-free guidance."
+ ),
+ InputParam(
+ name="action",
+ type_hint=CosmosActionCondition,
+ required=True,
+ description="Action-conditioning metadata and its reference visual input.",
+ ),
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
+ InputParam(
+ name="height",
+ type_hint=int,
+ default=None,
+ description="Height of the generated video or image in pixels.",
+ ),
+ InputParam(
+ name="width",
+ type_hint=int,
+ default=None,
+ description="Width of the generated video or image in pixels.",
+ ),
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
+ InputParam(
+ name="use_system_prompt",
+ type_hint=bool | None,
+ default=None,
+ description="Whether to prepend the Cosmos3 system prompt.",
+ ),
+ InputParam(
+ name="add_resolution_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add resolution metadata to the prompt.",
+ ),
+ InputParam(
+ name="add_duration_template",
+ type_hint=bool,
+ default=True,
+ description="Whether to add duration metadata to the prompt.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("action_mode", type_hint=str, description="Requested action-generation mode."),
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
+ OutputParam(
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self._check_inputs(block_state)
+ if block_state.use_system_prompt is None:
+ block_state.use_system_prompt = components.config.default_use_system_prompt
+
+ action = block_state.action
+ block_state.action_mode = action.mode
+ block_state.num_frames = action.chunk_size + 1
+ conditioning_clip = [action.image] if action.image is not None else action.video
+ probe = components.video_processor.preprocess_video(conditioning_clip)
+ source_h, source_w = int(probe.shape[-2]), int(probe.shape[-1])
+ resolution_key = str(action.resolution_tier)
+ block_state.height, block_state.width = VideoProcessor.classify_height_width_bin(
+ source_h, source_w, ratios=_ACTION_RESOLUTION_BINS[resolution_key]
+ )
+
+ if components.requires_safety_checker:
+ if getattr(components, "safety_checker", None) is None:
+ raise ValueError(
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
+ )
+ device = components._execution_device
+ components.safety_checker.to(device)
+ try:
+ if not components.safety_checker.check_text_safety(block_state.prompt):
+ raise ValueError(
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
+ )
+ finally:
+ components.safety_checker.to("cpu")
+
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
+ block_state.prompt,
+ block_state.negative_prompt,
+ num_frames=block_state.num_frames,
+ height=block_state.height,
+ width=block_state.width,
+ fps=block_state.fps,
+ use_system_prompt=block_state.use_system_prompt,
+ add_resolution_template=block_state.add_resolution_template,
+ add_duration_template=block_state.add_duration_template,
+ action_mode=block_state.action_mode,
+ action_view_point=action.view_point,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ImageVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Encodes non-action image-to-video conditioning into Cosmos3 vision latents."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="image", default=None, description="Reference image for image-to-video conditioning."),
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
+ InputParam(
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
+ ),
+ InputParam(
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "x0_tokens_vision",
+ type_hint=torch.Tensor,
+ description="Vision latents encoded from the conditioning image or video.",
+ ),
+ OutputParam(
+ "vision_condition_frames",
+ type_hint=list[int],
+ description="Latent-frame indexes fixed by visual conditioning.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ if block_state.image is None:
+ raise ValueError("`Cosmos3ImageVaeEncoderStep` requires an `image` input.")
+ if block_state.num_frames == 1:
+ raise ValueError(
+ "`image` conditioning requires `num_frames` > 1; image-to-image generation is not supported."
+ )
+ if block_state.num_frames < 1:
+ raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.")
+
+ sf = int(components.vae.config.scale_factor_spatial)
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
+ raise ValueError(
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
+ )
+
+ conditioning_frame_2d = components.video_processor.preprocess(
+ block_state.image, height=block_state.height, width=block_state.width
+ ).to(device=device, dtype=dtype)
+
+ vision_tensor = torch.zeros(
+ 1,
+ 3,
+ block_state.num_frames,
+ block_state.height,
+ block_state.width,
+ dtype=dtype,
+ device=device,
+ )
+ vision_tensor[:, :, 0] = conditioning_frame_2d
+ vision_tensor[:, :, 1:] = conditioning_frame_2d.unsqueeze(2).expand(-1, -1, block_state.num_frames - 1, -1, -1)
+
+ block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
+ block_state.vision_condition_frames = [0]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3VideoVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return "Encodes non-action video conditioning into Cosmos3 vision latents."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="video", default=None, description="Reference video for video-to-video conditioning."),
+ InputParam(
+ name="condition_frame_indexes_vision",
+ type_hint=tuple[int, ...] | list[int],
+ default=(0, 1),
+ description="Latent-frame indexes to preserve from the conditioning video.",
+ ),
+ InputParam(
+ name="condition_video_keep",
+ type_hint=str,
+ default="first",
+ description="Which end of a longer conditioning video to use: `first` or `last`.",
+ ),
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
+ InputParam(
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
+ ),
+ InputParam(
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "x0_tokens_vision",
+ type_hint=torch.Tensor,
+ description="Vision latents encoded from the conditioning image or video.",
+ ),
+ OutputParam(
+ "vision_condition_frames",
+ type_hint=list[int],
+ description="Latent-frame indexes fixed by visual conditioning.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ if block_state.video is None:
+ raise ValueError("`Cosmos3VideoVaeEncoderStep` requires a `video` input.")
+ if block_state.num_frames == 1:
+ raise ValueError("`video` conditioning requires `num_frames` > 1.")
+ if block_state.num_frames < 1:
+ raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.")
+
+ sf = int(components.vae.config.scale_factor_spatial)
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
+ raise ValueError(
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
+ )
+
+ if not isinstance(block_state.condition_frame_indexes_vision, (list, tuple)) or isinstance(
+ block_state.condition_frame_indexes_vision, (str, bytes)
+ ):
+ raise ValueError(
+ "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got "
+ f"{block_state.condition_frame_indexes_vision!r}."
+ )
+ if not all(isinstance(index, int) and index >= 0 for index in block_state.condition_frame_indexes_vision):
+ raise ValueError(
+ "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got "
+ f"{block_state.condition_frame_indexes_vision!r}."
+ )
+ if block_state.condition_video_keep not in {"first", "last"}:
+ raise ValueError("`condition_video_keep` must be either 'first' or 'last'.")
+
+ indexes = tuple(block_state.condition_frame_indexes_vision)
+ if not indexes:
+ raise ValueError("`condition_frame_indexes_vision` must contain at least one index.")
+ latent_t = (block_state.num_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1
+ if max(indexes) >= latent_t:
+ raise ValueError(
+ f"`condition_frame_indexes_vision` {indexes} contains an index outside the latent timeline "
+ f"(latent_frames={latent_t} for num_frames={block_state.num_frames})."
+ )
+
+ condition_indexes_vision = indexes
+ conditioning_frames_3d = components.video_processor.preprocess_video(
+ block_state.video, height=block_state.height, width=block_state.width
+ ).to(device=device, dtype=dtype)
+ temporal_compression = int(components.vae.config.scale_factor_temporal)
+ max_cond_frames = max(condition_indexes_vision) * temporal_compression + 1
+ if block_state.condition_video_keep == "first":
+ conditioning_frames_3d = conditioning_frames_3d[:, :, :max_cond_frames]
+ else:
+ conditioning_frames_3d = conditioning_frames_3d[:, :, -max_cond_frames:]
+
+ vision_tensor = torch.zeros(
+ 1,
+ 3,
+ block_state.num_frames,
+ block_state.height,
+ block_state.width,
+ dtype=dtype,
+ device=device,
+ )
+ t_fill = min(conditioning_frames_3d.shape[2], block_state.num_frames)
+ vision_tensor[:, :, :t_fill] = conditioning_frames_3d[:, :, :t_fill]
+ if t_fill < block_state.num_frames:
+ vision_tensor[:, :, t_fill:] = vision_tensor[:, :, t_fill - 1 : t_fill].expand(
+ -1, -1, block_state.num_frames - t_fill, -1, -1
+ )
+ vision_condition_frames = list(condition_indexes_vision)
+
+ block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
+ block_state.vision_condition_frames = vision_condition_frames
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3TransferChunkVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Per-chunk transfer VAE encode: slices + pads this chunk's control maps, seeds the target's conditioning "
+ "frames (first chunk from the input video, later chunks from the previous chunk's tail), and encodes both "
+ "the controls and the seeded target into clean Cosmos3 vision latents. Runs inside the autoregressive "
+ "chunk loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."),
+ InputParam(
+ name="previous_output",
+ default=None,
+ description="Decoded pixels of the previous chunk, used to seed later chunks.",
+ ),
+ InputParam(
+ name="control_frames",
+ type_hint=dict,
+ required=True,
+ description="Preprocessed, time-padded control maps in canonical hint order.",
+ ),
+ InputParam(name="chunk_frames", type_hint=int, required=True, description="Pixel frames per chunk."),
+ InputParam(
+ name="total_frames", type_hint=int, required=True, description="Total number of output frames."
+ ),
+ InputParam(name="stride", type_hint=int, required=True, description="Frame stride between chunks."),
+ InputParam(
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
+ ),
+ InputParam(
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
+ ),
+ InputParam(
+ name="video",
+ default=None,
+ description="Optional input video that seeds the first chunk's conditioning.",
+ ),
+ InputParam(
+ name="num_first_chunk_conditional_frames",
+ type_hint=int,
+ default=0,
+ description="Number of frames the first chunk reuses from the input video.",
+ ),
+ InputParam(
+ name="num_conditional_frames",
+ type_hint=int,
+ default=1,
+ description="Number of frames each later chunk reuses from the previous chunk's tail.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "control_latents",
+ type_hint=list[torch.Tensor],
+ description="Clean control latents for this chunk, one per hint in canonical order.",
+ ),
+ OutputParam(
+ "x0_tokens_vision",
+ type_hint=torch.Tensor,
+ description="Clean target vision latents encoded from the seeded target frames.",
+ ),
+ OutputParam(
+ "current_conditional_frames",
+ type_hint=int,
+ description="Number of pixel frames actually used to seed this chunk's target.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ chunk_id = block_state.chunk_id
+ chunk_frames = block_state.chunk_frames
+ height = block_state.height
+ width = block_state.width
+
+ # Slice this chunk's window out of the (padded) control maps and reflect-pad it up to a full chunk (repeat the
+ # last frame once too short to keep reflecting). control_frames is already in canonical hint order.
+ start_frame = chunk_id * block_state.stride
+ end_frame = min(start_frame + chunk_frames, block_state.total_frames)
+ chunk_controls = []
+ for frames in block_state.control_frames.values():
+ frames = frames[:, :, start_frame:end_frame]
+ while frames.shape[2] < chunk_frames:
+ pad_len = min(frames.shape[2] - 1, chunk_frames - frames.shape[2])
+ if pad_len <= 0:
+ pad_frame = frames[:, :, -1:].repeat(1, 1, chunk_frames - frames.shape[2], 1, 1)
+ frames = torch.cat([frames, pad_frame], dim=2)
+ break
+ frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
+ chunk_controls.append(frames)
+
+ # Seed the target with conditioning frames (first chunk from the input video, later chunks from the
+ # previous chunk's tail), repeat-padding the remaining frames so the whole clip is well-defined.
+ target = torch.zeros(1, 3, chunk_frames, height, width, device=device, dtype=dtype)
+ current_conditional_frames = 0
+ if chunk_id == 0 and block_state.num_first_chunk_conditional_frames > 0 and block_state.video is not None:
+ input_frames = components.video_processor.preprocess_video(
+ block_state.video, height=height, width=width
+ ).to(device=device, dtype=dtype)
+ current_conditional_frames = min(
+ block_state.num_first_chunk_conditional_frames, input_frames.shape[2], chunk_frames
+ )
+ if current_conditional_frames > 0:
+ target[:, :, :current_conditional_frames] = input_frames[:, :, :current_conditional_frames]
+ elif chunk_id > 0 and block_state.previous_output is not None:
+ current_conditional_frames = min(
+ block_state.num_conditional_frames, block_state.previous_output.shape[2], chunk_frames
+ )
+ if current_conditional_frames > 0:
+ target[:, :, :current_conditional_frames] = block_state.previous_output[
+ :, :, -current_conditional_frames:
+ ].to(device=device, dtype=dtype)
+ if 0 < current_conditional_frames < chunk_frames:
+ fill = target[:, :, current_conditional_frames - 1 : current_conditional_frames]
+ target[:, :, current_conditional_frames:] = fill.expand(
+ -1, -1, chunk_frames - current_conditional_frames, -1, -1
+ )
+
+ block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls]
+ block_state.x0_tokens_vision = components._encode_video(target).contiguous().float()
+ block_state.current_conditional_frames = current_conditional_frames
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Cosmos3ActionVisionVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "cosmos3-omni"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepares action-conditioned vision latents and action frame metadata. "
+ "Only the action visual reference (image/video) is VAE-encoded; action vectors are handled separately."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="action",
+ type_hint=CosmosActionCondition,
+ required=True,
+ description="Action-conditioning metadata and its reference visual input.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "x0_tokens_vision",
+ type_hint=torch.Tensor,
+ description="Vision latents encoded from the conditioning image or video.",
+ ),
+ OutputParam(
+ "vision_condition_frames",
+ type_hint=list[int],
+ description="Latent-frame indexes fixed by visual conditioning.",
+ ),
+ OutputParam(
+ "action_condition_frame_indexes",
+ type_hint=list[int],
+ description="Action-frame indexes fixed by action conditioning.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ action = block_state.action
+ target_frames = action.chunk_size + 1
+ conditioning_clip = [action.image] if action.image is not None else action.video
+ vision_tensor, action_image_size, _, _ = components._prepare_action_video_conditioning(
+ conditioning_clip,
+ action.resolution_tier,
+ target_frames,
+ device=device,
+ dtype=dtype,
+ )
+
+ if action.mode == "forward_dynamics":
+ vision_condition_frames = [0]
+ action_condition_frame_indexes = list(range(action.chunk_size))
+ elif action.mode == "policy":
+ vision_condition_frames = [0]
+ action_condition_frame_indexes = []
+ elif action.mode == "inverse_dynamics":
+ latent_frames = (target_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1
+ vision_condition_frames = list(range(latent_frames))
+ action_condition_frame_indexes = []
+ else:
+ raise ValueError(
+ f"Unsupported action_mode={action.mode!r}; expected one of ['forward_dynamics', 'inverse_dynamics', 'policy']."
+ )
+
+ x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
+ if action_image_size is not None:
+ x0_tokens_vision = components._remove_action_video_padding_from_latent(x0_tokens_vision, action_image_size)
+
+ block_state.x0_tokens_vision = x0_tokens_vision
+ block_state.vision_condition_frames = vision_condition_frames
+ block_state.action_condition_frame_indexes = action_condition_frame_indexes
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c39fcde150a060ed012760988838a69c585812f
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py
@@ -0,0 +1,1301 @@
+import torch
+
+from ..modular_pipeline import (
+ AutoPipelineBlocks,
+ ConditionalPipelineBlocks,
+ PipelineState,
+ SequentialPipelineBlocks,
+)
+from ..modular_pipeline_utils import InputParam, OutputParam
+from .after_decode import Cosmos3ActionOutputStep
+from .before_denoise import (
+ Cosmos3ActionDenoiseInputStep,
+ Cosmos3ActionPackSequenceStep,
+ Cosmos3ActionPrepareLatentsStep,
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3SetTimestepsStep,
+ Cosmos3SoundDenoiseInputStep,
+ Cosmos3SoundPackSequenceStep,
+ Cosmos3SoundPrepareLatentsStep,
+ Cosmos3TransferPackSequenceStep,
+ Cosmos3TransferPrepareLatentsStep,
+ Cosmos3TransferSetTimestepsStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionPrepareLatentsStep,
+)
+from .before_encoder import Cosmos3TransferSetupStep
+from .decoders import (
+ Cosmos3SoundDecodeStep,
+ Cosmos3TransferDecodeChunkStep,
+ Cosmos3TransferStitchStep,
+ Cosmos3VideoDecodeStep,
+)
+from .denoise import (
+ Cosmos3TransferDenoiseStep,
+ Cosmos3VisionActionDenoiseStep,
+ Cosmos3VisionDenoiseStep,
+ Cosmos3VisionSoundActionDenoiseStep,
+ Cosmos3VisionSoundDenoiseStep,
+)
+from .encoders import (
+ Cosmos3ActionTextStep,
+ Cosmos3ActionVisionVaeEncoderStep,
+ Cosmos3ImageVaeEncoderStep,
+ Cosmos3TextEncoderStep,
+ Cosmos3TransferChunkVaeEncoderStep,
+ Cosmos3TransferTextStep,
+ Cosmos3VideoVaeEncoderStep,
+)
+from .modular_pipeline import Cosmos3OmniModularPipeline
+
+
+# auto_docstring
+class Cosmos3TransferTextBlocks(SequentialPipelineBlocks):
+ """
+ Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in
+ transfer mode using the per-chunk frame count.
+
+ Components:
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
+
+ Inputs:
+ control_videos (`dict`):
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ num_frames (`int`, *optional*):
+ Optional cap on the number of output frames (defaults to the control video length).
+ num_video_frames_per_chunk (`int`, *optional*):
+ Number of pixel frames generated per autoregressive chunk.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each chunk reuses from the previous chunk's tail.
+ prompt (`str`):
+ The text prompt that guides Cosmos3 generation.
+ negative_prompt (`str`, *optional*):
+ The negative text prompt used for classifier-free guidance.
+ use_system_prompt (`bool`, *optional*, defaults to True):
+ Whether to prepend the Cosmos3 transfer system prompt.
+
+ Outputs:
+ height (`int`):
+ Resolved output height in pixels.
+ width (`int`):
+ Resolved output width in pixels.
+ control_frames (`dict`):
+ Preprocessed, time-padded control maps in canonical hint order.
+ total_frames (`int`):
+ Total number of output frames to generate.
+ chunk_frames (`int`):
+ Number of pixel frames per autoregressive chunk.
+ num_chunks (`int`):
+ Number of autoregressive chunks.
+ stride (`int`):
+ Frame stride between consecutive chunks.
+ cond_input_ids (`Tensor`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`Tensor`):
+ Token IDs for the unconditional prompt.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3TransferSetupStep, Cosmos3TransferTextStep]
+ block_names = ["setup", "transfer_text"]
+
+ @property
+ def description(self):
+ return (
+ "Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) "
+ "prompt in transfer mode using the per-chunk frame count."
+ )
+
+
+# auto_docstring
+class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks):
+ """
+ Auto text encoder block for Cosmos3.
+ - Cosmos3TransferTextBlocks runs when control_videos are provided.
+ - Cosmos3ActionTextStep runs when action is provided.
+ - Cosmos3TextEncoderStep runs otherwise.
+
+ Components:
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
+
+ Configs:
+ default_use_system_prompt (default: True) enable_safety_checker (default: True)
+
+ Inputs:
+ control_videos (`dict`, *optional*):
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ num_frames (`int`, *optional*):
+ Optional cap on the number of output frames (defaults to the control video length).
+ num_video_frames_per_chunk (`int`, *optional*):
+ Number of pixel frames generated per autoregressive chunk.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each chunk reuses from the previous chunk's tail.
+ prompt (`str`):
+ The text prompt that guides Cosmos3 generation.
+ negative_prompt (`str`, *optional*):
+ The negative text prompt used for classifier-free guidance.
+ use_system_prompt (`bool`, *optional*, defaults to True):
+ Whether to prepend the Cosmos3 transfer system prompt.
+ action (`CosmosActionCondition`, *optional*):
+ Action-conditioning metadata and its reference visual input.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ add_resolution_template (`bool`, *optional*, defaults to True):
+ Whether to add resolution metadata to the prompt.
+ add_duration_template (`bool`, *optional*, defaults to True):
+ Whether to add duration metadata to the prompt.
+
+ Outputs:
+ height (`int`):
+ Resolved output height in pixels.
+ width (`int`):
+ Resolved output width in pixels.
+ control_frames (`dict`):
+ Preprocessed, time-padded control maps in canonical hint order.
+ total_frames (`int`):
+ Total number of output frames to generate.
+ chunk_frames (`int`):
+ Number of pixel frames per autoregressive chunk.
+ num_chunks (`int`):
+ Number of autoregressive chunks.
+ stride (`int`):
+ Frame stride between consecutive chunks.
+ cond_input_ids (`Tensor`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`Tensor`):
+ Token IDs for the unconditional prompt.
+ action_mode (`str`):
+ Requested action-generation mode.
+ num_frames (`int`):
+ Number of frames to generate.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3TransferTextBlocks, Cosmos3ActionTextStep, Cosmos3TextEncoderStep]
+ block_names = ["transfer_text", "action_text", "text"]
+ block_trigger_inputs = ["control_videos", "action", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto text encoder block for Cosmos3.\n"
+ + " - Cosmos3TransferTextBlocks runs when control_videos are provided.\n"
+ + " - Cosmos3ActionTextStep runs when action is provided.\n"
+ + " - Cosmos3TextEncoderStep runs otherwise."
+ )
+
+
+# auto_docstring
+class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks):
+ """
+ Auto VAE conditioning block for Cosmos3.
+ - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.
+ - Cosmos3VideoVaeEncoderStep runs for the non-action video path.
+ - Cosmos3ImageVaeEncoderStep runs for the non-action image path.
+ - when no action, image, or video conditioning is provided, this block is skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ action (`CosmosActionCondition`, *optional*):
+ Action-conditioning metadata and its reference visual input.
+ video (`None`, *optional*):
+ Reference video for video-to-video conditioning.
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
+ Latent-frame indexes to preserve from the conditioning video.
+ condition_video_keep (`str`, *optional*, defaults to first):
+ Which end of a longer conditioning video to use: `first` or `last`.
+ num_frames (`int`, *optional*):
+ Number of frames to generate.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ image (`None`, *optional*):
+ Reference image for image-to-video conditioning.
+
+ Outputs:
+ x0_tokens_vision (`Tensor`):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`):
+ Latent-frame indexes fixed by visual conditioning.
+ action_condition_frame_indexes (`list`):
+ Action-frame indexes fixed by action conditioning.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3ActionVisionVaeEncoderStep, Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep]
+ block_names = ["action_conditioning", "video_conditioning", "image_conditioning"]
+ block_trigger_inputs = ["action", "video", "image", "control_videos"]
+ default_block_name = None
+
+ def select_block(self, **kwargs) -> str | None:
+ action = kwargs.get("action")
+ image = kwargs.get("image")
+ video = kwargs.get("video")
+ # Transfer preprocesses/encodes its control maps inside the denoise chunk loop, so the standard VAE
+ # conditioning stage is skipped when control_videos drive the workflow.
+ if kwargs.get("control_videos") is not None:
+ return None
+ if action is not None:
+ if image is not None or video is not None:
+ raise ValueError(
+ "Pass action conditioning via `action.image` / `action.video`, not top-level image/video."
+ )
+ return "action_conditioning"
+ if image is not None and video is not None:
+ raise ValueError("Pass either image or video, not both.")
+ if video is not None:
+ return "video_conditioning"
+ if image is not None:
+ return "image_conditioning"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Auto VAE conditioning block for Cosmos3.\n"
+ + " - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.\n"
+ + " - Cosmos3VideoVaeEncoderStep runs for the non-action video path.\n"
+ + " - Cosmos3ImageVaeEncoderStep runs for the non-action image path.\n"
+ + " - when no action, image, or video conditioning is provided, this block is skipped."
+ )
+
+
+# auto_docstring
+class Cosmos3AutoSoundDecodeStep(AutoPipelineBlocks):
+ """
+ Auto sound decoder block for Cosmos3.
+ - Cosmos3SoundDecodeStep runs when sound_latents are present.
+ - if sound_latents are not provided, this block is skipped.
+
+ Components:
+ sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
+
+ Inputs:
+ sound_latents (`Tensor`, *optional*):
+ Denoised sound latents to decode.
+
+ Outputs:
+ sound (`Tensor`):
+ Generated waveform.
+ sampling_rate (`int`):
+ Sample rate of the generated waveform in Hz.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3SoundDecodeStep]
+ block_names = ["decode"]
+ block_trigger_inputs = ["sound_latents"]
+
+ @property
+ def description(self):
+ return (
+ "Auto sound decoder block for Cosmos3.\n"
+ + " - Cosmos3SoundDecodeStep runs when sound_latents are present.\n"
+ + " - if sound_latents are not provided, this block is skipped."
+ )
+
+
+# auto_docstring
+class Cosmos3DecodeStep(SequentialPipelineBlocks):
+ """
+ Decodes denoised latents into modality outputs.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
+
+ Inputs:
+ latents (`Tensor`):
+ Denoised vision latents to decode.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ sound_latents (`Tensor`, *optional*):
+ Denoised sound latents to decode.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ sound (`Tensor`):
+ Generated waveform.
+ sampling_rate (`int`):
+ Sample rate of the generated waveform in Hz.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3VideoDecodeStep, Cosmos3AutoSoundDecodeStep]
+ block_names = ["video", "sound"]
+
+ @property
+ def description(self) -> str:
+ return "Decodes denoised latents into modality outputs."
+
+
+class Cosmos3AutoDecodeStep(ConditionalPipelineBlocks):
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3TransferStitchStep, Cosmos3DecodeStep]
+ block_names = ["transfer", "standard"]
+ block_trigger_inputs = ["control_videos"]
+ default_block_name = "standard"
+
+ def select_block(self, **kwargs) -> str | None:
+ if kwargs.get("control_videos") is not None:
+ return "transfer"
+ return "standard"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Selects the Cosmos3 decode workflow.\n"
+ + " - Cosmos3TransferStitchStep stitches the decoded transfer chunks when control_videos are provided.\n"
+ + " - Cosmos3DecodeStep decodes the denoised latents otherwise."
+ )
+
+
+# auto_docstring
+class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Runs the text-and-vision Cosmos3 denoising workflow.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+
+ Configs:
+ use_native_flow_schedule (default: False)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`):
+ Number of frames to generate.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for classifier-free guidance.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionPrepareLatentsStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3SetTimestepsStep,
+ Cosmos3VisionDenoiseStep,
+ ]
+ block_names = [
+ "prepare_text_segments",
+ "prepare_vision_latents",
+ "pack_vision_sequence",
+ "prepare_vision_denoiser_inputs",
+ "set_timesteps",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Runs the text-and-vision Cosmos3 denoising workflow."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Runs the text, vision, and sound Cosmos3 denoising workflow.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+
+ Configs:
+ use_native_flow_schedule (default: False)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`):
+ Number of frames to generate.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sound_latents (`Tensor`, *optional*):
+ Pre-generated noisy sound latents.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for classifier-free guidance.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ sound_latents (`Tensor`):
+ Denoised sound latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionPrepareLatentsStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3SetTimestepsStep,
+ Cosmos3SoundPrepareLatentsStep,
+ Cosmos3SoundPackSequenceStep,
+ Cosmos3SoundDenoiseInputStep,
+ Cosmos3VisionSoundDenoiseStep,
+ ]
+ block_names = [
+ "prepare_text_segments",
+ "prepare_vision_latents",
+ "pack_vision_sequence",
+ "prepare_vision_denoiser_inputs",
+ "set_timesteps",
+ "prepare_sound_latents",
+ "pack_sound_sequence",
+ "prepare_sound_denoiser_inputs",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Runs the text, vision, and sound Cosmos3 denoising workflow."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."),
+ ]
+
+
+# auto_docstring
+class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Runs the text, vision, and action Cosmos3 denoising workflow.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+
+ Configs:
+ use_native_flow_schedule (default: False)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`):
+ Number of frames to generate.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ action (`CosmosActionCondition`):
+ Action-conditioning metadata.
+ action_condition_frame_indexes (`list`, *optional*):
+ Action-frame indexes fixed by action conditioning.
+ action_latents (`Tensor`, *optional*):
+ Pre-generated noisy action latents.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for classifier-free guidance.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ action_latents (`Tensor`):
+ Denoised action latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionPrepareLatentsStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3SetTimestepsStep,
+ Cosmos3ActionPrepareLatentsStep,
+ Cosmos3ActionPackSequenceStep,
+ Cosmos3ActionDenoiseInputStep,
+ Cosmos3VisionActionDenoiseStep,
+ ]
+ block_names = [
+ "prepare_text_segments",
+ "prepare_vision_latents",
+ "pack_vision_sequence",
+ "prepare_vision_denoiser_inputs",
+ "set_timesteps",
+ "prepare_action_latents",
+ "pack_action_sequence",
+ "prepare_action_denoiser_inputs",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Runs the text, vision, and action Cosmos3 denoising workflow."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."),
+ ]
+
+
+# auto_docstring
+class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Runs the text, vision, sound, and action Cosmos3 denoising workflow.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+
+ Configs:
+ use_native_flow_schedule (default: False)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`):
+ Number of frames to generate.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sound_latents (`Tensor`, *optional*):
+ Pre-generated noisy sound latents.
+ action (`CosmosActionCondition`):
+ Action-conditioning metadata.
+ action_condition_frame_indexes (`list`, *optional*):
+ Action-frame indexes fixed by action conditioning.
+ action_latents (`Tensor`, *optional*):
+ Pre-generated noisy action latents.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for classifier-free guidance.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ sound_latents (`Tensor`):
+ Denoised sound latents.
+ action_latents (`Tensor`):
+ Denoised action latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionPrepareLatentsStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3SetTimestepsStep,
+ Cosmos3SoundPrepareLatentsStep,
+ Cosmos3SoundPackSequenceStep,
+ Cosmos3SoundDenoiseInputStep,
+ Cosmos3ActionPrepareLatentsStep,
+ Cosmos3ActionPackSequenceStep,
+ Cosmos3ActionDenoiseInputStep,
+ Cosmos3VisionSoundActionDenoiseStep,
+ ]
+ block_names = [
+ "prepare_text_segments",
+ "prepare_vision_latents",
+ "pack_vision_sequence",
+ "prepare_vision_denoiser_inputs",
+ "set_timesteps",
+ "prepare_sound_latents",
+ "pack_sound_sequence",
+ "prepare_sound_denoiser_inputs",
+ "prepare_action_latents",
+ "pack_action_sequence",
+ "prepare_action_denoiser_inputs",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Runs the text, vision, sound, and action Cosmos3 denoising workflow."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."),
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."),
+ ]
+
+
+# auto_docstring
+class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
+ """
+ Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf
+ LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk
+ cross-carry (previous_output, output_chunks) lives on PipelineState.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`Cosmos3OmniTransformer`) scheduler
+ (`UniPCMultistepScheduler`)
+
+ Inputs:
+ chunk_id (`int`, *optional*, defaults to 0):
+ Index of the current chunk.
+ previous_output (`None`, *optional*):
+ Decoded pixels of the previous chunk, used to seed later chunks.
+ control_frames (`dict`):
+ Preprocessed, time-padded control maps in canonical hint order.
+ chunk_frames (`int`):
+ Pixel frames per chunk.
+ total_frames (`int`):
+ Total number of output frames.
+ stride (`int`):
+ Frame stride between chunks.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ video (`None`, *optional*):
+ Optional input video that seeds the first chunk's conditioning.
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
+ Number of frames the first chunk reuses from the input video.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each later chunk reuses from the previous chunk's tail.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ cond_text_segment (`dict`):
+ Conditional text segment.
+ uncond_text_segment (`dict`):
+ Unconditional text segment.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for text classifier-free guidance.
+ control_guidance (`float`, *optional*, defaults to 1.0):
+ Scale for the control (structural) guidance axis.
+ guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
+ control_guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
+ output_chunks (`list`):
+ Decoded pixel chunks accumulated so far.
+ num_chunks (`int`):
+ Number of autoregressive chunks.
+
+ Outputs:
+ control_latents (`list`):
+ Clean control latents for this chunk, one per hint in canonical order.
+ x0_tokens_vision (`Tensor`):
+ Clean target vision latents encoded from the seeded target frames.
+ current_conditional_frames (`int`):
+ Number of pixel frames actually used to seed this chunk's target.
+ latents (`Tensor`):
+ Noisy target latents for this chunk.
+ velocity_mask (`Tensor`):
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
+ condition_latents (`Tensor`):
+ Clean target latents on the conditioned frames (the autoregressive seed).
+ target_condition_indexes (`list`):
+ Latent-frame indexes fixed by the chunk's conditioning.
+ cond_full_static (`dict`):
+ Conditional [control..., target] transfer sequence carrying every control item.
+ cond_no_control_static (`dict`):
+ Conditional [target] transfer sequence with the control items dropped.
+ uncond_full_static (`dict`):
+ Unconditional [control..., target] transfer sequence for text CFG.
+ num_noisy_vision_tokens (`int`):
+ Number of noisy target vision tokens denoised each step.
+ timesteps (`Tensor`):
+ Scheduler timesteps for this chunk.
+ num_warmup_steps (`int`):
+ Number of scheduler warmup steps for this chunk.
+ vision_tokens_full (`list`):
+ Token list for the [control..., target] forward passes.
+ vision_tokens_target (`list`):
+ Token list for the target-only (no-control) forward pass.
+ vision_timesteps (`Tensor`):
+ Timesteps for the noisy target tokens.
+ velocity (`Tensor`):
+ Predicted (masked) transfer velocity.
+ previous_output (`Tensor`):
+ Decoded pixels of this chunk, used to seed the next chunk.
+ output_chunks (`list`):
+ Decoded pixel chunks accumulated so far (with this chunk appended).
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3TransferChunkVaeEncoderStep,
+ Cosmos3TransferPrepareLatentsStep,
+ Cosmos3TransferPackSequenceStep,
+ Cosmos3TransferSetTimestepsStep,
+ Cosmos3TransferDenoiseStep,
+ Cosmos3TransferDecodeChunkStep,
+ ]
+ block_names = [
+ "encode_transfer_chunk",
+ "prepare_transfer_latents",
+ "pack_transfer_sequence",
+ "set_timesteps",
+ "denoise",
+ "decode_chunk",
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a "
+ "non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). "
+ "Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return super().inputs + [
+ InputParam(name="num_chunks", type_hint=int, required=True, description="Number of autoregressive chunks.")
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+ num_chunks = state.get("num_chunks")
+ state.set("output_chunks", [])
+ state.set("previous_output", None)
+ for chunk_id in range(num_chunks):
+ state.set("chunk_id", chunk_id)
+ for _, block in self.sub_blocks.items():
+ components, state = block(components, state)
+ return components, state
+
+
+# auto_docstring
+class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
+ (`UniPCMultistepScheduler`)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ chunk_id (`int`, *optional*, defaults to 0):
+ Index of the current chunk.
+ previous_output (`None`, *optional*):
+ Decoded pixels of the previous chunk, used to seed later chunks.
+ control_frames (`dict`):
+ Preprocessed, time-padded control maps in canonical hint order.
+ chunk_frames (`int`):
+ Pixel frames per chunk.
+ total_frames (`int`):
+ Total number of output frames.
+ stride (`int`):
+ Frame stride between chunks.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ video (`None`, *optional*):
+ Optional input video that seeds the first chunk's conditioning.
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
+ Number of frames the first chunk reuses from the input video.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each later chunk reuses from the previous chunk's tail.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for text classifier-free guidance.
+ control_guidance (`float`, *optional*, defaults to 1.0):
+ Scale for the control (structural) guidance axis.
+ guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
+ control_guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
+ output_chunks (`list`):
+ Decoded pixel chunks accumulated so far.
+ num_chunks (`int`):
+ Number of autoregressive chunks.
+
+ Outputs:
+ cond_text_segment (`dict`):
+ Conditional text segment for the denoiser.
+ uncond_text_segment (`dict`):
+ Unconditional text segment for the denoiser.
+ control_latents (`list`):
+ Clean control latents for this chunk, one per hint in canonical order.
+ x0_tokens_vision (`Tensor`):
+ Clean target vision latents encoded from the seeded target frames.
+ current_conditional_frames (`int`):
+ Number of pixel frames actually used to seed this chunk's target.
+ latents (`Tensor`):
+ Noisy target latents for this chunk.
+ velocity_mask (`Tensor`):
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
+ condition_latents (`Tensor`):
+ Clean target latents on the conditioned frames (the autoregressive seed).
+ target_condition_indexes (`list`):
+ Latent-frame indexes fixed by the chunk's conditioning.
+ cond_full_static (`dict`):
+ Conditional [control..., target] transfer sequence carrying every control item.
+ cond_no_control_static (`dict`):
+ Conditional [target] transfer sequence with the control items dropped.
+ uncond_full_static (`dict`):
+ Unconditional [control..., target] transfer sequence for text CFG.
+ num_noisy_vision_tokens (`int`):
+ Number of noisy target vision tokens denoised each step.
+ timesteps (`Tensor`):
+ Scheduler timesteps for this chunk.
+ num_warmup_steps (`int`):
+ Number of scheduler warmup steps for this chunk.
+ vision_tokens_full (`list`):
+ Token list for the [control..., target] forward passes.
+ vision_tokens_target (`list`):
+ Token list for the target-only (no-control) forward pass.
+ vision_timesteps (`Tensor`):
+ Timesteps for the noisy target tokens.
+ velocity (`Tensor`):
+ Predicted (masked) transfer velocity.
+ previous_output (`Tensor`):
+ Decoded pixels of this chunk, used to seed the next chunk.
+ output_chunks (`list`):
+ Decoded pixel chunks accumulated so far (with this chunk appended).
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3TransferChunkDenoiseStep,
+ ]
+ block_names = ["prepare_text_segments", "chunk_denoise"]
+
+ @property
+ def description(self) -> str:
+ return "Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop."
+
+
+# auto_docstring
+class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ """
+ Selects the Cosmos3 core denoising workflow.
+ - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.
+ - vision_sound_action runs when action and enable_sound are provided.
+ - vision_action runs when action is provided.
+ - vision_sound runs when enable_sound is true.
+ - vision runs otherwise.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
+ (`UniPCMultistepScheduler`)
+
+ Configs:
+ use_native_flow_schedule (default: False)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ chunk_id (`int`, *optional*, defaults to 0):
+ Index of the current chunk.
+ previous_output (`None`, *optional*):
+ Decoded pixels of the previous chunk, used to seed later chunks.
+ control_frames (`dict`, *optional*):
+ Preprocessed, time-padded control maps in canonical hint order.
+ chunk_frames (`int`, *optional*):
+ Pixel frames per chunk.
+ total_frames (`int`, *optional*):
+ Total number of output frames.
+ stride (`int`, *optional*):
+ Frame stride between chunks.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ video (`None`, *optional*):
+ Optional input video that seeds the first chunk's conditioning.
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
+ Number of frames the first chunk reuses from the input video.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each later chunk reuses from the previous chunk's tail.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for text classifier-free guidance.
+ control_guidance (`float`, *optional*, defaults to 1.0):
+ Scale for the control (structural) guidance axis.
+ guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
+ control_guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
+ output_chunks (`list`, *optional*):
+ Decoded pixel chunks accumulated so far.
+ num_chunks (`int`, *optional*):
+ Number of autoregressive chunks.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`, *optional*):
+ Number of frames to generate.
+ latents (`Tensor`):
+ Pre-generated noisy vision latents.
+ sound_latents (`Tensor`, *optional*):
+ Pre-generated noisy sound latents.
+ action (`CosmosActionCondition`, *optional*):
+ Action-conditioning metadata.
+ action_condition_frame_indexes (`list`, *optional*):
+ Action-frame indexes fixed by action conditioning.
+ action_latents (`Tensor`, *optional*):
+ Pre-generated noisy action latents.
+ enable_sound (`bool`, *optional*, defaults to False):
+ Whether to generate a synchronized sound track.
+
+ Outputs:
+ cond_text_segment (`dict`):
+ Conditional text segment for the denoiser.
+ uncond_text_segment (`dict`):
+ Unconditional text segment for the denoiser.
+ control_latents (`list`):
+ Clean control latents for this chunk, one per hint in canonical order.
+ x0_tokens_vision (`Tensor`):
+ Clean target vision latents encoded from the seeded target frames.
+ current_conditional_frames (`int`):
+ Number of pixel frames actually used to seed this chunk's target.
+ latents (`Tensor`):
+ Noisy target latents for this chunk.
+ velocity_mask (`Tensor`):
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
+ condition_latents (`Tensor`):
+ Clean target latents on the conditioned frames (the autoregressive seed).
+ target_condition_indexes (`list`):
+ Latent-frame indexes fixed by the chunk's conditioning.
+ cond_full_static (`dict`):
+ Conditional [control..., target] transfer sequence carrying every control item.
+ cond_no_control_static (`dict`):
+ Conditional [target] transfer sequence with the control items dropped.
+ uncond_full_static (`dict`):
+ Unconditional [control..., target] transfer sequence for text CFG.
+ num_noisy_vision_tokens (`int`):
+ Number of noisy target vision tokens denoised each step.
+ timesteps (`Tensor`):
+ Scheduler timesteps for this chunk.
+ num_warmup_steps (`int`):
+ Number of scheduler warmup steps for this chunk.
+ vision_tokens_full (`list`):
+ Token list for the [control..., target] forward passes.
+ vision_tokens_target (`list`):
+ Token list for the target-only (no-control) forward pass.
+ vision_timesteps (`Tensor`):
+ Timesteps for the noisy target tokens.
+ velocity (`Tensor`):
+ Predicted (masked) transfer velocity.
+ previous_output (`Tensor`):
+ Decoded pixels of this chunk, used to seed the next chunk.
+ output_chunks (`list`):
+ Decoded pixel chunks accumulated so far (with this chunk appended).
+ sound_latents (`Tensor`):
+ Denoised sound latents.
+ action_latents (`Tensor`):
+ Denoised action latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3TransferCoreDenoiseStep,
+ Cosmos3VisionSoundActionCoreDenoiseStep,
+ Cosmos3VisionActionCoreDenoiseStep,
+ Cosmos3VisionSoundCoreDenoiseStep,
+ Cosmos3VisionCoreDenoiseStep,
+ ]
+ block_names = ["transfer", "vision_sound_action", "vision_action", "vision_sound", "vision"]
+ block_trigger_inputs = ["action", "enable_sound", "control_videos"]
+ default_block_name = "vision"
+
+ @property
+ def inputs(self):
+ inputs = super().inputs
+ inputs.append(
+ InputParam(
+ name="enable_sound",
+ type_hint=bool,
+ default=False,
+ description="Whether to generate a synchronized sound track.",
+ )
+ )
+ return inputs
+
+ def select_block(self, **kwargs) -> str | None:
+ action = kwargs.get("action")
+ enable_sound = kwargs.get("enable_sound")
+ if kwargs.get("control_videos") is not None:
+ return "transfer"
+ if action is not None and enable_sound:
+ return "vision_sound_action"
+ if action is not None:
+ return "vision_action"
+ if enable_sound:
+ return "vision_sound"
+ return "vision"
+
+ @property
+ def description(self):
+ return (
+ "Selects the Cosmos3 core denoising workflow.\n"
+ + " - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.\n"
+ + " - vision_sound_action runs when action and enable_sound are provided.\n"
+ + " - vision_action runs when action is provided.\n"
+ + " - vision_sound runs when enable_sound is true.\n"
+ + " - vision runs otherwise."
+ )
+
+
+# auto_docstring
+class Cosmos3OmniBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for Cosmos3 generation modes.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`, `num_frames`
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `prompt`, `image`
+ - `video2video`: requires `prompt`, `video`
+ - `text2video_with_sound`: requires `prompt`, `enable_sound`
+ - `image2video_with_sound`: requires `prompt`, `image`, `enable_sound`
+ - `video2video_with_sound`: requires `prompt`, `video`, `enable_sound`
+ - `action_policy`: requires `prompt`, `action`
+ - `action_forward_dynamics`: requires `prompt`, `action`
+ - `action_inverse_dynamics`: requires `prompt`, `action`
+
+ Components:
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer
+ (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer
+ (`Cosmos3AVAEAudioTokenizer`)
+
+ Configs:
+ default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule
+ (default: False)
+
+ Inputs:
+ control_videos (`dict`, *optional*):
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ num_frames (`int`, *optional*):
+ Optional cap on the number of output frames (defaults to the control video length).
+ num_video_frames_per_chunk (`int`, *optional*):
+ Number of pixel frames generated per autoregressive chunk.
+ num_conditional_frames (`int`, *optional*, defaults to 1):
+ Number of frames each chunk reuses from the previous chunk's tail.
+ prompt (`str`):
+ The text prompt that guides Cosmos3 generation.
+ negative_prompt (`str`, *optional*):
+ The negative text prompt used for classifier-free guidance.
+ use_system_prompt (`bool`, *optional*, defaults to True):
+ Whether to prepend the Cosmos3 transfer system prompt.
+ action (`CosmosActionCondition`, *optional*):
+ Action-conditioning metadata and its reference visual input.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ add_resolution_template (`bool`, *optional*, defaults to True):
+ Whether to add resolution metadata to the prompt.
+ add_duration_template (`bool`, *optional*, defaults to True):
+ Whether to add duration metadata to the prompt.
+ video (`None`, *optional*):
+ Reference video for video-to-video conditioning.
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
+ Latent-frame indexes to preserve from the conditioning video.
+ condition_video_keep (`str`, *optional*, defaults to first):
+ Which end of a longer conditioning video to use: `first` or `last`.
+ image (`None`, *optional*):
+ Reference image for image-to-video conditioning.
+ chunk_id (`int`, *optional*, defaults to 0):
+ Index of the current chunk.
+ previous_output (`None`, *optional*):
+ Decoded pixels of the previous chunk, used to seed later chunks.
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
+ Number of frames the first chunk reuses from the input video.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ guidance_scale (`float`, *optional*, defaults to 6.0):
+ Scale for text classifier-free guidance.
+ control_guidance (`float`, *optional*, defaults to 1.0):
+ Scale for the control (structural) guidance axis.
+ guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
+ control_guidance_interval (`tuple`, *optional*):
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
+ output_chunks (`list`, *optional*):
+ Decoded pixel chunks accumulated so far.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ latents (`Tensor`):
+ Pre-generated noisy vision latents.
+ sound_latents (`Tensor`, *optional*):
+ Pre-generated noisy sound latents.
+ action_condition_frame_indexes (`list`, *optional*):
+ Action-frame indexes fixed by action conditioning.
+ action_latents (`Tensor`, *optional*):
+ Pre-generated noisy action latents.
+ enable_sound (`bool`, *optional*, defaults to False):
+ Whether to generate a synchronized sound track.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ sound (`Tensor`):
+ Generated waveform.
+ sampling_rate (`int`):
+ Sample rate of the generated waveform in Hz.
+ action (`list`):
+ Generated action vectors.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3AutoTextEncoderStep,
+ Cosmos3AutoVaeEncoderStep,
+ Cosmos3AutoCoreDenoiseStep,
+ Cosmos3AutoDecodeStep,
+ Cosmos3ActionOutputStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True, "num_frames": 1},
+ "text2video": {"prompt": True},
+ "image2video": {"prompt": True, "image": True},
+ "video2video": {"prompt": True, "video": True},
+ "text2video_with_sound": {"prompt": True, "enable_sound": True},
+ "image2video_with_sound": {"prompt": True, "image": True, "enable_sound": True},
+ "video2video_with_sound": {"prompt": True, "video": True, "enable_sound": True},
+ "action_policy": {"prompt": True, "action": True},
+ "action_forward_dynamics": {"prompt": True, "action": True},
+ "action_inverse_dynamics": {"prompt": True, "action": True},
+ }
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for Cosmos3 generation modes."
+
+ def get_workflow(self, workflow_name: str):
+ if workflow_name == "transfer":
+ raise NotImplementedError(
+ 'The standalone "transfer" workflow is temporarily unavailable because its nested autoregressive '
+ "chunk and denoising loops cannot be preserved by the current workflow extraction logic. Transfer "
+ "remains available through the full Cosmos3OmniBlocks pipeline. The standalone workflow will be "
+ "enabled after migration to the upcoming composable nested-loop abstraction."
+ )
+ return super().get_workflow(workflow_name)
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("videos"),
+ OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."),
+ OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."),
+ OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors."),
+ ]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c654827e5403d2942030e1b39a27b85512001e0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py
@@ -0,0 +1,240 @@
+from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ Cosmos3DistilledSetTimestepsStep,
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionPrepareLatentsStep,
+)
+from .decoders import Cosmos3VideoDecodeStep
+from .denoise import Cosmos3DistilledVisionDenoiseStep
+from .encoders import (
+ Cosmos3DistilledTextEncoderStep,
+ Cosmos3ImageVaeEncoderStep,
+ Cosmos3VideoVaeEncoderStep,
+)
+
+
+# auto_docstring
+class Cosmos3DistilledAutoVaeEncoderStep(ConditionalPipelineBlocks):
+ """
+ Auto VAE conditioning block for distilled Cosmos3.
+ - Cosmos3VideoVaeEncoderStep runs for the video path.
+ - Cosmos3ImageVaeEncoderStep runs for the image path.
+ - when no image or video conditioning is provided, this block is skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ video (`None`, *optional*):
+ Reference video for video-to-video conditioning.
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
+ Latent-frame indexes to preserve from the conditioning video.
+ condition_video_keep (`str`, *optional*, defaults to first):
+ Which end of a longer conditioning video to use: `first` or `last`.
+ num_frames (`int`, *optional*):
+ Number of frames to generate.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ image (`None`, *optional*):
+ Reference image for image-to-video conditioning.
+
+ Outputs:
+ x0_tokens_vision (`Tensor`):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`):
+ Latent-frame indexes fixed by visual conditioning.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep]
+ block_names = ["video_conditioning", "image_conditioning"]
+ block_trigger_inputs = ["video", "image"]
+ default_block_name = None
+
+ def select_block(self, **kwargs) -> str | None:
+ image = kwargs.get("image")
+ video = kwargs.get("video")
+ if image is not None and video is not None:
+ raise ValueError("Pass either image or video, not both.")
+ if video is not None:
+ return "video_conditioning"
+ if image is not None:
+ return "image_conditioning"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Auto VAE conditioning block for distilled Cosmos3.\n"
+ + " - Cosmos3VideoVaeEncoderStep runs for the video path.\n"
+ + " - Cosmos3ImageVaeEncoderStep runs for the image path.\n"
+ + " - when no image or video conditioning is provided, this block is skipped."
+ )
+
+
+# auto_docstring
+class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Runs the text-and-vision distilled Cosmos3 denoising workflow.
+
+ Components:
+ transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Configs:
+ is_distilled (default: True) distilled_sigmas (default: None)
+
+ Inputs:
+ cond_input_ids (`None`):
+ Token IDs for the conditional prompt.
+ uncond_input_ids (`None`):
+ Token IDs for the unconditional prompt.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ num_frames (`int`):
+ Number of frames to generate.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*):
+ The number of denoising steps.
+ guidance_scale (`float`, *optional*):
+ Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
+ forced to 1.0. Passing a value other than 1.0 raises an error.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3PrepareTextSegmentsStep,
+ Cosmos3VisionPrepareLatentsStep,
+ Cosmos3VisionPackSequenceStep,
+ Cosmos3VisionDenoiseInputStep,
+ Cosmos3DistilledSetTimestepsStep,
+ Cosmos3DistilledVisionDenoiseStep,
+ ]
+ block_names = [
+ "prepare_text_segments",
+ "prepare_vision_latents",
+ "pack_vision_sequence",
+ "prepare_vision_denoiser_inputs",
+ "set_timesteps",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Runs the text-and-vision distilled Cosmos3 denoising workflow."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`, `num_frames`
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `prompt`, `image`
+ - `video2video`: requires `prompt`, `video`
+
+ Components:
+ text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer
+ (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Configs:
+ default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True)
+ distilled_sigmas (default: None)
+
+ Inputs:
+ prompt (`str`):
+ The text prompt that guides Cosmos3 generation.
+ num_frames (`int`, *optional*):
+ Number of frames to generate.
+ height (`int`, *optional*):
+ Height of the generated video or image in pixels.
+ width (`int`, *optional*):
+ Width of the generated video or image in pixels.
+ fps (`float`, *optional*, defaults to 24.0):
+ Frame rate of the generated video.
+ use_system_prompt (`bool`, *optional*, defaults to True):
+ Whether to prepend the Cosmos3 system prompt.
+ add_resolution_template (`bool`, *optional*, defaults to True):
+ Whether to add resolution metadata to the prompt.
+ add_duration_template (`bool`, *optional*, defaults to True):
+ Whether to add duration metadata to the prompt.
+ video (`None`, *optional*):
+ Reference video for video-to-video conditioning.
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
+ Latent-frame indexes to preserve from the conditioning video.
+ condition_video_keep (`str`, *optional*, defaults to first):
+ Which end of a longer conditioning video to use: `first` or `last`.
+ image (`None`, *optional*):
+ Reference image for image-to-video conditioning.
+ x0_tokens_vision (`Tensor`, *optional*):
+ Vision latents encoded from the conditioning image or video.
+ vision_condition_frames (`list`, *optional*):
+ Latent-frame indexes fixed by visual conditioning.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy vision latents.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*):
+ The number of denoising steps.
+ guidance_scale (`float`, *optional*):
+ Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
+ forced to 1.0. Passing a value other than 1.0 raises an error.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "cosmos3-omni"
+ block_classes = [
+ Cosmos3DistilledTextEncoderStep,
+ Cosmos3DistilledAutoVaeEncoderStep,
+ Cosmos3DistilledVisionCoreDenoiseStep,
+ Cosmos3VideoDecodeStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True, "num_frames": 1},
+ "text2video": {"prompt": True},
+ "image2video": {"prompt": True, "image": True},
+ "video2video": {"prompt": True, "video": True},
+ }
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..33d341fdf9c3511ca3a8f3be97eaf87b6cf73bec
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py
@@ -0,0 +1,130 @@
+import torch
+
+from ...pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline, CosmosSafetyChecker
+from ..modular_pipeline import ModularPipeline
+
+
+class Cosmos3OmniModularPipeline(ModularPipeline):
+ """
+ A ModularPipeline for Cosmos 3 omni generation.
+ """
+
+ default_blocks_name = "Cosmos3OmniBlocks"
+
+ duration_template = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
+ image_resolution_template = "This image is of {height}x{width} resolution."
+ video_resolution_template = "This video is of {height}x{width} resolution."
+ inverse_duration_template = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS."
+ inverse_image_resolution_template = "This image is not of {height}x{width} resolution."
+ inverse_video_resolution_template = "This video is not of {height}x{width} resolution."
+
+ @property
+ def vae_scale_factor_spatial(self):
+ if getattr(self, "vae", None) is not None:
+ return int(self.vae.config.scale_factor_spatial)
+ return 16
+
+ @property
+ def vae_scale_factor_temporal(self):
+ if getattr(self, "vae", None) is not None:
+ return int(self.vae.config.scale_factor_temporal)
+ return 4
+
+ @property
+ def num_channels_latents(self):
+ if getattr(self, "transformer", None) is not None:
+ return int(self.transformer.config.latent_channel)
+ return 48
+
+ @property
+ def sound_sampling_rate(self):
+ if getattr(self, "sound_tokenizer", None) is not None:
+ return int(self.sound_tokenizer.config.sampling_rate)
+ return 48000
+
+ @property
+ def sound_hop_size(self):
+ if getattr(self, "sound_tokenizer", None) is not None:
+ return int(self.sound_tokenizer._hop_size)
+ return 1920
+
+ @property
+ def _vae_latents_mean(self):
+ return torch.tensor(self.vae.config.latents_mean, dtype=self.vae.dtype)
+
+ @property
+ def _vae_latents_inv_std(self):
+ return 1.0 / torch.tensor(self.vae.config.latents_std, dtype=self.vae.dtype)
+
+ @property
+ def llm_special_tokens(self):
+ if getattr(self, "text_tokenizer", None) is None:
+ return None
+ return {
+ "start_of_generation": self.text_tokenizer.convert_tokens_to_ids("<|vision_start|>"),
+ "eos_token_id": self.text_tokenizer.eos_token_id,
+ }
+
+ def enable_safety_checker(self, safety_checker=None):
+ if safety_checker is not None:
+ self.safety_checker = safety_checker
+ elif getattr(self, "safety_checker", None) is None:
+ self.safety_checker = CosmosSafetyChecker()
+ self._is_safety_checker_enabled = True
+
+ def disable_safety_checker(self):
+ self._is_safety_checker_enabled = False
+
+ @property
+ def requires_safety_checker(self):
+ return getattr(self, "_is_safety_checker_enabled", self.config.enable_safety_checker)
+
+ def _encode_video(self, x):
+ return Cosmos3OmniPipeline._encode_video(self, x)
+
+ def decode_sound(self, latent):
+ return Cosmos3OmniPipeline.decode_sound(self, latent)
+
+ def _prepare_text_segment(self, input_ids, device):
+ return Cosmos3OmniPipeline._prepare_text_segment(self, input_ids, device)
+
+ def _prepare_vision_segment(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._prepare_vision_segment(self, *args, **kwargs)
+
+ def _prepare_sound_segment(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._prepare_sound_segment(self, *args, **kwargs)
+
+ def _prepare_action_segment(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._prepare_action_segment(self, *args, **kwargs)
+
+ def _prepare_action_video_conditioning(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._prepare_action_video_conditioning(self, *args, **kwargs)
+
+ def _remove_action_video_padding_from_latent(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._remove_action_video_padding_from_latent(self, *args, **kwargs)
+
+ @staticmethod
+ def _build_action_json_prompt(*args, **kwargs):
+ return Cosmos3OmniPipeline._build_action_json_prompt(*args, **kwargs)
+
+ def tokenize_prompt(self, *args, **kwargs):
+ return Cosmos3OmniPipeline.tokenize_prompt(self, *args, **kwargs)
+
+ @staticmethod
+ def _mask_velocity_predictions(*args, **kwargs):
+ return Cosmos3OmniPipeline._mask_velocity_predictions(*args, **kwargs)
+
+ def _apply_video_safety_check(self, *args, **kwargs):
+ return Cosmos3OmniPipeline._apply_video_safety_check(self, *args, **kwargs)
+
+
+class Cosmos3DistilledModularPipeline(Cosmos3OmniModularPipeline):
+ """
+ A ModularPipeline for distilled (few-step) Cosmos 3 omni generation.
+
+ Distilled checkpoints bake classifier-free guidance into the weights and sample on a fixed schedule read from the
+ pipeline's `distilled_sigmas` config (populated from `modular_model_index.json`), so `guidance_scale` and
+ `num_inference_steps` are fixed and `negative_prompt` is not supported.
+ """
+
+ default_blocks_name = "Cosmos3DistilledBlocks"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b7050c6fe5d462b88feb332637757047199ecdd7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_ernie_image"] = ["ErnieImageAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["ErnieImageModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_ernie_image import ErnieImageAutoBlocks
+ from .modular_pipeline import ErnieImageModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..58b80d86ff4fef27923a904790355b6bedf9ba4c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py
@@ -0,0 +1,270 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...models import ErnieImageTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ErnieImageModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _pad_text(
+ text_hiddens: list[torch.Tensor], device: torch.device, dtype: torch.dtype, text_in_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Pad a list of variable-length text hidden states to a common length and return (padded, lengths)."""
+ batch_size = len(text_hiddens)
+ if batch_size == 0:
+ return (
+ torch.zeros((0, 0, text_in_dim), device=device, dtype=dtype),
+ torch.zeros((0,), device=device, dtype=torch.long),
+ )
+ normalized = [t.squeeze(1).to(device).to(dtype) if t.dim() == 3 else t.to(device).to(dtype) for t in text_hiddens]
+ lengths = torch.tensor([t.shape[0] for t in normalized], device=device, dtype=torch.long)
+ max_length = int(lengths.max().item())
+ padded = torch.zeros((batch_size, max_length, text_in_dim), device=device, dtype=dtype)
+ for i, t in enumerate(normalized):
+ padded[i, : t.shape[0], :] = t
+ return padded, lengths
+
+
+class ErnieImageTextInputStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that pads the variable-length text hidden states to a common length and "
+ "produces `text_bth` / `text_lens` tensors consumed by the denoiser."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=list,
+ description="List of per-prompt text embeddings from the text encoder step.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=list,
+ description="List of per-prompt negative text embeddings from the text encoder step.",
+ ),
+ InputParam(
+ "num_images_per_prompt",
+ type_hint=int,
+ default=1,
+ description="Number of images to generate per prompt.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("batch_size", type_hint=int, description="The number of prompts in the batch."),
+ OutputParam(
+ "text_bth",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Padded text hidden states of shape (B, T_max, H) fed into the transformer.",
+ ),
+ OutputParam(
+ "text_lens",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Actual per-prompt text lengths used to build the transformer attention mask.",
+ ),
+ OutputParam(
+ "negative_text_bth",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Padded negative text hidden states, when classifier-free guidance is enabled.",
+ ),
+ OutputParam(
+ "negative_text_lens",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Actual per-prompt negative text lengths, when classifier-free guidance is enabled.",
+ ),
+ ]
+
+ @staticmethod
+ def _expand(hiddens: list[torch.Tensor], num_images_per_prompt: int) -> list[torch.Tensor]:
+ if num_images_per_prompt == 1:
+ return list(hiddens)
+ return [h for h in hiddens for _ in range(num_images_per_prompt)]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+ text_in_dim = components.text_in_dim
+ num_images_per_prompt = block_state.num_images_per_prompt
+
+ prompt_embeds = block_state.prompt_embeds
+ block_state.batch_size = len(prompt_embeds)
+
+ prompt_embeds = self._expand(prompt_embeds, num_images_per_prompt)
+ text_bth, text_lens = _pad_text(prompt_embeds, device, dtype, text_in_dim)
+ block_state.text_bth = text_bth
+ block_state.text_lens = text_lens
+
+ negative_prompt_embeds = block_state.negative_prompt_embeds
+ if negative_prompt_embeds is not None:
+ negative_prompt_embeds = self._expand(negative_prompt_embeds, num_images_per_prompt)
+ negative_text_bth, negative_text_lens = _pad_text(negative_prompt_embeds, device, dtype, text_in_dim)
+ block_state.negative_text_bth = negative_text_bth
+ block_state.negative_text_lens = negative_text_lens
+ else:
+ block_state.negative_text_bth = None
+ block_state.negative_text_lens = None
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ErnieImageSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference using a linear sigma schedule."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "num_inference_steps",
+ type_hint=int,
+ default=50,
+ description="Number of denoising steps.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference."),
+ OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ num_inference_steps = block_state.num_inference_steps
+
+ sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1]
+ components.scheduler.set_timesteps(sigmas=sigmas, device=device)
+
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.num_inference_steps = num_inference_steps
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ErnieImagePrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return "Prepare random noise latents for the ErnieImage text-to-image denoising process."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height", type_hint=int, description="The height in pixels of the generated image."),
+ InputParam("width", type_hint=int, description="The width in pixels of the generated image."),
+ InputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="Pre-generated noisy latents. If provided, skips noise sampling.",
+ ),
+ InputParam(
+ "generator",
+ type_hint=torch.Generator,
+ description="Torch generator for deterministic noise sampling.",
+ ),
+ InputParam(
+ "text_bth",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Padded text hidden states; used to derive the total batch size for the latents.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor, description="The initial noise latents to denoise."),
+ OutputParam("height", type_hint=int, description="The resolved image height in pixels."),
+ OutputParam("width", type_hint=int, description="The resolved image width in pixels."),
+ ]
+
+ @staticmethod
+ def _check_inputs(components: ErnieImageModularPipeline, height: int, width: int) -> None:
+ vae_scale_factor = components.vae_scale_factor
+ if height % vae_scale_factor != 0 or width % vae_scale_factor != 0:
+ raise ValueError(
+ f"`height` and `width` must be divisible by {vae_scale_factor}, got {height} and {width}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ self._check_inputs(components, height, width)
+
+ total_batch_size = block_state.text_bth.shape[0]
+ latent_h = height // components.vae_scale_factor
+ latent_w = width // components.vae_scale_factor
+ num_channels_latents = components.num_channels_latents
+
+ shape = (total_batch_size, num_channels_latents, latent_h, latent_w)
+ if block_state.latents is None:
+ block_state.latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=dtype)
+ else:
+ block_state.latents = block_state.latents.to(device=device, dtype=dtype)
+
+ block_state.height = height
+ block_state.width = width
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..abefc715f6f7392902ec41e039945f192e512480
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py
@@ -0,0 +1,92 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKLFlux2
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ErnieImageModularPipeline, ErnieImagePachifier
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ErnieImageVaeDecoderStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images (unpachify, BN denormalization, VAE decode)."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLFlux2),
+ ComponentSpec(
+ "pachifier",
+ ErnieImagePachifier,
+ config=FrozenDict({"patch_size": 2}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to decode into images.",
+ ),
+ InputParam(
+ "output_type",
+ type_hint=str,
+ default="pil",
+ description="Output format: 'pil', 'np', or 'pt'.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("images", type_hint=list, description="The generated images.")]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae = components.vae
+ device = block_state.latents.device
+
+ latents = block_state.latents
+ bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device=device, dtype=latents.dtype)
+ bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + 1e-5).to(device=device, dtype=latents.dtype)
+ latents = latents * bn_std + bn_mean
+
+ latents = components.pachifier.unpack_latents(latents)
+
+ images = vae.decode(latents.to(vae.dtype), return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(images, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..193135d3015de04f6ae569f9fc1add9a44ae8931
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py
@@ -0,0 +1,236 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import ErnieImageTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ErnieImageModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ErnieImageLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that prepares the latent model input and timestep tensor. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `ErnieImageDenoiseLoopWrapper`)."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to denoise.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents = block_state.latents
+ block_state.latent_model_input = latents.to(components.transformer.dtype)
+ block_state.timestep = t.expand(latents.shape[0]).to(components.transformer.dtype)
+ return components, block_state
+
+
+class ErnieImageLoopDenoiser(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", ErnieImageTransformer2DModel),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that runs the ErnieImage transformer with classifier-free guidance via "
+ "the configured guider."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "text_bth",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Padded text hidden states fed into the transformer.",
+ ),
+ InputParam(
+ "text_lens",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt text lengths used by the transformer attention mask.",
+ ),
+ InputParam(
+ "negative_text_bth",
+ type_hint=torch.Tensor,
+ description="Padded negative text hidden states for classifier-free guidance.",
+ ),
+ InputParam(
+ "negative_text_lens",
+ type_hint=torch.Tensor,
+ description="Per-prompt negative text lengths for classifier-free guidance.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="Total number of denoising steps. Used by the guider for step-aware scheduling.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ guider_inputs = {
+ "text_bth": (block_state.text_bth, block_state.negative_text_bth),
+ "text_lens": (block_state.text_lens, block_state.negative_text_lens),
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs.keys()}
+ noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=block_state.timestep,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ guider_state_batch.noise_pred = noise_pred
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state)[0]
+ return components, block_state
+
+
+class ErnieImageLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that updates the latents using the scheduler step."
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred, t, block_state.latents, return_dict=False
+ )[0]
+ if block_state.latents.dtype != latents_dtype and torch.backends.mps.is_available():
+ block_state.latents = block_state.latents.to(latents_dtype)
+ return components, block_state
+
+
+class ErnieImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoises the latents over `timesteps`. "
+ "The specific steps within each iteration can be customized with `sub_blocks` attribute."
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", ErnieImageTransformer2DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for inference.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of denoising steps.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents.")]
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ progress_bar.update()
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ErnieImageDenoiseStep(ErnieImageDenoiseLoopWrapper):
+ block_classes = [
+ ErnieImageLoopBeforeDenoiser,
+ ErnieImageLoopDenoiser,
+ ErnieImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents. At each iteration it runs:\n"
+ " - `ErnieImageLoopBeforeDenoiser`\n"
+ " - `ErnieImageLoopDenoiser`\n"
+ " - `ErnieImageLoopAfterDenoiser`"
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2c4bc402395e4f651e3c7128a04fb5365352050
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py
@@ -0,0 +1,264 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import torch
+from transformers import AutoTokenizer, Mistral3Model
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...utils import logging
+from ...utils.import_utils import is_transformers_version
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ErnieImageModularPipeline
+
+
+if is_transformers_version("<", "5.0.0"):
+ raise ImportError("`ErnieImageModularPipeline` requires `transformers>=5.0.0` for `Ministral3ForCausalLM`.")
+
+from transformers import Ministral3ForCausalLM # noqa: E402
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ErnieImagePromptEnhancerStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return "Prompt enhancer step that rewrites the input prompt using a causal language model (PE)."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pe", Ministral3ForCausalLM),
+ ComponentSpec("pe_tokenizer", AutoTokenizer),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "prompt",
+ required=True,
+ type_hint=str,
+ description="The prompt or prompts to guide image generation.",
+ ),
+ InputParam("height", type_hint=int, description="The height in pixels of the generated image."),
+ InputParam("width", type_hint=int, description="The width in pixels of the generated image."),
+ InputParam(
+ "pe_system_prompt",
+ type_hint=str,
+ default=None,
+ description="Optional system prompt passed to the prompt enhancer.",
+ ),
+ InputParam(
+ "pe_temperature",
+ type_hint=float,
+ default=0.6,
+ description="Sampling temperature used when generating with the prompt enhancer.",
+ ),
+ InputParam(
+ "pe_top_p",
+ type_hint=float,
+ default=0.95,
+ description="Nucleus sampling `top_p` used when generating with the prompt enhancer.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("prompt", type_hint=list, description="The prompt list after prompt-enhancer rewriting."),
+ OutputParam("height", type_hint=int, description="The resolved image height in pixels."),
+ OutputParam("width", type_hint=int, description="The resolved image width in pixels."),
+ ]
+
+ @staticmethod
+ def _enhance_prompt(
+ pe: Ministral3ForCausalLM,
+ pe_tokenizer: AutoTokenizer,
+ prompt: str,
+ device: torch.device,
+ width: int,
+ height: int,
+ system_prompt: str | None,
+ temperature: float,
+ top_p: float,
+ ) -> str:
+ user_content = json.dumps({"prompt": prompt, "width": width, "height": height}, ensure_ascii=False)
+ messages = []
+ if system_prompt is not None:
+ messages.append({"role": "system", "content": system_prompt})
+ messages.append({"role": "user", "content": user_content})
+
+ input_text = pe_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+ inputs = pe_tokenizer(input_text, return_tensors="pt").to(device)
+ output_ids = pe.generate(
+ **inputs,
+ max_new_tokens=pe_tokenizer.model_max_length,
+ do_sample=temperature != 1.0 or top_p != 1.0,
+ temperature=temperature,
+ top_p=top_p,
+ pad_token_id=pe_tokenizer.pad_token_id,
+ eos_token_id=pe_tokenizer.eos_token_id,
+ )
+ generated_ids = output_ids[0][inputs["input_ids"].shape[1] :]
+ return pe_tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ prompt = block_state.prompt
+ if isinstance(prompt, str):
+ prompt = [prompt]
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+
+ revised = [
+ self._enhance_prompt(
+ pe=components.pe,
+ pe_tokenizer=components.pe_tokenizer,
+ prompt=p,
+ device=device,
+ width=width,
+ height=height,
+ system_prompt=block_state.pe_system_prompt,
+ temperature=block_state.pe_temperature,
+ top_p=block_state.pe_top_p,
+ )
+ for p in prompt
+ ]
+
+ block_state.prompt = revised
+ block_state.height = height
+ block_state.width = width
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ErnieImageTextEncoderStep(ModularPipelineBlocks):
+ model_name = "ernie-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text encoder step that encodes prompts into variable-length hidden states for the ErnieImage transformer."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Mistral3Model),
+ ComponentSpec("tokenizer", AutoTokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt", type_hint=str, description="The prompt or prompts to guide image generation."),
+ InputParam(
+ "negative_prompt",
+ type_hint=str,
+ description="The prompt or prompts to avoid during image generation.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=list,
+ kwargs_type="denoiser_input_fields",
+ description="List of per-prompt text embeddings of shape (T, H).",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=list,
+ kwargs_type="denoiser_input_fields",
+ description="List of per-prompt negative text embeddings for classifier-free guidance.",
+ ),
+ ]
+
+ @staticmethod
+ def _encode(
+ text_encoder: Mistral3Model,
+ tokenizer: AutoTokenizer,
+ prompt: list[str],
+ device: torch.device,
+ ) -> list[torch.Tensor]:
+ text_hiddens = []
+ for p in prompt:
+ ids = tokenizer(p, add_special_tokens=True, truncation=True, padding=False)["input_ids"]
+ if len(ids) == 0:
+ ids = [tokenizer.bos_token_id if tokenizer.bos_token_id is not None else 0]
+ input_ids = torch.tensor([ids], device=device)
+ outputs = text_encoder(input_ids=input_ids, output_hidden_states=True)
+ text_hiddens.append(outputs.hidden_states[-2][0])
+ return text_hiddens
+
+ @torch.no_grad()
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ prompt = block_state.prompt
+ if prompt is None:
+ prompt = [""]
+ if isinstance(prompt, str):
+ prompt = [prompt]
+
+ block_state.prompt_embeds = self._encode(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ device=device,
+ )
+
+ if components.requires_unconditional_embeds:
+ negative_prompt = block_state.negative_prompt
+ if negative_prompt is None:
+ negative_prompt = ""
+ if isinstance(negative_prompt, str):
+ negative_prompt = [negative_prompt] * len(prompt)
+ if len(negative_prompt) != len(prompt):
+ raise ValueError(
+ f"`negative_prompt` must have the same length as `prompt` ({len(prompt)}), "
+ f"got {len(negative_prompt)}."
+ )
+ block_state.negative_prompt_embeds = self._encode(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=negative_prompt,
+ device=device,
+ )
+ else:
+ block_state.negative_prompt_embeds = None
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..06e02b7b2f709b00bbb8da2f26e76ae1e360c865
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py
@@ -0,0 +1,200 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ ErnieImagePrepareLatentsStep,
+ ErnieImageSetTimestepsStep,
+ ErnieImageTextInputStep,
+)
+from .decoders import ErnieImageVaeDecoderStep
+from .denoise import ErnieImageDenoiseStep
+from .encoders import ErnieImagePromptEnhancerStep, ErnieImageTextEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# auto_docstring
+class ErnieImageAutoPromptEnhancerStep(ConditionalPipelineBlocks):
+ """
+ Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.
+ - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.
+ - If `use_pe` is `None` or `False`, the step is skipped.
+
+ Components:
+ pe (`Ministral3ForCausalLM`) pe_tokenizer (`AutoTokenizer`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ pe_system_prompt (`str`, *optional*):
+ Optional system prompt passed to the prompt enhancer.
+ pe_temperature (`float`, *optional*, defaults to 0.6):
+ Sampling temperature used when generating with the prompt enhancer.
+ pe_top_p (`float`, *optional*, defaults to 0.95):
+ Nucleus sampling `top_p` used when generating with the prompt enhancer.
+
+ Outputs:
+ prompt (`list`):
+ The prompt list after prompt-enhancer rewriting.
+ height (`int`):
+ The resolved image height in pixels.
+ width (`int`):
+ The resolved image width in pixels.
+ """
+
+ model_name = "ernie-image"
+ block_classes = [ErnieImagePromptEnhancerStep]
+ block_names = ["prompt_enhancer"]
+ block_trigger_inputs = ["use_pe"]
+
+ def select_block(self, use_pe=None) -> str | None:
+ if use_pe:
+ return "prompt_enhancer"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.\n"
+ " - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.\n"
+ " - If `use_pe` is `None` or `False`, the step is skipped."
+ )
+
+
+# auto_docstring
+class ErnieImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block that takes encoded conditions and runs the denoising process for ErnieImage.
+
+ Components:
+ transformer (`ErnieImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt_embeds (`list`):
+ List of per-prompt text embeddings from the text encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ List of per-prompt negative text embeddings from the text encoder step.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ Number of images to generate per prompt.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ Number of denoising steps.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents. If provided, skips noise sampling.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic noise sampling.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "ernie-image"
+ block_classes = [
+ ErnieImageTextInputStep,
+ ErnieImageSetTimestepsStep,
+ ErnieImagePrepareLatentsStep,
+ ErnieImageDenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block that takes encoded conditions and runs the denoising process for ErnieImage."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class ErnieImageAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto modular pipeline for ErnieImage text-to-image generation. Supports an optional prompt enhancer when the `pe`
+ components are loaded and `use_pe=True`.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+
+ Components:
+ pe (`Ministral3ForCausalLM`) pe_tokenizer (`AutoTokenizer`) text_encoder (`Mistral3Model`) tokenizer
+ (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer (`ErnieImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLFlux2`) pachifier (`ErnieImagePachifier`)
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ pe_system_prompt (`str`, *optional*):
+ Optional system prompt passed to the prompt enhancer.
+ pe_temperature (`float`, *optional*, defaults to 0.6):
+ Sampling temperature used when generating with the prompt enhancer.
+ pe_top_p (`float`, *optional*, defaults to 0.95):
+ Nucleus sampling `top_p` used when generating with the prompt enhancer.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts to avoid during image generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ Number of images to generate per prompt.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ Number of denoising steps.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents. If provided, skips noise sampling.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic noise sampling.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', or 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "ernie-image"
+ block_classes = [
+ ErnieImageAutoPromptEnhancerStep,
+ ErnieImageTextEncoderStep,
+ ErnieImageCoreDenoiseStep,
+ ErnieImageVaeDecoderStep,
+ ]
+ block_names = ["prompt_enhancer", "text_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self):
+ return (
+ "Auto modular pipeline for ErnieImage text-to-image generation. Supports an optional prompt enhancer "
+ "when the `pe` components are loaded and `use_pe=True`."
+ )
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..44d92f2349c75723eec6818f2ef4f1e26f8a704d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_pipeline.py
@@ -0,0 +1,110 @@
+# Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...loaders import ErnieImageLoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ErnieImagePachifier(ConfigMixin):
+ """
+ A class to pack and unpack latents for ErnieImage.
+ """
+
+ config_name = "config.json"
+
+ @register_to_config
+ def __init__(self, patch_size: int = 2):
+ super().__init__()
+
+ def pack_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ batch_size, num_channels, height, width = latents.shape
+ patch_size = self.config.patch_size
+
+ if height % patch_size != 0 or width % patch_size != 0:
+ raise ValueError(
+ f"Latent height and width must be divisible by {patch_size}, but got {height} and {width}"
+ )
+
+ latents = latents.view(
+ batch_size, num_channels, height // patch_size, patch_size, width // patch_size, patch_size
+ )
+ latents = latents.permute(0, 1, 3, 5, 2, 4)
+ return latents.reshape(
+ batch_size, num_channels * patch_size * patch_size, height // patch_size, width // patch_size
+ )
+
+ def unpack_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ batch_size, num_channels, height, width = latents.shape
+ patch_size = self.config.patch_size
+
+ latents = latents.reshape(
+ batch_size, num_channels // (patch_size * patch_size), patch_size, patch_size, height, width
+ )
+ latents = latents.permute(0, 1, 4, 2, 5, 3)
+ return latents.reshape(
+ batch_size, num_channels // (patch_size * patch_size), height * patch_size, width * patch_size
+ )
+
+
+class ErnieImageModularPipeline(ModularPipeline, ErnieImageLoraLoaderMixin):
+ """
+ A ModularPipeline for ErnieImage.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "ErnieImageAutoBlocks"
+
+ @property
+ def default_height(self):
+ return 1024
+
+ @property
+ def default_width(self):
+ return 1024
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 16
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** len(self.vae.config.block_out_channels)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 128
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels
+ return num_channels_latents
+
+ @property
+ def text_in_dim(self):
+ text_in_dim = 3584
+ if hasattr(self, "transformer") and self.transformer is not None:
+ text_in_dim = self.transformer.config.text_in_dim
+ return text_in_dim
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+ return requires_unconditional_embeds
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..934e5b8a3c3e0142a06576468b5cac7b3b24fd97
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/__init__.py
@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_flux"] = ["FluxAutoBlocks"]
+ _import_structure["modular_blocks_flux_kontext"] = ["FluxKontextAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["FluxKontextModularPipeline", "FluxModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_flux import FluxAutoBlocks
+ from .modular_blocks_flux_kontext import FluxKontextAutoBlocks
+ from .modular_pipeline import FluxKontextModularPipeline, FluxModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..69c79223345d58f74b5c2447832fb78ec8b09cb3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/before_denoise.py
@@ -0,0 +1,618 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...pipelines import FluxPipeline
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import FluxModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def _get_initial_timesteps_and_optionals(
+ transformer,
+ scheduler,
+ batch_size,
+ height,
+ width,
+ vae_scale_factor,
+ num_inference_steps,
+ guidance_scale,
+ sigmas,
+ device,
+):
+ image_seq_len = (int(height) // vae_scale_factor // 2) * (int(width) // vae_scale_factor // 2)
+
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+ if hasattr(scheduler.config, "use_flow_sigmas") and scheduler.config.use_flow_sigmas:
+ sigmas = None
+ mu = calculate_shift(
+ image_seq_len,
+ scheduler.config.get("base_image_seq_len", 256),
+ scheduler.config.get("max_image_seq_len", 4096),
+ scheduler.config.get("base_shift", 0.5),
+ scheduler.config.get("max_shift", 1.15),
+ )
+ timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu)
+ if transformer.config.guidance_embeds:
+ guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+ guidance = guidance.expand(batch_size)
+ else:
+ guidance = None
+
+ return timesteps, num_inference_steps, sigmas, guidance
+
+
+class FluxSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ InputParam("guidance_scale", default=3.5),
+ InputParam("latents", type_hint=torch.Tensor),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time",
+ ),
+ OutputParam("guidance", type_hint=torch.Tensor, description="Optional guidance to be used."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ scheduler = components.scheduler
+ transformer = components.transformer
+
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+ timesteps, num_inference_steps, sigmas, guidance = _get_initial_timesteps_and_optionals(
+ transformer,
+ scheduler,
+ batch_size,
+ block_state.height,
+ block_state.width,
+ components.vae_scale_factor,
+ block_state.num_inference_steps,
+ block_state.guidance_scale,
+ block_state.sigmas,
+ block_state.device,
+ )
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = num_inference_steps
+ block_state.sigmas = sigmas
+ block_state.guidance = guidance
+
+ # We set the index here to remove DtoH sync, helpful especially during compilation.
+ # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+ components.scheduler.set_begin_index(0)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxImg2ImgSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ InputParam("strength", default=0.6),
+ InputParam("guidance_scale", default=3.5),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time",
+ ),
+ OutputParam("guidance", type_hint=torch.Tensor, description="Optional guidance to be used."),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps with self.scheduler->scheduler
+ def get_timesteps(scheduler, num_inference_steps, strength, device):
+ # get the original timestep using init_timestep
+ init_timestep = min(num_inference_steps * strength, num_inference_steps)
+
+ t_start = int(max(num_inference_steps - init_timestep, 0))
+ timesteps = scheduler.timesteps[t_start * scheduler.order :]
+ if hasattr(scheduler, "set_begin_index"):
+ scheduler.set_begin_index(t_start * scheduler.order)
+
+ return timesteps, num_inference_steps - t_start
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ scheduler = components.scheduler
+ transformer = components.transformer
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+ timesteps, num_inference_steps, sigmas, guidance = _get_initial_timesteps_and_optionals(
+ transformer,
+ scheduler,
+ batch_size,
+ block_state.height,
+ block_state.width,
+ components.vae_scale_factor,
+ block_state.num_inference_steps,
+ block_state.guidance_scale,
+ block_state.sigmas,
+ block_state.device,
+ )
+ timesteps, num_inference_steps = self.get_timesteps(
+ scheduler, num_inference_steps, block_state.strength, block_state.device
+ )
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = num_inference_steps
+ block_state.sigmas = sigmas
+ block_state.guidance = guidance
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the latents for the text-to-image generation process"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam("latents", type_hint=torch.Tensor | None),
+ InputParam("num_images_per_prompt", type_hint=int, default=1),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (block_state.height is not None and block_state.height % (components.vae_scale_factor * 2) != 0) or (
+ block_state.width is not None and block_state.width % (components.vae_scale_factor * 2) != 0
+ ):
+ logger.warning(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor} but are {block_state.height} and {block_state.width}."
+ )
+
+ @staticmethod
+ def prepare_latents(
+ comp,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ height = 2 * (int(height) // (comp.vae_scale_factor * 2))
+ width = 2 * (int(width) // (comp.vae_scale_factor * 2))
+
+ shape = (batch_size, num_channels_latents, height, width)
+
+ if latents is not None:
+ return latents.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ # TODO: move packing latents code to a patchifier similar to Qwen
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ latents = FluxPipeline._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+ block_state.device = components._execution_device
+ block_state.num_channels_latents = components.num_channels_latents
+
+ self.check_inputs(components, block_state)
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+ block_state.latents = self.prepare_latents(
+ components,
+ batch_size,
+ block_state.num_channels_latents,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ block_state.latents,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class FluxImg2ImgPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return "Step that adds noise to image latents for image-to-image. Should be run after `set_timesteps`,"
+ " `prepare_latents`. Both noise and image latents should already be patchified."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial random noised, can be generated in prepare latent step.",
+ ),
+ InputParam(
+ name="image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The image latents to use for the denoising process. Can be generated in vae encoder and packed in input step.",
+ ),
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="initial_noise",
+ type_hint=torch.Tensor,
+ description="The initial random noised used for inpainting denoising.",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(image_latents, latents):
+ if image_latents.shape[0] != latents.shape[0]:
+ raise ValueError(
+ f"`image_latents` must have have same batch size as `latents`, but got {image_latents.shape[0]} and {latents.shape[0]}"
+ )
+
+ if image_latents.ndim != 3:
+ raise ValueError(f"`image_latents` must have 3 dimensions (patchified), but got {image_latents.ndim}")
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(image_latents=block_state.image_latents, latents=block_state.latents)
+
+ # prepare latent timestep
+ latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
+
+ # make copy of initial_noise
+ block_state.initial_noise = block_state.latents
+
+ # scale noise
+ block_state.latents = components.scheduler.scale_noise(
+ block_state.image_latents, latent_timestep, block_state.latents
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class FluxRoPEInputsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the RoPE inputs for the denoising process. Should be placed after text encoder and latent preparation steps."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="height", required=True),
+ InputParam(name="width", required=True),
+ InputParam(name="prompt_embeds"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="txt_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the prompt embeds, used for RoPE calculation.",
+ ),
+ OutputParam(
+ name="img_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the image latents, used for RoPE calculation.",
+ ),
+ ]
+
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_embeds = block_state.prompt_embeds
+ device, dtype = prompt_embeds.device, prompt_embeds.dtype
+ block_state.txt_ids = torch.zeros(prompt_embeds.shape[1], 3).to(
+ device=prompt_embeds.device, dtype=prompt_embeds.dtype
+ )
+
+ height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+ width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+ block_state.img_ids = FluxPipeline._prepare_latent_image_ids(None, height // 2, width // 2, device, dtype)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class FluxKontextRoPEInputsStep(ModularPipelineBlocks):
+ model_name = "flux-kontext"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the RoPE inputs for the denoising process of Flux Kontext. Should be placed after text encoder and latent preparation steps."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="image_height"),
+ InputParam(name="image_width"),
+ InputParam(name="height"),
+ InputParam(name="width"),
+ InputParam(name="prompt_embeds"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="txt_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the prompt embeds, used for RoPE calculation.",
+ ),
+ OutputParam(
+ name="img_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the image latents, used for RoPE calculation.",
+ ),
+ ]
+
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_embeds = block_state.prompt_embeds
+ device, dtype = prompt_embeds.device, prompt_embeds.dtype
+ block_state.txt_ids = torch.zeros(prompt_embeds.shape[1], 3).to(
+ device=prompt_embeds.device, dtype=prompt_embeds.dtype
+ )
+
+ img_ids = None
+ if (
+ getattr(block_state, "image_height", None) is not None
+ and getattr(block_state, "image_width", None) is not None
+ ):
+ image_latent_height = 2 * (int(block_state.image_height) // (components.vae_scale_factor * 2))
+ image_latent_width = 2 * (int(block_state.image_width) // (components.vae_scale_factor * 2))
+ img_ids = FluxPipeline._prepare_latent_image_ids(
+ None, image_latent_height // 2, image_latent_width // 2, device, dtype
+ )
+ # image ids are the same as latent ids with the first dimension set to 1 instead of 0
+ img_ids[..., 0] = 1
+
+ height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+ width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+ latent_ids = FluxPipeline._prepare_latent_image_ids(None, height // 2, width // 2, device, dtype)
+
+ if img_ids is not None:
+ latent_ids = torch.cat([latent_ids, img_ids], dim=0)
+
+ block_state.img_ids = latent_ids
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..233e46200fc4006a180aa2b698fe64337eec92d0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/decoders.py
@@ -0,0 +1,109 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKL
+from ...utils import logging
+from ...video_processor import VaeImageProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _unpack_latents(latents, height, width, vae_scale_factor):
+ batch_size, num_patches, channels = latents.shape
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = 2 * (int(height) // (vae_scale_factor * 2))
+ width = 2 * (int(width) // (vae_scale_factor * 2))
+
+ latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+ latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+ return latents
+
+
+class FluxDecodeStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("output_type", default="pil"),
+ InputParam("height", default=1024),
+ InputParam("width", default=1024),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents from the denoising step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | torch.Tensor | np.ndarray,
+ description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae = components.vae
+
+ if not block_state.output_type == "latent":
+ latents = block_state.latents
+ latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor)
+ latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
+ block_state.images = vae.decode(latents, return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(
+ block_state.images, output_type=block_state.output_type
+ )
+ else:
+ block_state.images = block_state.latents
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2c0c6695bca393072fdbd106fcee688f18c5245
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/denoise.py
@@ -0,0 +1,322 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...models import FluxTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import FluxModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class FluxLoopDenoiser(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", FluxTransformer2DModel)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `FluxDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("joint_attention_kwargs"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "guidance",
+ required=False,
+ type_hint=torch.Tensor,
+ description="Guidance scale as a tensor",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Prompt embeddings",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pooled prompt embeddings",
+ ),
+ InputParam(
+ "txt_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="IDs computed from text sequence needed for RoPE",
+ ),
+ InputParam(
+ "img_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="IDs computed from image sequence needed for RoPE",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ noise_pred = components.transformer(
+ hidden_states=block_state.latents,
+ timestep=t.flatten() / 1000,
+ guidance=block_state.guidance,
+ encoder_hidden_states=block_state.prompt_embeds,
+ pooled_projections=block_state.pooled_prompt_embeds,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ txt_ids=block_state.txt_ids,
+ img_ids=block_state.img_ids,
+ return_dict=False,
+ )[0]
+ block_state.noise_pred = noise_pred
+
+ return components, block_state
+
+
+class FluxKontextLoopDenoiser(ModularPipelineBlocks):
+ model_name = "flux-kontext"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", FluxTransformer2DModel)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents for Flux Kontext. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `FluxDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("joint_attention_kwargs"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Image latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "guidance",
+ required=False,
+ type_hint=torch.Tensor,
+ description="Guidance scale as a tensor",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Prompt embeddings",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pooled prompt embeddings",
+ ),
+ InputParam(
+ "txt_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="IDs computed from text sequence needed for RoPE",
+ ),
+ InputParam(
+ "img_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="IDs computed from latent sequence needed for RoPE",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ latents = block_state.latents
+ latent_model_input = latents
+ image_latents = block_state.image_latents
+ if image_latents is not None:
+ latent_model_input = torch.cat([latent_model_input, image_latents], dim=1)
+
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+ noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=block_state.guidance,
+ encoder_hidden_states=block_state.prompt_embeds,
+ pooled_projections=block_state.pooled_prompt_embeds,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ txt_ids=block_state.txt_ids,
+ img_ids=block_state.img_ids,
+ return_dict=False,
+ )[0]
+ noise_pred = noise_pred[:, : latents.size(1)]
+ block_state.noise_pred = noise_pred
+
+ return components, block_state
+
+
+class FluxLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that update the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `FluxDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # Perform scheduler step using the predicted output
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class FluxDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoise the latents over `timesteps`. "
+ "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", FluxTransformer2DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class FluxDenoiseStep(FluxDenoiseLoopWrapper):
+ block_classes = [FluxLoopDenoiser, FluxLoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `FluxDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `FluxLoopDenoiser`\n"
+ " - `FluxLoopAfterDenoiser`\n"
+ "This block supports both text2image and img2img tasks."
+ )
+
+
+class FluxKontextDenoiseStep(FluxDenoiseLoopWrapper):
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextLoopDenoiser, FluxLoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `FluxDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `FluxKontextLoopDenoiser`\n"
+ " - `FluxLoopAfterDenoiser`\n"
+ "This block supports both text2image and img2img tasks."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa811647e8b8a6af582f9f2492dfe38b0df155aa
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/encoders.py
@@ -0,0 +1,480 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import html
+
+import regex as re
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist
+from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin
+from ...models import AutoencoderKL
+from ...utils import USE_PEFT_BACKEND, is_ftfy_available, logging, scale_lora_layers, unscale_lora_layers
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import FluxModularPipeline
+
+
+if is_ftfy_available():
+ import ftfy
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def basic_clean(text):
+ text = ftfy.fix_text(text)
+ text = html.unescape(html.unescape(text))
+ return text.strip()
+
+
+def whitespace_clean(text):
+ text = re.sub(r"\s+", " ", text)
+ text = text.strip()
+ return text
+
+
+def prompt_clean(text):
+ text = whitespace_clean(basic_clean(text))
+ return text
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def encode_vae_image(vae: AutoencoderKL, image: torch.Tensor, generator: torch.Generator, sample_mode="sample"):
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode)
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
+
+ image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+ return image_latents
+
+
+class FluxProcessImagesInputStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [InputParam("resized_image"), InputParam("image"), InputParam("height"), InputParam("width")]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="processed_image")]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ if block_state.resized_image is None and block_state.image is None:
+ raise ValueError("`resized_image` and `image` cannot be None at the same time")
+
+ if block_state.resized_image is None:
+ image = block_state.image
+ self.check_inputs(
+ height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
+ )
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ else:
+ width, height = block_state.resized_image[0].size
+ image = block_state.resized_image
+
+ block_state.processed_image = components.image_processor.preprocess(image=image, height=height, width=width)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxKontextProcessImagesInputStep(ModularPipelineBlocks):
+ model_name = "flux-kontext"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Image preprocess step for Flux Kontext. The preprocessed image goes to the VAE.\n"
+ "Kontext works as a T2I model, too, in case no input image is provided."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [InputParam("image"), InputParam("_auto_resize", type_hint=bool, default=True)]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="processed_image")]
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState):
+ from ...pipelines.flux.pipeline_flux_kontext import PREFERRED_KONTEXT_RESOLUTIONS
+
+ block_state = self.get_block_state(state)
+ images = block_state.image
+
+ if images is None:
+ block_state.processed_image = None
+
+ else:
+ multiple_of = components.image_processor.config.vae_scale_factor
+
+ if not is_valid_image_imagelist(images):
+ raise ValueError(f"Images must be image or list of images but are {type(images)}")
+
+ if is_valid_image(images):
+ images = [images]
+
+ img = images[0]
+ image_height, image_width = components.image_processor.get_default_height_width(img)
+ aspect_ratio = image_width / image_height
+ _auto_resize = block_state._auto_resize
+ if _auto_resize:
+ # Kontext is trained on specific resolutions, using one of them is recommended
+ _, image_width, image_height = min(
+ (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
+ )
+ image_width = image_width // multiple_of * multiple_of
+ image_height = image_height // multiple_of * multiple_of
+ images = components.image_processor.resize(images, image_height, image_width)
+ block_state.processed_image = components.image_processor.preprocess(images, image_height, image_width)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ def __init__(
+ self, input_name: str = "processed_image", output_name: str = "image_latents", sample_mode: str = "sample"
+ ):
+ """Initialize a VAE encoder step for converting images to latent representations.
+
+ Both the input and output names are configurable so this block can be configured to process to different image
+ inputs (e.g., "processed_image" -> "image_latents", "processed_control_image" -> "control_image_latents").
+
+ Args:
+ input_name (str, optional): Name of the input image tensor. Defaults to "processed_image".
+ Examples: "processed_image" or "processed_control_image"
+ output_name (str, optional): Name of the output latent tensor. Defaults to "image_latents".
+ Examples: "image_latents" or "control_image_latents"
+ sample_mode (str, optional): Sampling mode to be used.
+
+ Examples:
+ # Basic usage with default settings (includes image processor): # FluxImageVaeEncoderDynamicStep()
+
+ # Custom input/output names for control image: # FluxImageVaeEncoderDynamicStep(
+ input_name="processed_control_image", output_name="control_image_latents"
+ )
+ """
+ self._image_input_name = input_name
+ self._image_latents_output_name = output_name
+ self.sample_mode = sample_mode
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ components = [ComponentSpec("vae", AutoencoderKL)]
+ return components
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [InputParam(self._image_input_name), InputParam("generator")]
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ self._image_latents_output_name,
+ type_hint=torch.Tensor,
+ description="The latents representing the reference image",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ image = getattr(block_state, self._image_input_name)
+
+ if image is None:
+ setattr(block_state, self._image_latents_output_name, None)
+ else:
+ device = components._execution_device
+ dtype = components.vae.dtype
+ image = image.to(device=device, dtype=dtype)
+
+ # Encode image into latents
+ image_latents = encode_vae_image(
+ image=image, vae=components.vae, generator=block_state.generator, sample_mode=self.sample_mode
+ )
+ setattr(block_state, self._image_latents_output_name, image_latents)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class FluxTextEncoderStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generate text_embeddings to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", CLIPTextModel),
+ ComponentSpec("tokenizer", CLIPTokenizer),
+ ComponentSpec("text_encoder_2", T5EncoderModel),
+ ComponentSpec("tokenizer_2", T5TokenizerFast),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("prompt_2"),
+ InputParam("max_sequence_length", type_hint=int, default=512, required=False),
+ InputParam("joint_attention_kwargs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "pooled_prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="pooled text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ for prompt in [block_state.prompt, block_state.prompt_2]:
+ if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` or `prompt_2` has to be of type `str` or `list` but is {type(prompt)}")
+
+ @staticmethod
+ def _get_t5_prompt_embeds(components, prompt: str | list[str], max_sequence_length: int, device: torch.device):
+ dtype = components.text_encoder_2.dtype
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ if isinstance(components, TextualInversionLoaderMixin):
+ prompt = components.maybe_convert_prompt(prompt, components.tokenizer_2)
+
+ text_inputs = components.tokenizer_2(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_length=False,
+ return_overflowing_tokens=False,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+
+ untruncated_ids = components.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = components.tokenizer_2.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ prompt_embeds = components.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+ return prompt_embeds
+
+ @staticmethod
+ def _get_clip_prompt_embeds(components, prompt: str | list[str], device: torch.device):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ if isinstance(components, TextualInversionLoaderMixin):
+ prompt = components.maybe_convert_prompt(prompt, components.tokenizer)
+
+ text_inputs = components.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=components.tokenizer.model_max_length,
+ truncation=True,
+ return_overflowing_tokens=False,
+ return_length=False,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids
+ tokenizer_max_length = components.tokenizer.model_max_length
+ untruncated_ids = components.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = components.tokenizer.batch_decode(untruncated_ids[:, tokenizer_max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {tokenizer_max_length} tokens: {removed_text}"
+ )
+ prompt_embeds = components.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+ # Use pooled output of CLIPTextModel
+ prompt_embeds = prompt_embeds.pooler_output
+ prompt_embeds = prompt_embeds.to(dtype=components.text_encoder.dtype, device=device)
+
+ return prompt_embeds
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str | list[str],
+ prompt_2: str | list[str],
+ device: torch.device | None = None,
+ prompt_embeds: torch.FloatTensor | None = None,
+ pooled_prompt_embeds: torch.FloatTensor | None = None,
+ max_sequence_length: int = 512,
+ lora_scale: float | None = None,
+ ):
+ device = device or components._execution_device
+
+ # set lora scale so that monkey patched LoRA
+ # function of text encoder can correctly access it
+ if lora_scale is not None and isinstance(components, FluxLoraLoaderMixin):
+ components._lora_scale = lora_scale
+
+ # dynamically adjust the LoRA scale
+ if components.text_encoder is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(components.text_encoder, lora_scale)
+ if components.text_encoder_2 is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(components.text_encoder_2, lora_scale)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ if prompt_embeds is None:
+ prompt_2 = prompt_2 or prompt
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+ # We only use the pooled prompt output from the CLIPTextModel
+ pooled_prompt_embeds = FluxTextEncoderStep._get_clip_prompt_embeds(
+ components,
+ prompt=prompt,
+ device=device,
+ )
+ prompt_embeds = FluxTextEncoderStep._get_t5_prompt_embeds(
+ components,
+ prompt=prompt_2,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ if components.text_encoder is not None:
+ if isinstance(components, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(components.text_encoder, lora_scale)
+
+ if components.text_encoder_2 is not None:
+ if isinstance(components, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(components.text_encoder_2, lora_scale)
+
+ return prompt_embeds, pooled_prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ # Get inputs and intermediates
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ # Encode input prompt
+ block_state.text_encoder_lora_scale = (
+ block_state.joint_attention_kwargs.get("scale", None)
+ if block_state.joint_attention_kwargs is not None
+ else None
+ )
+ block_state.prompt_embeds, block_state.pooled_prompt_embeds = self.encode_prompt(
+ components,
+ prompt=block_state.prompt,
+ prompt_2=None,
+ prompt_embeds=None,
+ pooled_prompt_embeds=None,
+ device=block_state.device,
+ max_sequence_length=block_state.max_sequence_length,
+ lora_scale=block_state.text_encoder_lora_scale,
+ )
+
+ # Add outputs
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/inputs.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/inputs.py
new file mode 100644
index 0000000000000000000000000000000000000000..40688663278ae367592dab1464efd814c6c078fd
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/inputs.py
@@ -0,0 +1,363 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...pipelines import FluxPipeline
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import InputParam, OutputParam
+
+# TODO: consider making these common utilities for modular if they are not pipeline-specific.
+from ..qwenimage.inputs import calculate_dimension_from_latents, repeat_tensor_to_batch_size
+from .modular_pipeline import FluxModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class FluxTextInputStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text input processing step that standardizes text embeddings for the pipeline.\n"
+ "This step:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Pre-generated pooled text embeddings. Can be generated from text_encoder step.",
+ ),
+ # TODO: support negative embeddings?
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `prompt_embeds`)",
+ ),
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="pooled text embeddings used to guide the image generation",
+ ),
+ # TODO: support negative embeddings?
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.prompt_embeds is not None and block_state.pooled_prompt_embeds is not None:
+ if block_state.prompt_embeds.shape[0] != block_state.pooled_prompt_embeds.shape[0]:
+ raise ValueError(
+ "`prompt_embeds` and `pooled_prompt_embeds` must have the same batch size when passed directly, but"
+ f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `pooled_prompt_embeds`"
+ f" {block_state.pooled_prompt_embeds.shape}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ # TODO: consider adding negative embeddings?
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+ pooled_prompt_embeds = block_state.pooled_prompt_embeds.repeat(1, block_state.num_images_per_prompt)
+ block_state.pooled_prompt_embeds = pooled_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, -1
+ )
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# Adapted from `QwenImageAdditionalInputsStep`
+class FluxAdditionalInputsStep(ModularPipelineBlocks):
+ model_name = "flux"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[str] = ["image_latents"],
+ additional_batch_inputs: list[str] = [],
+ ):
+ if not isinstance(image_latent_inputs, list):
+ image_latent_inputs = [image_latent_inputs]
+ if not isinstance(additional_batch_inputs, list):
+ additional_batch_inputs = [additional_batch_inputs]
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ # Functionality section
+ summary_section = (
+ "Input processing step that:\n"
+ " 1. For image latent inputs: Updates height/width if None, patchifies latents, and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ # Inputs info
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {self._image_latent_inputs}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {self._additional_batch_inputs}"
+
+ # Placement guidance
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(name="num_images_per_prompt", default=1),
+ InputParam(name="batch_size", required=True),
+ InputParam(name="height"),
+ InputParam(name="width"),
+ ]
+
+ # Add image latent inputs
+ for image_latent_input_name in self._image_latent_inputs:
+ inputs.append(InputParam(name=image_latent_input_name))
+
+ # Add additional batch inputs
+ for input_name in self._additional_batch_inputs:
+ inputs.append(InputParam(name=input_name))
+
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="image_height", type_hint=int, description="The height of the image latents"),
+ OutputParam(name="image_width", type_hint=int, description="The width of the image latents"),
+ ]
+
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs (height/width calculation, patchify, and batch expansion)
+ for image_latent_input_name in self._image_latent_inputs:
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate height/width from latents
+ height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor)
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ if not hasattr(block_state, "image_height"):
+ block_state.image_height = height
+ if not hasattr(block_state, "image_width"):
+ block_state.image_width = width
+
+ # 2. Patchify the image latent tensor
+ # TODO: Implement patchifier for Flux.
+ latent_height, latent_width = image_latent_tensor.shape[2:]
+ image_latent_tensor = FluxPipeline._pack_latents(
+ image_latent_tensor, block_state.batch_size, image_latent_tensor.shape[1], latent_height, latent_width
+ )
+
+ # 3. Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=image_latent_input_name,
+ input_tensor=image_latent_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_name in self._additional_batch_inputs:
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ # Only expand batch size
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxKontextAdditionalInputsStep(FluxAdditionalInputsStep):
+ model_name = "flux-kontext"
+
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs (height/width calculation, patchify, and batch expansion)
+ for image_latent_input_name in self._image_latent_inputs:
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate height/width from latents
+ # Unlike the `FluxAdditionalInputsStep`, we don't overwrite the `block.height` and `block.width`
+ height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor)
+ if not hasattr(block_state, "image_height"):
+ block_state.image_height = height
+ if not hasattr(block_state, "image_width"):
+ block_state.image_width = width
+
+ # 2. Patchify the image latent tensor
+ # TODO: Implement patchifier for Flux.
+ latent_height, latent_width = image_latent_tensor.shape[2:]
+ image_latent_tensor = FluxPipeline._pack_latents(
+ image_latent_tensor, block_state.batch_size, image_latent_tensor.shape[1], latent_height, latent_width
+ )
+
+ # 3. Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=image_latent_input_name,
+ input_tensor=image_latent_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_name in self._additional_batch_inputs:
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ # Only expand batch size
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class FluxKontextSetResolutionStep(ModularPipelineBlocks):
+ model_name = "flux-kontext"
+
+ @property
+ def description(self):
+ return (
+ "Determines the height and width to be used during the subsequent computations.\n"
+ "It should always be placed _before_ the latent preparation step."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(name="height"),
+ InputParam(name="width"),
+ InputParam(name="max_area", type_hint=int, default=1024**2),
+ ]
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="height", type_hint=int, description="The height of the initial noisy latents"),
+ OutputParam(name="width", type_hint=int, description="The width of the initial noisy latents"),
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ def __call__(self, components: FluxModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ self.check_inputs(height, width, components.vae_scale_factor)
+
+ original_height, original_width = height, width
+ max_area = block_state.max_area
+ aspect_ratio = width / height
+ width = round((max_area * aspect_ratio) ** 0.5)
+ height = round((max_area / aspect_ratio) ** 0.5)
+
+ multiple_of = components.vae_scale_factor * 2
+ width = width // multiple_of * multiple_of
+ height = height // multiple_of * multiple_of
+
+ if height != original_height or width != original_width:
+ logger.warning(
+ f"Generation `height` and `width` have been adjusted to {height} and {width} to fit the model requirements."
+ )
+
+ block_state.height = height
+ block_state.width = width
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py
new file mode 100644
index 0000000000000000000000000000000000000000..51b13ef47b5211e661ba4a5c95f4e88150fa8b92
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py
@@ -0,0 +1,586 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ FluxImg2ImgPrepareLatentsStep,
+ FluxImg2ImgSetTimestepsStep,
+ FluxPrepareLatentsStep,
+ FluxRoPEInputsStep,
+ FluxSetTimestepsStep,
+)
+from .decoders import FluxDecodeStep
+from .denoise import FluxDenoiseStep
+from .encoders import (
+ FluxProcessImagesInputStep,
+ FluxTextEncoderStep,
+ FluxVaeEncoderStep,
+)
+from .inputs import (
+ FluxAdditionalInputsStep,
+ FluxTextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# vae encoder (run before before_denoise)
+
+
+# auto_docstring
+class FluxImg2ImgVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that preprocess andencode the image inputs into their latent representations.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux"
+
+ block_classes = [FluxProcessImagesInputStep(), FluxVaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
+
+
+# auto_docstring
+class FluxAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+ This is an auto pipeline block that works for img2img tasks.
+ - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided. - if `image` is not provided,
+ step will be skipped.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux"
+ block_classes = [FluxImg2ImgVaeEncoderStep]
+ block_names = ["img2img"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block that works for img2img tasks.\n"
+ + " - `FluxImg2ImgVaeEncoderStep` (img2img) is used when only `image` is provided."
+ + " - if `image` is not provided, step will be skipped."
+ )
+
+
+# before_denoise: text2img
+# auto_docstring
+class FluxBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepares the inputs for the denoise step in text-to-image generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepares the inputs for the denoise step in text-to-image generation."
+
+
+# before_denoise: img2img
+# auto_docstring
+class FluxImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step for img2img task.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The image latents to use for the denoising process. Can be generated in vae encoder and packed in input
+ step.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [
+ FluxPrepareLatentsStep(),
+ FluxImg2ImgSetTimestepsStep(),
+ FluxImg2ImgPrepareLatentsStep(),
+ FluxRoPEInputsStep(),
+ ]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_img2img_latents", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs for the denoise step for img2img task."
+
+
+# before_denoise: all task (text2img, img2img)
+# auto_docstring
+class FluxAutoBeforeDenoiseStep(AutoPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step.
+ This is an auto pipeline block that works for text2image.
+ - `FluxBeforeDenoiseStep` (text2image) is used.
+ - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`):
+ TODO: Add description.
+ width (`int`):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ The image latents to use for the denoising process. Can be generated in vae encoder and packed in input
+ step.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxImg2ImgBeforeDenoiseStep, FluxBeforeDenoiseStep]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is an auto pipeline block that works for text2image.\n"
+ + " - `FluxBeforeDenoiseStep` (text2image) is used.\n"
+ + " - `FluxImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n"
+ )
+
+
+# inputs: text2image/img2img
+
+
+# auto_docstring
+class FluxImg2ImgInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the img2img denoising step. It:
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux"
+ block_classes = [FluxTextInputStep(), FluxAdditionalInputsStep()]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return "Input step that prepares the inputs for the img2img denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+
+
+# auto_docstring
+class FluxAutoInputStep(AutoPipelineBlocks):
+ """
+ Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size,
+ and patchified.
+ This is an auto pipeline block that works for text2image/img2img tasks.
+ - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.
+ - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux"
+
+ block_classes = [FluxImg2ImgInputStep, FluxTextInputStep]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
+ " This is an auto pipeline block that works for text2image/img2img tasks.\n"
+ + " - `FluxImg2ImgInputStep` (img2img) is used when `image_latents` is provided.\n"
+ + " - `FluxTextInputStep` (text2image) is used when `image_latents` are not provided.\n"
+ )
+
+
+# auto_docstring
+class FluxCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core step that performs the denoising process for Flux.
+ This step supports text-to-image and image-to-image tasks for Flux:
+ - for image-to-image generation, you need to provide `image_latents`
+ - for text-to-image generation, all you need to provide is prompt embeddings.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux"
+ block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxDenoiseStep]
+ block_names = ["input", "before_denoise", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process for Flux.\n"
+ + "This step supports text-to-image and image-to-image tasks for Flux:\n"
+ + " - for image-to-image generation, you need to provide `image_latents`\n"
+ + " - for text-to-image generation, all you need to provide is prompt embeddings."
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Auto blocks (text2image and img2img)
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", FluxTextEncoderStep()),
+ ("vae_encoder", FluxAutoVaeEncoderStep()),
+ ("denoise", FluxCoreDenoiseStep()),
+ ("decode", FluxDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class FluxAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-to-image using Flux.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2
+ (`T5Tokenizer`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ resized_image (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image and image-to-image using Flux."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py
new file mode 100644
index 0000000000000000000000000000000000000000..3de6d282a1357fe68af7c892a44a0c8dee3b372f
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py
@@ -0,0 +1,585 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ FluxKontextRoPEInputsStep,
+ FluxPrepareLatentsStep,
+ FluxRoPEInputsStep,
+ FluxSetTimestepsStep,
+)
+from .decoders import FluxDecodeStep
+from .denoise import FluxKontextDenoiseStep
+from .encoders import (
+ FluxKontextProcessImagesInputStep,
+ FluxTextEncoderStep,
+ FluxVaeEncoderStep,
+)
+from .inputs import (
+ FluxKontextAdditionalInputsStep,
+ FluxKontextSetResolutionStep,
+ FluxTextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Flux Kontext vae encoder (run before before_denoise)
+# auto_docstring
+class FluxKontextVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that preprocess andencode the image inputs into their latent representations.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextProcessImagesInputStep(), FluxVaeEncoderStep(sample_mode="argmax")]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
+
+
+# auto_docstring
+class FluxKontextAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+ This is an auto pipeline block that works for image-conditioned tasks.
+ - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided. - if `image` is not
+ provided, step will be skipped.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ processed_image (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextVaeEncoderStep]
+ block_names = ["image_conditioned"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block that works for image-conditioned tasks.\n"
+ + " - `FluxKontextVaeEncoderStep` (image_conditioned) is used when only `image` is provided."
+ + " - if `image` is not provided, step will be skipped."
+ )
+
+
+# before_denoise: text2img
+# auto_docstring
+class FluxKontextBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepares the inputs for the denoise step for Flux Kontext
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepares the inputs for the denoise step for Flux Kontext\n"
+ "for text-to-image tasks."
+
+
+# before_denoise: image-conditioned
+# auto_docstring
+class FluxKontextImageConditionedBeforeDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step for Flux Kontext
+ for image-conditioned tasks.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_height (`None`, *optional*):
+ TODO: Add description.
+ image_width (`None`, *optional*):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxPrepareLatentsStep(), FluxSetTimestepsStep(), FluxKontextRoPEInputsStep()]
+ block_names = ["prepare_latents", "set_timesteps", "prepare_rope_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step for Flux Kontext\n"
+ "for image-conditioned tasks."
+ )
+
+
+# auto_docstring
+class FluxKontextAutoBeforeDenoiseStep(AutoPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs for the denoise step.
+ This is an auto pipeline block that works for text2image.
+ - `FluxKontextBeforeDenoiseStep` (text2image) is used.
+ - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is
+ provided.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.
+ Can be generated in input step.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ image_height (`None`, *optional*):
+ TODO: Add description.
+ image_width (`None`, *optional*):
+ TODO: Add description.
+ prompt_embeds (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ timesteps (`Tensor`):
+ The timesteps to use for inference
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time
+ guidance (`Tensor`):
+ Optional guidance to be used.
+ txt_ids (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation.
+ img_ids (`list`):
+ The sequence lengths of the image latents, used for RoPE calculation.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = [FluxKontextImageConditionedBeforeDenoiseStep, FluxKontextBeforeDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is an auto pipeline block that works for text2image.\n"
+ + " - `FluxKontextBeforeDenoiseStep` (text2image) is used.\n"
+ + " - `FluxKontextImageConditionedBeforeDenoiseStep` (image_conditioned) is used when only `image_latents` is provided.\n"
+ )
+
+
+# inputs: Flux Kontext
+# auto_docstring
+class FluxKontextInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the both text2img and img2img denoising step. It:
+ - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).
+ - update height/width based `image_latents`, patchify `image_latents`.
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ height (`int`):
+ The height of the initial noisy latents
+ width (`int`):
+ The width of the initial noisy latents
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextSetResolutionStep(), FluxTextInputStep(), FluxKontextAdditionalInputsStep()]
+ block_names = ["set_resolution", "text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the both text2img and img2img denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+ )
+
+
+# auto_docstring
+class FluxKontextAutoInputStep(AutoPipelineBlocks):
+ """
+ Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size,
+ and patchified.
+ This is an auto pipeline block that works for text2image/img2img tasks.
+ - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided.
+ - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present.
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ height (`int`):
+ The height of the initial noisy latents
+ width (`int`):
+ The width of the initial noisy latents
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `prompt_embeds`)
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation
+ pooled_prompt_embeds (`Tensor`):
+ pooled text embeddings used to guide the image generation
+ image_height (`int`):
+ The height of the image latents
+ image_width (`int`):
+ The width of the image latents
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextInputStep, FluxTextInputStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Input step that standardize the inputs for the denoising step, e.g. make sure inputs have consistent batch size, and patchified. \n"
+ " This is an auto pipeline block that works for text2image/img2img tasks.\n"
+ + " - `FluxKontextInputStep` (image_conditioned) is used when `image_latents` is provided.\n"
+ + " - `FluxKontextInputStep` is also capable of handling text2image task when `image_latent` isn't present."
+ )
+
+
+# auto_docstring
+class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core step that performs the denoising process for Flux Kontext.
+ This step supports text-to-image and image-conditioned tasks for Flux Kontext:
+ - for image-conditioned generation, you need to provide `image_latents`
+ - for text-to-image generation, all you need to provide is prompt embeddings.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be generated from text_encoder step.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux-kontext"
+ block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextDenoiseStep]
+ block_names = ["input", "before_denoise", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process for Flux Kontext.\n"
+ + "This step supports text-to-image and image-conditioned tasks for Flux Kontext:\n"
+ + " - for image-conditioned generation, you need to provide `image_latents`\n"
+ + " - for text-to-image generation, all you need to provide is prompt embeddings."
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+AUTO_BLOCKS_KONTEXT = InsertableDict(
+ [
+ ("text_encoder", FluxTextEncoderStep()),
+ ("vae_encoder", FluxKontextAutoVaeEncoderStep()),
+ ("denoise", FluxKontextCoreDenoiseStep()),
+ ("decode", FluxDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class FluxKontextAutoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for image-to-image using Flux Kontext.
+
+ Supported workflows:
+ - `image_conditioned`: requires `image`, `prompt`
+ - `text2image`: requires `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) tokenizer (`CLIPTokenizer`) text_encoder_2 (`T5EncoderModel`) tokenizer_2
+ (`T5Tokenizer`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`FluxTransformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ _auto_resize (`bool`, *optional*, defaults to True):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ max_area (`int`, *optional*, defaults to 1048576):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 3.5):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux-kontext"
+
+ block_classes = AUTO_BLOCKS_KONTEXT.values()
+ block_names = AUTO_BLOCKS_KONTEXT.keys()
+ _workflow_map = {
+ "image_conditioned": {"image": True, "prompt": True},
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Modular pipeline for image-to-image using Flux Kontext."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b03721aed61aab6035f110bc1fc873f868a1401
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux/modular_pipeline.py
@@ -0,0 +1,67 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class FluxModularPipeline(ModularPipeline, FluxLoraLoaderMixin, TextualInversionLoaderMixin):
+ """
+ A ModularPipeline for Flux.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "FluxAutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ return 128
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if getattr(self, "vae", None) is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if getattr(self, "transformer", None):
+ num_channels_latents = self.transformer.config.in_channels // 4
+ return num_channels_latents
+
+
+class FluxKontextModularPipeline(FluxModularPipeline):
+ """
+ A ModularPipeline for Flux Kontext.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "FluxKontextAutoBlocks"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c88abb345551ce40144bc5e691d6dc7cd1cc614
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/__init__.py
@@ -0,0 +1,57 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["encoders"] = ["Flux2RemoteTextEncoderStep"]
+ _import_structure["modular_blocks_flux2"] = ["Flux2AutoBlocks"]
+ _import_structure["modular_blocks_flux2_klein"] = ["Flux2KleinAutoBlocks"]
+ _import_structure["modular_blocks_flux2_klein_base"] = ["Flux2KleinBaseAutoBlocks"]
+ _import_structure["modular_pipeline"] = [
+ "Flux2KleinBaseModularPipeline",
+ "Flux2KleinModularPipeline",
+ "Flux2ModularPipeline",
+ ]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .encoders import Flux2RemoteTextEncoderStep
+ from .modular_blocks_flux2 import Flux2AutoBlocks
+ from .modular_blocks_flux2_klein import Flux2KleinAutoBlocks
+ from .modular_blocks_flux2_klein_base import Flux2KleinBaseAutoBlocks
+ from .modular_pipeline import Flux2KleinBaseModularPipeline, Flux2KleinModularPipeline, Flux2ModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..23f6cac0670e41421e21cafe55c38b41813b52d7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/before_denoise.py
@@ -0,0 +1,591 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...models import Flux2Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Flux2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
+ """Compute empirical mu for Flux2 timestep scheduling."""
+ a1, b1 = 8.73809524e-05, 1.89833333
+ a2, b2 = 0.00016927, 0.45666666
+
+ if image_seq_len > 4300:
+ mu = a2 * image_seq_len + b2
+ return float(mu)
+
+ m_200 = a2 * image_seq_len + b2
+ m_10 = a1 * image_seq_len + b1
+
+ a = (m_200 - m_10) / 190.0
+ b = m_200 - 200.0 * a
+ mu = a * num_steps + b
+
+ return float(mu)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class Flux2SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", Flux2Transformer2DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for Flux2 inference using empirical mu calculation"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ InputParam("latents", type_hint=torch.Tensor),
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ scheduler = components.scheduler
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ vae_scale_factor = components.vae_scale_factor
+
+ latent_height = 2 * (int(height) // (vae_scale_factor * 2))
+ latent_width = 2 * (int(width) // (vae_scale_factor * 2))
+ image_seq_len = (latent_height // 2) * (latent_width // 2)
+
+ num_inference_steps = block_state.num_inference_steps
+ sigmas = block_state.sigmas
+ timesteps = block_state.timesteps
+
+ if timesteps is None and sigmas is None:
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+ if hasattr(scheduler.config, "use_flow_sigmas") and scheduler.config.use_flow_sigmas:
+ sigmas = None
+
+ mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps)
+
+ timesteps, num_inference_steps = retrieve_timesteps(
+ scheduler,
+ num_inference_steps,
+ device,
+ timesteps=timesteps,
+ sigmas=sigmas,
+ mu=mu,
+ )
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = num_inference_steps
+
+ components.scheduler.set_begin_index(0)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2PrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the initial noise latents for Flux2 text-to-image generation"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam("latents", type_hint=torch.Tensor | None),
+ InputParam("num_images_per_prompt", type_hint=int, default=1),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ ),
+ OutputParam("latent_ids", type_hint=torch.Tensor, description="Position IDs for the latents (for RoPE)"),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ vae_scale_factor = components.vae_scale_factor
+ if (block_state.height is not None and block_state.height % (vae_scale_factor * 2) != 0) or (
+ block_state.width is not None and block_state.width % (vae_scale_factor * 2) != 0
+ ):
+ logger.warning(
+ f"`height` and `width` have to be divisible by {vae_scale_factor * 2} but are {block_state.height} and {block_state.width}."
+ )
+
+ @staticmethod
+ def _prepare_latent_ids(latents: torch.Tensor):
+ """
+ Generates 4D position coordinates (T, H, W, L) for latent tensors.
+
+ Args:
+ latents: Latent tensor of shape (B, C, H, W)
+
+ Returns:
+ Position IDs tensor of shape (B, H*W, 4)
+ """
+ batch_size, _, height, width = latents.shape
+
+ t = torch.arange(1)
+ h = torch.arange(height)
+ w = torch.arange(width)
+ l = torch.arange(1)
+
+ latent_ids = torch.cartesian_prod(t, h, w, l)
+ latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+ return latent_ids
+
+ @staticmethod
+ def _pack_latents(latents):
+ """Pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)"""
+ batch_size, num_channels, height, width = latents.shape
+ latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1)
+ return latents
+
+ @staticmethod
+ def prepare_latents(
+ comp,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ height = 2 * (int(height) // (comp.vae_scale_factor * 2))
+ width = 2 * (int(width) // (comp.vae_scale_factor * 2))
+
+ shape = (batch_size, num_channels_latents * 4, height // 2, width // 2)
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ latents = latents.to(device=device, dtype=dtype)
+
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+ block_state.device = components._execution_device
+ block_state.num_channels_latents = components.num_channels_latents
+
+ self.check_inputs(components, block_state)
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+ latents = self.prepare_latents(
+ components,
+ batch_size,
+ block_state.num_channels_latents,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ block_state.latents,
+ )
+
+ latent_ids = self._prepare_latent_ids(latents)
+ latent_ids = latent_ids.to(block_state.device)
+
+ latents = self._pack_latents(latents)
+
+ block_state.latents = latents
+ block_state.latent_ids = latent_ids
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2RoPEInputsStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the 4D RoPE position IDs for Flux2 denoising. Should be placed after text encoder and latent preparation steps."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="prompt_embeds", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="txt_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.",
+ ),
+ ]
+
+ @staticmethod
+ def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
+ """Prepare 4D position IDs for text tokens."""
+ B, L, _ = x.shape
+ out_ids = []
+
+ for i in range(B):
+ t = torch.arange(1) if t_coord is None else t_coord[i]
+ h = torch.arange(1)
+ w = torch.arange(1)
+ seq_l = torch.arange(L)
+
+ coords = torch.cartesian_prod(t, h, w, seq_l)
+ out_ids.append(coords)
+
+ return torch.stack(out_ids)
+
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_embeds = block_state.prompt_embeds
+ device = prompt_embeds.device
+
+ block_state.txt_ids = self._prepare_text_ids(prompt_embeds)
+ block_state.txt_ids = block_state.txt_ids.to(device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2KleinBaseRoPEInputsStep(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the 4D RoPE position IDs for Flux2-Klein base model denoising. Should be placed after text encoder and latent preparation steps."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="prompt_embeds", required=True),
+ InputParam(name="negative_prompt_embeds", required=False),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="txt_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.",
+ ),
+ OutputParam(
+ name="negative_txt_ids",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="4D position IDs (T, H, W, L) for negative text tokens, used for RoPE calculation.",
+ ),
+ ]
+
+ @staticmethod
+ def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
+ """Prepare 4D position IDs for text tokens."""
+ B, L, _ = x.shape
+ out_ids = []
+
+ for i in range(B):
+ t = torch.arange(1) if t_coord is None else t_coord[i]
+ h = torch.arange(1)
+ w = torch.arange(1)
+ seq_l = torch.arange(L)
+
+ coords = torch.cartesian_prod(t, h, w, seq_l)
+ out_ids.append(coords)
+
+ return torch.stack(out_ids)
+
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_embeds = block_state.prompt_embeds
+ device = prompt_embeds.device
+
+ block_state.txt_ids = self._prepare_text_ids(prompt_embeds)
+ block_state.txt_ids = block_state.txt_ids.to(device)
+
+ block_state.negative_txt_ids = None
+ if block_state.negative_prompt_embeds is not None:
+ block_state.negative_txt_ids = self._prepare_text_ids(block_state.negative_prompt_embeds)
+ block_state.negative_txt_ids = block_state.negative_txt_ids.to(device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2PrepareImageLatentsStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares image latents and their position IDs for Flux2 image conditioning."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image_latents", type_hint=list[torch.Tensor]),
+ InputParam("batch_size", required=True, type_hint=int),
+ InputParam("num_images_per_prompt", default=1, type_hint=int),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Packed image latents for conditioning",
+ ),
+ OutputParam(
+ "image_latent_ids",
+ type_hint=torch.Tensor,
+ description="Position IDs for image latents",
+ ),
+ ]
+
+ @staticmethod
+ def _prepare_image_ids(image_latents: list[torch.Tensor], scale: int = 10):
+ """
+ Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents.
+
+ Args:
+ image_latents: A list of image latent feature tensors of shape (1, C, H, W).
+ scale: Factor used to define the time separation between latents.
+
+ Returns:
+ Combined coordinate tensor of shape (1, N_total, 4)
+ """
+ if not isinstance(image_latents, list):
+ raise ValueError(f"Expected `image_latents` to be a list, got {type(image_latents)}.")
+
+ t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))]
+ t_coords = [t.view(-1) for t in t_coords]
+
+ image_latent_ids = []
+ for x, t in zip(image_latents, t_coords):
+ x = x.squeeze(0)
+ _, height, width = x.shape
+
+ x_ids = torch.cartesian_prod(t, torch.arange(height), torch.arange(width), torch.arange(1))
+ image_latent_ids.append(x_ids)
+
+ image_latent_ids = torch.cat(image_latent_ids, dim=0)
+ image_latent_ids = image_latent_ids.unsqueeze(0)
+
+ return image_latent_ids
+
+ @staticmethod
+ def _pack_latents(latents):
+ """Pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)"""
+ batch_size, num_channels, height, width = latents.shape
+ latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1)
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ image_latents = block_state.image_latents
+
+ if image_latents is None:
+ block_state.image_latents = None
+ block_state.image_latent_ids = None
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+ device = components._execution_device
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+ image_latent_ids = self._prepare_image_ids(image_latents)
+
+ packed_latents = []
+ for latent in image_latents:
+ packed = self._pack_latents(latent)
+ packed = packed.squeeze(0)
+ packed_latents.append(packed)
+
+ image_latents = torch.cat(packed_latents, dim=0)
+ image_latents = image_latents.unsqueeze(0)
+
+ image_latents = image_latents.repeat(batch_size, 1, 1)
+ image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1)
+ image_latent_ids = image_latent_ids.to(device)
+
+ block_state.image_latents = image_latents
+ block_state.image_latent_ids = image_latent_ids
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2PrepareGuidanceStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the guidance scale tensor for Flux2 inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("guidance_scale", default=4.0),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("guidance", type_hint=torch.Tensor, description="Guidance scale tensor"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+ guidance = torch.full([1], block_state.guidance_scale, device=device, dtype=torch.float32)
+ guidance = guidance.expand(batch_size)
+ block_state.guidance = guidance
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d0d89306e09b7522f97bb0b562bf355a6fc1541
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/decoders.py
@@ -0,0 +1,185 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, Union
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLFlux2
+from ...pipelines.flux2.image_processor import Flux2ImageProcessor
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class Flux2UnpackLatentsStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "Step that unpacks the latents from the denoising step"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents from the denoising step",
+ ),
+ InputParam(
+ "latent_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Position IDs for the latents, used for unpacking",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The denoise latents from denoising step, unpacked with position IDs.",
+ )
+ ]
+
+ @staticmethod
+ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> torch.Tensor:
+ """
+ Unpack latents using position IDs to scatter tokens into place.
+
+ Args:
+ x: Packed latents tensor of shape (B, seq_len, C)
+ x_ids: Position IDs tensor of shape (B, seq_len, 4) with (T, H, W, L) coordinates
+
+ Returns:
+ Unpacked latents tensor of shape (B, C, H, W)
+ """
+ x_list = []
+ for data, pos in zip(x, x_ids):
+ _, ch = data.shape # noqa: F841
+ h_ids = pos[:, 1].to(torch.int64)
+ w_ids = pos[:, 2].to(torch.int64)
+
+ h = torch.max(h_ids) + 1
+ w = torch.max(w_ids) + 1
+
+ flat_ids = h_ids * w + w_ids
+
+ out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype)
+ out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data)
+
+ out = out.view(h, w, ch).permute(2, 0, 1)
+ x_list.append(out)
+
+ return torch.stack(x_list, dim=0)
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latents = block_state.latents
+ latent_ids = block_state.latent_ids
+
+ latents = self._unpack_latents_with_ids(latents, latent_ids)
+
+ block_state.latents = latents
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2DecodeStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLFlux2),
+ ComponentSpec(
+ "image_processor",
+ Flux2ImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images using Flux2 VAE with batch norm denormalization"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("output_type", default="pil"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents from the denoising step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "images",
+ type_hint=Union[list[PIL.Image.Image], torch.Tensor, np.ndarray],
+ description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array",
+ )
+ ]
+
+ @staticmethod
+ def _unpatchify_latents(latents):
+ """Convert patchified latents back to regular format."""
+ batch_size, num_channels_latents, height, width = latents.shape
+ latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
+ latents = latents.permute(0, 1, 4, 2, 5, 3)
+ latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae = components.vae
+
+ latents = block_state.latents
+
+ latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+ latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
+ latents.device, latents.dtype
+ )
+ latents = latents * latents_bn_std + latents_bn_mean
+
+ latents = self._unpatchify_latents(latents)
+
+ block_state.images = vae.decode(latents, return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(
+ block_state.images, output_type=block_state.output_type
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..53802600bf6616605860afb8dd6dfbe84dbaf0eb
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/denoise.py
@@ -0,0 +1,501 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import Flux2Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import is_torch_xla_available, logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline
+
+
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm
+
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class Flux2LoopDenoiser(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Flux2Transformer2DModel)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoises the latents for Flux2. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `Flux2DenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("joint_attention_kwargs"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to denoise. Shape: (B, seq_len, C)",
+ ),
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)",
+ ),
+ InputParam(
+ "image_latent_ids",
+ type_hint=torch.Tensor,
+ description="Position IDs for image latents. Shape: (B, img_seq_len, 4)",
+ ),
+ InputParam(
+ "guidance",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Guidance scale as a tensor",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Text embeddings from Mistral3",
+ ),
+ InputParam(
+ "txt_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for text tokens (T, H, W, L)",
+ ),
+ InputParam(
+ "latent_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for latent tokens (T, H, W, L)",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ latents = block_state.latents
+ latent_model_input = latents.to(components.transformer.dtype)
+ img_ids = block_state.latent_ids
+
+ image_latents = getattr(block_state, "image_latents", None)
+ if image_latents is not None:
+ latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype)
+ image_latent_ids = block_state.image_latent_ids
+ img_ids = torch.cat([img_ids, image_latent_ids], dim=1)
+
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=block_state.guidance,
+ encoder_hidden_states=block_state.prompt_embeds,
+ txt_ids=block_state.txt_ids,
+ img_ids=img_ids,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ noise_pred = noise_pred[:, : latents.size(1)]
+ block_state.noise_pred = noise_pred
+
+ return components, block_state
+
+
+# same as Flux2LoopDenoiser but guidance=None
+class Flux2KleinLoopDenoiser(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Flux2Transformer2DModel)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoises the latents for Flux2. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `Flux2DenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("joint_attention_kwargs"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to denoise. Shape: (B, seq_len, C)",
+ ),
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)",
+ ),
+ InputParam(
+ "image_latent_ids",
+ type_hint=torch.Tensor,
+ description="Position IDs for image latents. Shape: (B, img_seq_len, 4)",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Text embeddings from Qwen3",
+ ),
+ InputParam(
+ "txt_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for text tokens (T, H, W, L)",
+ ),
+ InputParam(
+ "latent_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for latent tokens (T, H, W, L)",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ latents = block_state.latents
+ latent_model_input = latents.to(components.transformer.dtype)
+ img_ids = block_state.latent_ids
+
+ image_latents = getattr(block_state, "image_latents", None)
+ if image_latents is not None:
+ latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype)
+ image_latent_ids = block_state.image_latent_ids
+ img_ids = torch.cat([img_ids, image_latent_ids], dim=1)
+
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=None,
+ encoder_hidden_states=block_state.prompt_embeds,
+ txt_ids=block_state.txt_ids,
+ img_ids=img_ids,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ noise_pred = noise_pred[:, : latents.size(1)]
+ block_state.noise_pred = noise_pred
+
+ return components, block_state
+
+
+# support CFG for Flux2-Klein base model
+class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", Flux2Transformer2DModel),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="is_distilled", default=False),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoises the latents for Flux2. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `Flux2DenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("joint_attention_kwargs"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to denoise. Shape: (B, seq_len, C)",
+ ),
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)",
+ ),
+ InputParam(
+ "image_latent_ids",
+ type_hint=torch.Tensor,
+ description="Position IDs for image latents. Shape: (B, img_seq_len, 4)",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Text embeddings from Qwen3",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ required=False,
+ type_hint=torch.Tensor,
+ description="Negative text embeddings from Qwen3",
+ ),
+ InputParam(
+ "txt_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for text tokens (T, H, W, L)",
+ ),
+ InputParam(
+ "negative_txt_ids",
+ required=False,
+ type_hint=torch.Tensor,
+ description="4D position IDs for negative text tokens (T, H, W, L)",
+ ),
+ InputParam(
+ "latent_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="4D position IDs for latent tokens (T, H, W, L)",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ latents = block_state.latents
+ latent_model_input = latents.to(components.transformer.dtype)
+ img_ids = block_state.latent_ids
+
+ image_latents = getattr(block_state, "image_latents", None)
+ if image_latents is not None:
+ latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype)
+ image_latent_ids = block_state.image_latent_ids
+ img_ids = torch.cat([img_ids, image_latent_ids], dim=1)
+
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ guider_inputs = {
+ "encoder_hidden_states": (
+ getattr(block_state, "prompt_embeds", None),
+ getattr(block_state, "negative_prompt_embeds", None),
+ ),
+ "txt_ids": (
+ getattr(block_state, "txt_ids", None),
+ getattr(block_state, "negative_txt_ids", None),
+ ),
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+ noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep / 1000,
+ guidance=None,
+ img_ids=img_ids,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ guider_state_batch.noise_pred = noise_pred[:, : latents.size(1)]
+ components.guider.cleanup_models(components.transformer)
+
+ # perform guidance
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class Flux2LoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that updates the latents after denoising. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `Flux2DenoiseLoopWrapper`)"
+ )
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class Flux2DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoises the latents over `timesteps`. "
+ "The specific steps within each iteration can be customized with `sub_blocks` attribute"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", Flux2Transformer2DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2DenoiseStep(Flux2DenoiseLoopWrapper):
+ block_classes = [Flux2LoopDenoiser, Flux2LoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents for Flux2. \n"
+ "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `Flux2LoopDenoiser`\n"
+ " - `Flux2LoopAfterDenoiser`\n"
+ "This block supports both text-to-image and image-conditioned generation."
+ )
+
+
+class Flux2KleinDenoiseStep(Flux2DenoiseLoopWrapper):
+ block_classes = [Flux2KleinLoopDenoiser, Flux2LoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents for Flux2. \n"
+ "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `Flux2KleinLoopDenoiser`\n"
+ " - `Flux2LoopAfterDenoiser`\n"
+ "This block supports both text-to-image and image-conditioned generation."
+ )
+
+
+class Flux2KleinBaseDenoiseStep(Flux2DenoiseLoopWrapper):
+ block_classes = [Flux2KleinBaseLoopDenoiser, Flux2LoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents for Flux2. \n"
+ "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `Flux2KleinBaseLoopDenoiser`\n"
+ " - `Flux2LoopAfterDenoiser`\n"
+ "This block supports both text-to-image and image-conditioned generation."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c312ec025bf532541b00b136cbf52f89dc557b6
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/encoders.py
@@ -0,0 +1,608 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from transformers import AutoProcessor, Mistral3ForConditionalGeneration, Qwen2TokenizerFast, Qwen3ForCausalLM
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import AutoencoderKLFlux2
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def format_text_input(prompts: list[str], system_message: str = None):
+ """Format prompts for Mistral3 chat template."""
+ cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts]
+
+ return [
+ [
+ {
+ "role": "system",
+ "content": [{"type": "text", "text": system_message}],
+ },
+ {"role": "user", "content": [{"type": "text", "text": prompt}]},
+ ]
+ for prompt in cleaned_txt
+ ]
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class Flux2TextEncoderStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ # fmt: off
+ DEFAULT_SYSTEM_MESSAGE = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation."
+ # fmt: on
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings using Mistral3 to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Mistral3ForConditionalGeneration),
+ ComponentSpec("tokenizer", AutoProcessor),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("max_sequence_length", type_hint=int, default=512, required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(10, 20, 30), required=False),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Text embeddings from Mistral3 used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ prompt = block_state.prompt
+ if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ @staticmethod
+ def _get_mistral_3_prompt_embeds(
+ text_encoder: Mistral3ForConditionalGeneration,
+ tokenizer: AutoProcessor,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ max_sequence_length: int = 512,
+ # fmt: off
+ system_message: str = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation.",
+ # fmt: on
+ hidden_states_layers: tuple[int] = (10, 20, 30),
+ ):
+ dtype = text_encoder.dtype if dtype is None else dtype
+ device = text_encoder.device if device is None else device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ messages_batch = format_text_input(prompts=prompt, system_message=system_message)
+
+ inputs = tokenizer.apply_chat_template(
+ messages_batch,
+ add_generation_prompt=False,
+ tokenize=True,
+ return_dict=True,
+ return_tensors="pt",
+ padding="max_length",
+ truncation=True,
+ max_length=max_sequence_length,
+ )
+
+ input_ids = inputs["input_ids"].to(device)
+ attention_mask = inputs["attention_mask"].to(device)
+
+ output = text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+
+ out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
+ out = out.to(dtype=dtype, device=device)
+
+ batch_size, num_channels, seq_len, hidden_dim = out.shape
+ prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
+
+ return prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ prompt = block_state.prompt
+ if prompt is None:
+ prompt = ""
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ block_state.prompt_embeds = self._get_mistral_3_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ device=block_state.device,
+ max_sequence_length=block_state.max_sequence_length,
+ system_message=self.DEFAULT_SYSTEM_MESSAGE,
+ hidden_states_layers=block_state.text_encoder_out_layers,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2RemoteTextEncoderStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ REMOTE_URL = "https://remote-text-encoder-flux-2.huggingface.co/predict"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings using a remote API endpoint"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Text embeddings from remote API used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ prompt = block_state.prompt
+ if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ import io
+
+ import requests
+ from huggingface_hub import get_token
+
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ prompt = block_state.prompt
+ if prompt is None:
+ prompt = ""
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ response = requests.post(
+ self.REMOTE_URL,
+ json={"prompt": prompt},
+ headers={
+ "Authorization": f"Bearer {get_token()}",
+ "Content-Type": "application/json",
+ },
+ )
+ response.raise_for_status()
+
+ block_state.prompt_embeds = torch.load(io.BytesIO(response.content), weights_only=True)
+ block_state.prompt_embeds = block_state.prompt_embeds.to(block_state.device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2KleinTextEncoderStep(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3ForCausalLM),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="is_distilled", default=True),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("max_sequence_length", type_hint=int, default=512, required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Text embeddings from qwen3 used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ prompt = block_state.prompt
+
+ if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds
+ def _get_qwen3_prompt_embeds(
+ text_encoder: Qwen3ForCausalLM,
+ tokenizer: Qwen2TokenizerFast,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ max_sequence_length: int = 512,
+ hidden_states_layers: list[int] = (9, 18, 27),
+ ):
+ dtype = text_encoder.dtype if dtype is None else dtype
+ device = text_encoder.device if device is None else device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ all_input_ids = []
+ all_attention_masks = []
+
+ for single_prompt in prompt:
+ messages = [{"role": "user", "content": single_prompt}]
+ text = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=False,
+ )
+ inputs = tokenizer(
+ text,
+ return_tensors="pt",
+ padding="max_length",
+ truncation=True,
+ max_length=max_sequence_length,
+ )
+
+ all_input_ids.append(inputs["input_ids"])
+ all_attention_masks.append(inputs["attention_mask"])
+
+ input_ids = torch.cat(all_input_ids, dim=0).to(device)
+ attention_mask = torch.cat(all_attention_masks, dim=0).to(device)
+
+ # Forward pass through the model
+ output = text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+
+ # Only use outputs from intermediate layers and stack them
+ out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
+ out = out.to(dtype=dtype, device=device)
+
+ batch_size, num_channels, seq_len, hidden_dim = out.shape
+ prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
+
+ return prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ device = components._execution_device
+
+ prompt = block_state.prompt
+ if prompt is None:
+ prompt = ""
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ block_state.prompt_embeds = self._get_qwen3_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ device=device,
+ max_sequence_length=block_state.max_sequence_length,
+ hidden_states_layers=block_state.text_encoder_out_layers,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2KleinBaseTextEncoderStep(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3ForCausalLM),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(name="is_distilled", default=False),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("max_sequence_length", type_hint=int, default=512, required=False),
+ InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Text embeddings from qwen3 used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Negative text embeddings from qwen3 used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ prompt = block_state.prompt
+
+ if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds
+ def _get_qwen3_prompt_embeds(
+ text_encoder: Qwen3ForCausalLM,
+ tokenizer: Qwen2TokenizerFast,
+ prompt: str | list[str],
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ max_sequence_length: int = 512,
+ hidden_states_layers: list[int] = (9, 18, 27),
+ ):
+ dtype = text_encoder.dtype if dtype is None else dtype
+ device = text_encoder.device if device is None else device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ all_input_ids = []
+ all_attention_masks = []
+
+ for single_prompt in prompt:
+ messages = [{"role": "user", "content": single_prompt}]
+ text = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=False,
+ )
+ inputs = tokenizer(
+ text,
+ return_tensors="pt",
+ padding="max_length",
+ truncation=True,
+ max_length=max_sequence_length,
+ )
+
+ all_input_ids.append(inputs["input_ids"])
+ all_attention_masks.append(inputs["attention_mask"])
+
+ input_ids = torch.cat(all_input_ids, dim=0).to(device)
+ attention_mask = torch.cat(all_attention_masks, dim=0).to(device)
+
+ # Forward pass through the model
+ output = text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+
+ # Only use outputs from intermediate layers and stack them
+ out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1)
+ out = out.to(dtype=dtype, device=device)
+
+ batch_size, num_channels, seq_len, hidden_dim = out.shape
+ prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
+
+ return prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ device = components._execution_device
+
+ prompt = block_state.prompt
+ if prompt is None:
+ prompt = ""
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ block_state.prompt_embeds = self._get_qwen3_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ device=device,
+ max_sequence_length=block_state.max_sequence_length,
+ hidden_states_layers=block_state.text_encoder_out_layers,
+ )
+
+ if components.requires_unconditional_embeds:
+ negative_prompt = [""] * len(prompt)
+ block_state.negative_prompt_embeds = self._get_qwen3_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=negative_prompt,
+ device=device,
+ max_sequence_length=block_state.max_sequence_length,
+ hidden_states_layers=block_state.text_encoder_out_layers,
+ )
+ else:
+ block_state.negative_prompt_embeds = None
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2VaeEncoderStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "VAE Encoder step that encodes preprocessed images into latent representations for Flux2."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKLFlux2)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("condition_images", type_hint=list[torch.Tensor]),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=list[torch.Tensor],
+ description="List of latent representations for each reference image",
+ ),
+ ]
+
+ @staticmethod
+ def _patchify_latents(latents):
+ """Convert latents to patchified format for Flux2."""
+ batch_size, num_channels_latents, height, width = latents.shape
+ latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+ latents = latents.permute(0, 1, 3, 5, 2, 4)
+ latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2)
+ return latents
+
+ def _encode_vae_image(self, vae: AutoencoderKLFlux2, image: torch.Tensor, generator: torch.Generator):
+ """Encode a single image using Flux2 VAE with batch norm normalization."""
+ if image.ndim != 4:
+ raise ValueError(f"Expected image dims 4, got {image.ndim}.")
+
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode="argmax")
+ image_latents = self._patchify_latents(image_latents)
+
+ latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+ latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
+ latents_bn_std = latents_bn_std.to(image_latents.device, image_latents.dtype)
+ image_latents = (image_latents - latents_bn_mean) / latents_bn_std
+
+ return image_latents
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ condition_images = block_state.condition_images
+
+ if condition_images is None:
+ return components, state
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ image_latents = []
+ for image in condition_images:
+ image = image.to(device=device, dtype=dtype)
+ latent = self._encode_vae_image(
+ vae=components.vae,
+ image=image,
+ generator=block_state.generator,
+ )
+ image_latents.append(latent)
+
+ block_state.image_latents = image_latents
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/inputs.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/inputs.py
new file mode 100644
index 0000000000000000000000000000000000000000..28dcf2027e465a74c524c6ae39388437281ba030
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/inputs.py
@@ -0,0 +1,242 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...pipelines.flux2.image_processor import Flux2ImageProcessor
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Flux2ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class Flux2TextInputStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `prompt_embeds`)",
+ ),
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2KleinBaseTextInputStep(ModularPipelineBlocks):
+ model_name = "flux2-klein"
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ required=False,
+ kwargs_type="denoiser_input_fields",
+ type_hint=torch.Tensor,
+ description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `prompt_embeds`)",
+ ),
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Negative text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class Flux2ProcessImagesInputStep(ModularPipelineBlocks):
+ model_name = "flux2"
+
+ @property
+ def description(self) -> str:
+ return "Image preprocess step for Flux2. Validates and preprocesses reference images."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ Flux2ImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image"),
+ InputParam("height"),
+ InputParam("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="condition_images", type_hint=list[torch.Tensor])]
+
+ @torch.no_grad()
+ def __call__(self, components: Flux2ModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+ images = block_state.image
+
+ if images is None:
+ block_state.condition_images = None
+ self.set_block_state(state, block_state)
+ return components, state
+
+ if not isinstance(images, list):
+ images = [images]
+
+ condition_images = []
+ for img in images:
+ components.image_processor.check_image_input(img)
+
+ image_width, image_height = img.size
+ if image_width * image_height > 1024 * 1024:
+ img = components.image_processor._resize_to_target_area(img, 1024 * 1024)
+ image_width, image_height = img.size
+
+ multiple_of = components.vae_scale_factor * 2
+ image_width = (image_width // multiple_of) * multiple_of
+ image_height = (image_height // multiple_of) * multiple_of
+ condition_img = components.image_processor.preprocess(
+ img, height=image_height, width=image_width, resize_mode="crop"
+ )
+ condition_images.append(condition_img)
+
+ if block_state.height is None:
+ block_state.height = image_height
+ if block_state.width is None:
+ block_state.width = image_width
+
+ block_state.condition_images = condition_images
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a3fc0fb4b90adb6397ab0bc36fca8b045ddf684
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py
@@ -0,0 +1,356 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Flux2PrepareGuidanceStep,
+ Flux2PrepareImageLatentsStep,
+ Flux2PrepareLatentsStep,
+ Flux2RoPEInputsStep,
+ Flux2SetTimestepsStep,
+)
+from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
+from .denoise import Flux2DenoiseStep
+from .encoders import (
+ Flux2TextEncoderStep,
+ Flux2VaeEncoderStep,
+)
+from .inputs import (
+ Flux2ProcessImagesInputStep,
+ Flux2TextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# auto_docstring
+class Flux2VaeEncoderSequentialStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2"
+
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning."
+
+
+# auto_docstring
+class Flux2AutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ block_classes = [Flux2VaeEncoderSequentialStep]
+ block_names = ["img_conditioning"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image inputs into their latent representations.\n"
+ "This is an auto pipeline block that works for image conditioning tasks.\n"
+ " - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+Flux2CoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2TextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_guidance", Flux2PrepareGuidanceStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2DenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-dev.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2"
+
+ block_classes = Flux2CoreDenoiseBlocks.values()
+ block_names = Flux2CoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-dev."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+Flux2ImageConditionedCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2TextInputStep()),
+ ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_guidance", Flux2PrepareGuidanceStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2DenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-dev with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2"
+
+ block_classes = Flux2ImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2ImageConditionedCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-dev with image conditioning."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+class Flux2AutoCoreDenoiseStep(AutoPipelineBlocks):
+ model_name = "flux2"
+
+ block_classes = [Flux2ImageConditionedCoreDenoiseStep, Flux2CoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoise step that performs the denoising process for Flux2-dev."
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation."
+ " - `Flux2CoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2ImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
+ )
+
+
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", Flux2TextEncoderStep()),
+ ("vae_encoder", Flux2AutoVaeEncoderStep()),
+ ("denoise", Flux2AutoCoreDenoiseStep()),
+ ("decode", Flux2DecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2AutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Mistral3ForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
+ (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
+ (`Flux2Transformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (10, 20, 30)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ guidance_scale (`None`, *optional*, defaults to 4.0):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux2"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("images"),
+ ]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py
new file mode 100644
index 0000000000000000000000000000000000000000..9bab9aab065c4ee8f4999750a0427f49147b8fb4
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py
@@ -0,0 +1,399 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Flux2PrepareImageLatentsStep,
+ Flux2PrepareLatentsStep,
+ Flux2RoPEInputsStep,
+ Flux2SetTimestepsStep,
+)
+from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
+from .denoise import Flux2KleinDenoiseStep
+from .encoders import (
+ Flux2KleinTextEncoderStep,
+ Flux2VaeEncoderStep,
+)
+from .inputs import (
+ Flux2ProcessImagesInputStep,
+ Flux2TextInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+################
+# VAE encoder
+################
+
+
+# auto_docstring
+class Flux2KleinVaeEncoderSequentialStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that preprocesses and encodes the image inputs into their latent representations.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2-klein"
+
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations."
+
+
+# auto_docstring
+class Flux2KleinAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2-klein"
+
+ block_classes = [Flux2KleinVaeEncoderSequentialStep]
+ block_names = ["img_conditioning"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image inputs into their latent representations.\n"
+ "This is an auto pipeline block that works for image conditioning tasks.\n"
+ " - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+###
+### Core denoise
+###
+
+Flux2KleinCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2TextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2KleinDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image
+ generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+
+ block_classes = Flux2KleinCoreDenoiseBlocks.values()
+ block_names = Flux2KleinCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image generation."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+Flux2KleinImageConditionedCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2TextInputStep()),
+ ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2RoPEInputsStep()),
+ ("denoise", Flux2KleinDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+
+ block_classes = Flux2KleinImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2KleinImageConditionedCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# auto_docstring
+class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto core denoise step that performs the denoising process for Flux2-Klein.
+ This is an auto pipeline block that works for text-to-image and image-conditioned generation.
+ - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.
+ - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [Flux2KleinImageConditionedCoreDenoiseStep, Flux2KleinCoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoise step that performs the denoising process for Flux2-Klein.\n"
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n"
+ " - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
+ )
+
+
+###
+### Auto blocks
+###
+
+
+# auto_docstring
+class Flux2KleinAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2Tokenizer`) image_processor (`Flux2ImageProcessor`) vae
+ (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Configs:
+ is_distilled (default: True)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [
+ Flux2KleinTextEncoderStep(),
+ Flux2KleinAutoVaeEncoderStep(),
+ Flux2KleinAutoCoreDenoiseStep(),
+ Flux2DecodeStep(),
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("images"),
+ ]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ff9a764ac0551b0fac4390c19c0a4fc96e2e5d2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py
@@ -0,0 +1,413 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Flux2KleinBaseRoPEInputsStep,
+ Flux2PrepareImageLatentsStep,
+ Flux2PrepareLatentsStep,
+ Flux2SetTimestepsStep,
+)
+from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep
+from .denoise import Flux2KleinBaseDenoiseStep
+from .encoders import (
+ Flux2KleinBaseTextEncoderStep,
+ Flux2VaeEncoderStep,
+)
+from .inputs import (
+ Flux2KleinBaseTextInputStep,
+ Flux2ProcessImagesInputStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+################
+# VAE encoder
+################
+
+
+# auto_docstring
+class Flux2KleinBaseVaeEncoderSequentialStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that preprocesses and encodes the image inputs into their latent representations.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ model_name = "flux2"
+
+ block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations."
+
+
+# auto_docstring
+class Flux2KleinBaseAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image inputs into their latent representations.
+ This is an auto pipeline block that works for image conditioning tasks.
+ - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`)
+
+ Inputs:
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ condition_images (`list`):
+ TODO: Add description.
+ image_latents (`list`):
+ List of latent representations for each reference image
+ """
+
+ block_classes = [Flux2KleinBaseVaeEncoderSequentialStep]
+ block_names = ["img_conditioning"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image inputs into their latent representations.\n"
+ "This is an auto pipeline block that works for image conditioning tasks.\n"
+ " - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+###
+### Core denoise
+###
+
+Flux2KleinBaseCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2KleinBaseTextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()),
+ ("denoise", Flux2KleinBaseDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinBaseCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ Packed image latents for conditioning. Shape: (B, img_seq_len, C)
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = Flux2KleinBaseCoreDenoiseBlocks.values()
+ block_names = Flux2KleinBaseCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+Flux2KleinBaseImageConditionedCoreDenoiseBlocks = InsertableDict(
+ [
+ ("input", Flux2KleinBaseTextInputStep()),
+ ("prepare_latents", Flux2PrepareLatentsStep()),
+ ("prepare_image_latents", Flux2PrepareImageLatentsStep()),
+ ("set_timesteps", Flux2SetTimestepsStep()),
+ ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()),
+ ("denoise", Flux2KleinBaseDenoiseStep()),
+ ("after_denoise", Flux2UnpackLatentsStep()),
+ ]
+)
+
+
+# auto_docstring
+class Flux2KleinBaseImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.values()
+ block_names = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.keys()
+
+ @property
+ def description(self):
+ return "Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# auto_docstring
+class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto core denoise step that performs the denoising process for Flux2-Klein (base model).
+ This is an auto pipeline block that works for text-to-image and image-conditioned generation.
+ - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.
+ - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider
+ (`ClassifierFreeGuidance`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [Flux2KleinBaseImageConditionedCoreDenoiseStep, Flux2KleinBaseCoreDenoiseStep]
+ block_names = ["image_conditioned", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoise step that performs the denoising process for Flux2-Klein (base model).\n"
+ "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n"
+ " - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.\n"
+ " - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n"
+ )
+
+
+###
+### Auto blocks
+###
+
+
+# auto_docstring
+class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model).
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image_conditioned`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`)
+ image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
+
+ Configs:
+ is_distilled (default: False)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ TODO: Add description.
+ text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ image_latents (`list`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ timesteps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ joint_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_latent_ids (`Tensor`, *optional*):
+ Position IDs for image latents. Shape: (B, img_seq_len, 4)
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "flux2-klein"
+ block_classes = [
+ Flux2KleinBaseTextEncoderStep(),
+ Flux2KleinBaseAutoVaeEncoderStep(),
+ Flux2KleinBaseAutoCoreDenoiseStep(),
+ Flux2DecodeStep(),
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image_conditioned": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model)."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("images"),
+ ]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..d838cda6aff78b28240e30dcffd5bdf2c9b9426c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/flux2/modular_pipeline.py
@@ -0,0 +1,99 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...loaders import Flux2LoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class Flux2ModularPipeline(ModularPipeline, Flux2LoraLoaderMixin):
+ """
+ A ModularPipeline for Flux2.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "Flux2AutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ return 128
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if getattr(self, "vae", None) is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 32
+ if getattr(self, "transformer", None):
+ num_channels_latents = self.transformer.config.in_channels // 4
+ return num_channels_latents
+
+
+class Flux2KleinModularPipeline(Flux2ModularPipeline):
+ """
+ A ModularPipeline for Flux2-Klein (distilled model).
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "Flux2KleinAutoBlocks"
+
+ @property
+ def requires_unconditional_embeds(self):
+ if hasattr(self.config, "is_distilled") and self.config.is_distilled:
+ return False
+
+ requires_unconditional_embeds = False
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
+
+
+class Flux2KleinBaseModularPipeline(Flux2ModularPipeline):
+ """
+ A ModularPipeline for Flux2-Klein (base model).
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "Flux2KleinBaseAutoBlocks"
+
+ @property
+ def requires_unconditional_embeds(self):
+ if hasattr(self.config, "is_distilled") and self.config.is_distilled:
+ return False
+
+ requires_unconditional_embeds = False
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c86c1c64bd945dbac59554e4da823d402b6c4b6
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/__init__.py
@@ -0,0 +1,59 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_helios"] = ["HeliosAutoBlocks"]
+ _import_structure["modular_blocks_helios_pyramid"] = ["HeliosPyramidAutoBlocks"]
+ _import_structure["modular_blocks_helios_pyramid_distilled"] = ["HeliosPyramidDistilledAutoBlocks"]
+ _import_structure["modular_pipeline"] = [
+ "HeliosModularPipeline",
+ "HeliosPyramidDistilledModularPipeline",
+ "HeliosPyramidModularPipeline",
+ ]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_helios import HeliosAutoBlocks
+ from .modular_blocks_helios_pyramid import HeliosPyramidAutoBlocks
+ from .modular_blocks_helios_pyramid_distilled import HeliosPyramidDistilledAutoBlocks
+ from .modular_pipeline import (
+ HeliosModularPipeline,
+ HeliosPyramidDistilledModularPipeline,
+ HeliosPyramidModularPipeline,
+ )
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0271f0ab23b647725f0fa5e43576cf461cfbe6b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/before_denoise.py
@@ -0,0 +1,836 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+
+from ...models import HeliosTransformer3DModel
+from ...schedulers import HeliosScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import HeliosModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+class HeliosTextInputStep(ModularPipelineBlocks):
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_videos_per_prompt`\n\n"
+ "All input tensors are expected to have either batch_size=1 or match the batch_size\n"
+ "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n"
+ "have a final batch_size of batch_size * num_videos_per_prompt."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "num_videos_per_prompt",
+ default=1,
+ type_hint=int,
+ description="Number of videos to generate per prompt.",
+ ),
+ InputParam.template("prompt_embeds"),
+ InputParam.template("negative_prompt_embeds"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `prompt_embeds.dtype`)",
+ ),
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.prompt_embeds is not None and block_state.negative_prompt_embeds is not None:
+ if block_state.prompt_embeds.shape != block_state.negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {block_state.negative_prompt_embeds.shape}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_videos_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_videos_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# Copied from diffusers.modular_pipelines.wan.before_denoise.repeat_tensor_to_batch_size
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_videos_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_videos_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_videos_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_videos_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_videos_per_prompt (int, optional): Number of videos to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_videos_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_videos_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_videos_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_videos_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_videos_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_videos_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+# Copied from diffusers.modular_pipelines.wan.before_denoise.calculate_dimension_from_latents
+def calculate_dimension_from_latents(
+ latents: torch.Tensor, vae_scale_factor_temporal: int, vae_scale_factor_spatial: int
+) -> tuple[int, int]:
+ """Calculate image dimensions from latent tensor dimensions.
+
+ This function converts latent temporal and spatial dimensions to image temporal and spatial dimensions by
+ multiplying the latent num_frames/height/width by the VAE scale factor.
+
+ Args:
+ latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions.
+ Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width]
+ vae_scale_factor_temporal (int): The scale factor used by the VAE to compress temporal dimension.
+ Typically 4 for most VAEs (video is 4x larger than latents in temporal dimension)
+ vae_scale_factor_spatial (int): The scale factor used by the VAE to compress spatial dimension.
+ Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
+
+ Returns:
+ tuple[int, int]: The calculated image dimensions as (height, width)
+
+ Raises:
+ ValueError: If latents tensor doesn't have 4 or 5 dimensions
+
+ """
+ if latents.ndim != 5:
+ raise ValueError(f"latents must have 5 dimensions, but got {latents.ndim}")
+
+ _, _, num_latent_frames, latent_height, latent_width = latents.shape
+
+ num_frames = (num_latent_frames - 1) * vae_scale_factor_temporal + 1
+ height = latent_height * vae_scale_factor_spatial
+ width = latent_width * vae_scale_factor_spatial
+
+ return num_frames, height, width
+
+
+class HeliosAdditionalInputsStep(ModularPipelineBlocks):
+ """Configurable step that standardizes inputs for the denoising step.
+
+ This step handles:
+ 1. For encoded image latents: Computes height/width from latents and expands batch size
+ 2. For additional_batch_inputs: Expands batch dimensions to match final batch size
+ """
+
+ model_name = "helios"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
+ ):
+ if image_latent_inputs is None:
+ image_latent_inputs = [InputParam.template("image_latents")]
+ if additional_batch_inputs is None:
+ additional_batch_inputs = []
+
+ if not isinstance(image_latent_inputs, list):
+ raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}")
+ else:
+ for input_param in image_latent_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}")
+
+ if not isinstance(additional_batch_inputs, list):
+ raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}")
+ else:
+ for input_param in additional_batch_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(
+ f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}"
+ )
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ summary_section = (
+ "Input processing step that:\n"
+ " 1. For image latent inputs: Computes height/width from latents and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}"
+
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(name="num_videos_per_prompt", default=1),
+ InputParam(name="batch_size", required=True),
+ ]
+ inputs += self._image_latent_inputs + self._additional_batch_inputs
+
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ outputs = [
+ OutputParam("height", type_hint=int),
+ OutputParam("width", type_hint=int),
+ ]
+
+ for input_param in self._image_latent_inputs:
+ outputs.append(OutputParam(input_param.name, type_hint=torch.Tensor))
+
+ for input_param in self._additional_batch_inputs:
+ outputs.append(OutputParam(input_param.name, type_hint=torch.Tensor))
+
+ return outputs
+
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ for input_param in self._image_latent_inputs:
+ image_latent_tensor = getattr(block_state, input_param.name)
+ if image_latent_tensor is None:
+ continue
+
+ # Calculate height/width from latents
+ _, height, width = calculate_dimension_from_latents(
+ image_latent_tensor, components.vae_scale_factor_temporal, components.vae_scale_factor_spatial
+ )
+ block_state.height = height
+ block_state.width = width
+
+ # Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=input_param.name,
+ input_tensor=image_latent_tensor,
+ num_videos_per_prompt=block_state.num_videos_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_param.name, image_latent_tensor)
+
+ for input_param in self._additional_batch_inputs:
+ input_tensor = getattr(block_state, input_param.name)
+ if input_tensor is None:
+ continue
+
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_param.name,
+ input_tensor=input_tensor,
+ num_videos_per_prompt=block_state.num_videos_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_param.name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosAddNoiseToImageLatentsStep(ModularPipelineBlocks):
+ """Adds noise to image_latents and fake_image_latents for I2V conditioning.
+
+ Applies single-sigma noise to image_latents (using image_noise_sigma range) and single-sigma noise to
+ fake_image_latents (using video_noise_sigma range).
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Adds noise to image_latents and fake_image_latents for I2V conditioning. "
+ "Uses random sigma from configured ranges for each."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image_latents"),
+ InputParam(
+ "fake_image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Fake image latents used as history seed for I2V generation.",
+ ),
+ InputParam(
+ "image_noise_sigma_min",
+ default=0.111,
+ type_hint=float,
+ description="Minimum sigma for image latent noise.",
+ ),
+ InputParam(
+ "image_noise_sigma_max",
+ default=0.135,
+ type_hint=float,
+ description="Maximum sigma for image latent noise.",
+ ),
+ InputParam(
+ "video_noise_sigma_min",
+ default=0.111,
+ type_hint=float,
+ description="Minimum sigma for video/fake-image latent noise.",
+ ),
+ InputParam(
+ "video_noise_sigma_max",
+ default=0.135,
+ type_hint=float,
+ description="Maximum sigma for video/fake-image latent noise.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("image_latents"),
+ OutputParam("fake_image_latents", type_hint=torch.Tensor, description="Noisy fake image latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ image_latents = block_state.image_latents
+ fake_image_latents = block_state.fake_image_latents
+
+ # Add noise to image_latents
+ image_noise_sigma = (
+ torch.rand(1, device=device, generator=block_state.generator)
+ * (block_state.image_noise_sigma_max - block_state.image_noise_sigma_min)
+ + block_state.image_noise_sigma_min
+ )
+ image_latents = (
+ image_noise_sigma * randn_tensor(image_latents.shape, generator=block_state.generator, device=device)
+ + (1 - image_noise_sigma) * image_latents
+ )
+
+ # Add noise to fake_image_latents
+ fake_image_noise_sigma = (
+ torch.rand(1, device=device, generator=block_state.generator)
+ * (block_state.video_noise_sigma_max - block_state.video_noise_sigma_min)
+ + block_state.video_noise_sigma_min
+ )
+ fake_image_latents = (
+ fake_image_noise_sigma
+ * randn_tensor(fake_image_latents.shape, generator=block_state.generator, device=device)
+ + (1 - fake_image_noise_sigma) * fake_image_latents
+ )
+
+ block_state.image_latents = image_latents.to(device=device, dtype=torch.float32)
+ block_state.fake_image_latents = fake_image_latents.to(device=device, dtype=torch.float32)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosAddNoiseToVideoLatentsStep(ModularPipelineBlocks):
+ """Adds noise to image_latents and video_latents for V2V conditioning.
+
+ Applies single-sigma noise to image_latents (using image_noise_sigma range) and per-frame noise to video_latents in
+ chunks (using video_noise_sigma range).
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Adds noise to image_latents and video_latents for V2V conditioning. "
+ "Uses single-sigma noise for image_latents and per-frame noise for video chunks."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image_latents"),
+ InputParam(
+ "video_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Encoded video latents for V2V generation.",
+ ),
+ InputParam(
+ "num_latent_frames_per_chunk",
+ default=9,
+ type_hint=int,
+ description="Number of latent frames per temporal chunk.",
+ ),
+ InputParam(
+ "image_noise_sigma_min",
+ default=0.111,
+ type_hint=float,
+ description="Minimum sigma for image latent noise.",
+ ),
+ InputParam(
+ "image_noise_sigma_max",
+ default=0.135,
+ type_hint=float,
+ description="Maximum sigma for image latent noise.",
+ ),
+ InputParam(
+ "video_noise_sigma_min",
+ default=0.111,
+ type_hint=float,
+ description="Minimum sigma for video latent noise.",
+ ),
+ InputParam(
+ "video_noise_sigma_max",
+ default=0.135,
+ type_hint=float,
+ description="Maximum sigma for video latent noise.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("image_latents"),
+ OutputParam("video_latents", type_hint=torch.Tensor, description="Noisy video latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ image_latents = block_state.image_latents
+ video_latents = block_state.video_latents
+ num_latent_frames_per_chunk = block_state.num_latent_frames_per_chunk
+
+ # Add noise to first frame (single sigma)
+ image_noise_sigma = (
+ torch.rand(1, device=device, generator=block_state.generator)
+ * (block_state.image_noise_sigma_max - block_state.image_noise_sigma_min)
+ + block_state.image_noise_sigma_min
+ )
+ image_latents = (
+ image_noise_sigma * randn_tensor(image_latents.shape, generator=block_state.generator, device=device)
+ + (1 - image_noise_sigma) * image_latents
+ )
+
+ # Add per-frame noise to video chunks
+ noisy_latents_chunks = []
+ num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk
+ for i in range(num_latent_chunks):
+ chunk_start = i * num_latent_frames_per_chunk
+ chunk_end = chunk_start + num_latent_frames_per_chunk
+ latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :]
+
+ chunk_frames = latent_chunk.shape[2]
+ frame_sigmas = (
+ torch.rand(chunk_frames, device=device, generator=block_state.generator)
+ * (block_state.video_noise_sigma_max - block_state.video_noise_sigma_min)
+ + block_state.video_noise_sigma_min
+ )
+ frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1)
+
+ noisy_chunk = (
+ frame_sigmas * randn_tensor(latent_chunk.shape, generator=block_state.generator, device=device)
+ + (1 - frame_sigmas) * latent_chunk
+ )
+ noisy_latents_chunks.append(noisy_chunk)
+ video_latents = torch.cat(noisy_latents_chunks, dim=2)
+
+ block_state.image_latents = image_latents.to(device=device, dtype=torch.float32)
+ block_state.video_latents = video_latents.to(device=device, dtype=torch.float32)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosPrepareHistoryStep(ModularPipelineBlocks):
+ """Prepares chunk/history indices and initializes history state for the chunk loop."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepares the chunk loop by computing latent dimensions, number of chunks, "
+ "history indices, and initializing history state (history_latents, image_latents, latent_chunks)."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HeliosTransformer3DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", default=384),
+ InputParam.template("width", default=640),
+ InputParam(
+ "num_frames", default=132, type_hint=int, description="Total number of video frames to generate."
+ ),
+ InputParam("batch_size", required=True, type_hint=int),
+ InputParam(
+ "num_latent_frames_per_chunk",
+ default=9,
+ type_hint=int,
+ description="Number of latent frames per temporal chunk.",
+ ),
+ InputParam(
+ "history_sizes",
+ default=[16, 2, 1],
+ type_hint=list,
+ description="Sizes of long/mid/short history buffers for temporal context.",
+ ),
+ InputParam(
+ "keep_first_frame",
+ default=True,
+ type_hint=bool,
+ description="Whether to keep the first frame as a prefix in history.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("num_latent_chunk", type_hint=int, description="Number of temporal chunks"),
+ OutputParam("latent_shape", type_hint=tuple, description="Shape of latent tensor per chunk"),
+ OutputParam("history_sizes", type_hint=list, description="Adjusted history sizes (sorted, descending)"),
+ OutputParam("indices_hidden_states", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"),
+ OutputParam("indices_latents_history_short", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"),
+ OutputParam("indices_latents_history_mid", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"),
+ OutputParam("indices_latents_history_long", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"),
+ OutputParam("history_latents", type_hint=torch.Tensor, description="Initialized zero history latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ batch_size = block_state.batch_size
+ device = components._execution_device
+
+ block_state.num_frames = max(block_state.num_frames, 1)
+ history_sizes = sorted(block_state.history_sizes, reverse=True)
+
+ num_channels_latents = components.num_channels_latents
+ h_latent = block_state.height // components.vae_scale_factor_spatial
+ w_latent = block_state.width // components.vae_scale_factor_spatial
+
+ # Compute number of chunks
+ block_state.window_num_frames = (
+ block_state.num_latent_frames_per_chunk - 1
+ ) * components.vae_scale_factor_temporal + 1
+ block_state.num_latent_chunk = max(
+ 1, (block_state.num_frames + block_state.window_num_frames - 1) // block_state.window_num_frames
+ )
+
+ # Modify history_sizes for non-keep_first_frame (matching pipeline behavior)
+ if not block_state.keep_first_frame:
+ history_sizes = history_sizes.copy()
+ history_sizes[-1] = history_sizes[-1] + 1
+
+ # Compute indices ONCE (same structure for all chunks)
+ if block_state.keep_first_frame:
+ indices = torch.arange(0, sum([1, *history_sizes, block_state.num_latent_frames_per_chunk]))
+ (
+ indices_prefix,
+ indices_latents_history_long,
+ indices_latents_history_mid,
+ indices_latents_history_1x,
+ indices_hidden_states,
+ ) = indices.split([1, *history_sizes, block_state.num_latent_frames_per_chunk], dim=0)
+ indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0)
+ else:
+ indices = torch.arange(0, sum([*history_sizes, block_state.num_latent_frames_per_chunk]))
+ (
+ indices_latents_history_long,
+ indices_latents_history_mid,
+ indices_latents_history_short,
+ indices_hidden_states,
+ ) = indices.split([*history_sizes, block_state.num_latent_frames_per_chunk], dim=0)
+
+ # Latent shape per chunk
+ block_state.latent_shape = (
+ batch_size,
+ num_channels_latents,
+ block_state.num_latent_frames_per_chunk,
+ h_latent,
+ w_latent,
+ )
+
+ # Set outputs
+ block_state.history_sizes = history_sizes
+ block_state.indices_hidden_states = indices_hidden_states.unsqueeze(0)
+ block_state.indices_latents_history_short = indices_latents_history_short.unsqueeze(0)
+ block_state.indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0)
+ block_state.indices_latents_history_long = indices_latents_history_long.unsqueeze(0)
+ block_state.history_latents = torch.zeros(
+ batch_size,
+ num_channels_latents,
+ sum(history_sizes),
+ h_latent,
+ w_latent,
+ device=device,
+ dtype=torch.float32,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class HeliosI2VSeedHistoryStep(ModularPipelineBlocks):
+ """Seeds history_latents with fake_image_latents for I2V pipelines.
+
+ This small additive step runs after HeliosPrepareHistoryStep and appends fake_image_latents to the initialized
+ history_latents tensor.
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "I2V history seeding: appends fake_image_latents to history_latents."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("history_latents", required=True, type_hint=torch.Tensor),
+ InputParam("fake_image_latents", required=True, type_hint=torch.Tensor),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "history_latents", type_hint=torch.Tensor, description="History latents seeded with fake_image_latents"
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.history_latents = torch.cat([block_state.history_latents, block_state.fake_image_latents], dim=2)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosV2VSeedHistoryStep(ModularPipelineBlocks):
+ """Seeds history_latents with video_latents for V2V pipelines.
+
+ This step runs after HeliosPrepareHistoryStep and replaces the tail of history_latents with video_latents. If the
+ video has fewer frames than the history, the beginning of history is preserved.
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "V2V history seeding: replaces the tail of history_latents with video_latents."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("history_latents", required=True, type_hint=torch.Tensor),
+ InputParam("video_latents", required=True, type_hint=torch.Tensor),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "history_latents", type_hint=torch.Tensor, description="History latents seeded with video_latents"
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ history_latents = block_state.history_latents
+ video_latents = block_state.video_latents
+
+ history_frames = history_latents.shape[2]
+ video_frames = video_latents.shape[2]
+ if video_frames < history_frames:
+ keep_frames = history_frames - video_frames
+ history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2)
+ else:
+ history_latents = video_latents
+
+ block_state.history_latents = history_latents
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosSetTimestepsStep(ModularPipelineBlocks):
+ """Computes scheduler parameters (mu, sigmas) for the chunk loop."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "Computes scheduler shift parameter (mu) and default sigmas for the Helios chunk loop."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HeliosTransformer3DModel),
+ ComponentSpec("scheduler", HeliosScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latent_shape", required=True, type_hint=tuple),
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("mu", type_hint=float, description="Scheduler shift parameter"),
+ OutputParam("sigmas", type_hint=list, description="Sigma schedule for diffusion"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ patch_size = components.transformer.config.patch_size
+ latent_shape = block_state.latent_shape
+ image_seq_len = (latent_shape[-1] * latent_shape[-2] * latent_shape[-3]) // (
+ patch_size[0] * patch_size[1] * patch_size[2]
+ )
+
+ if block_state.sigmas is None:
+ block_state.sigmas = np.linspace(0.999, 0.0, block_state.num_inference_steps + 1)[:-1]
+
+ block_state.mu = calculate_shift(
+ image_seq_len,
+ components.scheduler.config.get("base_image_seq_len", 256),
+ components.scheduler.config.get("max_image_seq_len", 4096),
+ components.scheduler.config.get("base_shift", 0.5),
+ components.scheduler.config.get("max_shift", 1.15),
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..79bd569e7741f171da7abc57b05306fbaa6c36ce
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/decoders.py
@@ -0,0 +1,112 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLWan
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class HeliosDecodeStep(ModularPipelineBlocks):
+ """Decode all chunk latents with VAE, trim frames, and postprocess into final video output."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Decodes all chunk latents with the VAE, concatenates them, "
+ "trims to the target frame count, and postprocesses into the final video output."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latent_chunks", required=True, type_hint=list, description="List of per-chunk denoised latent tensors"
+ ),
+ InputParam("num_frames", required=True, type_hint=int, description="The target number of output frames"),
+ InputParam.template("output_type", default="np"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "videos",
+ type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray],
+ description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae = components.vae
+ device = components._execution_device
+ decode_dtype = vae.dtype
+
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(device, decode_dtype)
+ )
+ latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(
+ device, decode_dtype
+ )
+
+ history_video = None
+ for chunk_latents in block_state.latent_chunks:
+ current_latents = chunk_latents.to(device=device, dtype=decode_dtype) / latents_std + latents_mean
+ current_video = vae.decode(current_latents, return_dict=False)[0]
+
+ if history_video is None:
+ history_video = current_video
+ else:
+ history_video = torch.cat([history_video, current_video], dim=2)
+
+ # Trim to proper frame count
+ generated_frames = history_video.size(2)
+ generated_frames = (
+ generated_frames - 1
+ ) // components.vae_scale_factor_temporal * components.vae_scale_factor_temporal + 1
+ history_video = history_video[:, :, :generated_frames]
+
+ block_state.videos = components.video_processor.postprocess_video(
+ history_video, output_type=block_state.output_type
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..7dd8659c9d13bbce2860c7b628afe25940ba1463
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/denoise.py
@@ -0,0 +1,1069 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import math
+
+import torch
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance, ClassifierFreeZeroStarGuidance
+from ...models import HeliosTransformer3DModel
+from ...schedulers import HeliosScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .before_denoise import calculate_shift
+from .modular_pipeline import HeliosModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def sample_block_noise(
+ batch_size,
+ channel,
+ num_frames,
+ height,
+ width,
+ gamma,
+ patch_size=(1, 2, 2),
+ device=None,
+ generator=None,
+):
+ """Generate spatially-correlated block noise for pyramid upsampling correction.
+
+ Uses a multivariate normal distribution with covariance based on `gamma` to produce noise with block structure,
+ matching the upsampling artifacts that need correction.
+ """
+ # NOTE: A generator must be provided to ensure correct and reproducible results.
+ # Creating a default generator here is a fallback only — without a fixed seed,
+ # the output will be non-deterministic and may produce incorrect results in CP context.
+ if generator is None:
+ generator = torch.Generator(device=device)
+ elif isinstance(generator, list):
+ generator = generator[0]
+
+ _, ph, pw = patch_size
+ block_size = ph * pw
+
+ cov = (
+ torch.eye(block_size, device=device) * (1 + gamma) - torch.ones(block_size, block_size, device=device) * gamma
+ )
+ cov += torch.eye(block_size, device=device) * 1e-8
+ cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16.
+
+ L = torch.linalg.cholesky(cov)
+ block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
+ z = torch.randn(block_number, block_size, device=generator.device, generator=generator).to(device)
+ noise = z @ L.T
+
+ noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw)
+ noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width)
+ return noise
+
+
+# ========================================
+# Chunk Loop Leaf Blocks
+# ========================================
+
+
+class HeliosChunkHistorySliceStep(ModularPipelineBlocks):
+ """Slices history latents into short/mid/long for a T2V chunk.
+
+ At k==0 with no image_latents, creates a zero prefix. Otherwise uses image_latents (either provided or captured
+ from first chunk by HeliosChunkUpdateStep).
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "T2V history slice: splits history into long/mid/short. At k==0 with no image_latents, "
+ "creates a zero prefix; otherwise uses image_latents as prefix for short history."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "keep_first_frame",
+ default=True,
+ type_hint=bool,
+ description="Whether to keep the first frame as a prefix in history.",
+ ),
+ InputParam(
+ "history_sizes",
+ required=True,
+ type_hint=list,
+ description="Sizes of long/mid/short history buffers for temporal context.",
+ ),
+ InputParam(
+ "history_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Accumulated history latents from previous chunks.",
+ ),
+ InputParam("latent_shape", required=True, type_hint=tuple),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return []
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ keep_first_frame = block_state.keep_first_frame
+ history_sizes = block_state.history_sizes
+ image_latents = block_state.image_latents
+ device = components._execution_device
+
+ batch_size, num_channels_latents, _, h_latent, w_latent = block_state.latent_shape
+
+ if keep_first_frame:
+ latents_history_long, latents_history_mid, latents_history_1x = block_state.history_latents[
+ :, :, -sum(history_sizes) :
+ ].split(history_sizes, dim=2)
+ if image_latents is None and k == 0:
+ latents_prefix = torch.zeros(
+ batch_size,
+ num_channels_latents,
+ 1,
+ h_latent,
+ w_latent,
+ device=device,
+ dtype=torch.float32,
+ )
+ else:
+ latents_prefix = image_latents
+ latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2)
+ else:
+ latents_history_long, latents_history_mid, latents_history_short = block_state.history_latents[
+ :, :, -sum(history_sizes) :
+ ].split(history_sizes, dim=2)
+
+ block_state.latents_history_short = latents_history_short
+ block_state.latents_history_mid = latents_history_mid
+ block_state.latents_history_long = latents_history_long
+
+ return components, block_state
+
+
+class HeliosI2VChunkHistorySliceStep(ModularPipelineBlocks):
+ """Slices history latents into short/mid/long for an I2V chunk.
+
+ Always uses image_latents as prefix (assumes history pre-seeded with fake_image_latents).
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "I2V history slice: splits pre-seeded history into long/mid/short, "
+ "always using image_latents as prefix for short history."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "keep_first_frame",
+ default=True,
+ type_hint=bool,
+ description="Whether to keep the first frame as a prefix in history.",
+ ),
+ InputParam(
+ "history_sizes",
+ required=True,
+ type_hint=list,
+ description="Sizes of long/mid/short history buffers for temporal context.",
+ ),
+ InputParam(
+ "history_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Accumulated history latents from previous chunks.",
+ ),
+ InputParam(
+ "image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="First-frame latents used as prefix for short history.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return []
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ keep_first_frame = block_state.keep_first_frame
+ history_sizes = block_state.history_sizes
+ image_latents = block_state.image_latents
+
+ if keep_first_frame:
+ latents_history_long, latents_history_mid, latents_history_1x = block_state.history_latents[
+ :, :, -sum(history_sizes) :
+ ].split(history_sizes, dim=2)
+ latents_history_short = torch.cat([image_latents, latents_history_1x], dim=2)
+ else:
+ latents_history_long, latents_history_mid, latents_history_short = block_state.history_latents[
+ :, :, -sum(history_sizes) :
+ ].split(history_sizes, dim=2)
+
+ block_state.latents_history_short = latents_history_short
+ block_state.latents_history_mid = latents_history_mid
+ block_state.latents_history_long = latents_history_long
+
+ return components, block_state
+
+
+class HeliosChunkNoiseGenStep(ModularPipelineBlocks):
+ """Generates noise latents for a chunk using randn_tensor."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "Generates random noise latents at full resolution for a single chunk."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latent_shape", required=True, type_hint=tuple),
+ InputParam.template("generator"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ device = components._execution_device
+ block_state.latents = randn_tensor(
+ block_state.latent_shape, generator=block_state.generator, device=device, dtype=torch.float32
+ )
+ return components, block_state
+
+
+class HeliosPyramidChunkNoiseGenStep(ModularPipelineBlocks):
+ """Generates noise latents and downsamples to smallest pyramid level."""
+
+ model_name = "helios-pyramid"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Generates random noise at full resolution, then downsamples to the smallest "
+ "pyramid level via bilinear interpolation."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latent_shape", required=True, type_hint=tuple),
+ InputParam(
+ "pyramid_num_inference_steps_list",
+ default=[10, 10, 10],
+ type_hint=list,
+ description="Number of denoising steps per pyramid stage.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ device = components._execution_device
+ batch_size, num_channels_latents, num_latent_frames, h_latent, w_latent = block_state.latent_shape
+
+ latents = randn_tensor(
+ block_state.latent_shape, generator=block_state.generator, device=device, dtype=torch.float32
+ )
+
+ # Downsample to smallest pyramid level
+ h, w = h_latent, w_latent
+ latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_latent_frames, num_channels_latents, h, w)
+ for _ in range(len(block_state.pyramid_num_inference_steps_list) - 1):
+ h //= 2
+ w //= 2
+ latents = F.interpolate(latents, size=(h, w), mode="bilinear") * 2
+ block_state.latents = latents.reshape(batch_size, num_latent_frames, num_channels_latents, h, w).permute(
+ 0, 2, 1, 3, 4
+ )
+
+ return components, block_state
+
+
+class HeliosChunkSchedulerResetStep(ModularPipelineBlocks):
+ """Resets the scheduler with timesteps for a single chunk."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "Resets the scheduler with the correct timesteps and shift parameter (mu) for this chunk."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", HeliosScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("mu", required=True, type_hint=float),
+ InputParam.template("sigmas", required=True),
+ InputParam.template("num_inference_steps"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ device = components._execution_device
+ components.scheduler.set_timesteps(
+ block_state.num_inference_steps, device=device, sigmas=block_state.sigmas, mu=block_state.mu
+ )
+ block_state.timesteps = components.scheduler.timesteps
+
+ return components, block_state
+
+
+# ========================================
+# Inner Denoising Blocks
+# ========================================
+
+
+class HeliosChunkDenoiseInner(ModularPipelineBlocks):
+ """Inner timestep loop for denoising a single chunk, using guider for guidance."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Inner denoising loop that iterates over timesteps for a single chunk. "
+ "Uses the guider to manage conditional/unconditional forward passes with cache_context, "
+ "applies guidance, and runs scheduler step."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HeliosTransformer3DModel),
+ ComponentSpec("scheduler", HeliosScheduler),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam.template("timesteps"),
+ InputParam("prompt_embeds", type_hint=torch.Tensor),
+ InputParam("negative_prompt_embeds", type_hint=torch.Tensor),
+ InputParam.template("denoiser_input_fields"),
+ InputParam.template("num_inference_steps"),
+ InputParam.template("attention_kwargs"),
+ InputParam.template("generator"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ latents = block_state.latents
+ timesteps = block_state.timesteps
+ num_inference_steps = block_state.num_inference_steps
+
+ transformer_dtype = components.transformer.dtype
+ num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order
+
+ # Guider inputs: only encoder_hidden_states differs between cond/uncond
+ guider_inputs = {
+ "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds),
+ }
+
+ # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones)
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+ shared_kwargs = {}
+ for field_name, field_value in block_state.denoiser_input_fields.items():
+ if field_name in transformer_args and field_name not in guider_inputs:
+ shared_kwargs[field_name] = field_value
+
+ # Add loop-internal history latents with dtype casting
+ shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype)
+ shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype)
+ shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype)
+ shared_kwargs["attention_kwargs"] = block_state.attention_kwargs
+
+ with tqdm(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ timestep = t.expand(latents.shape[0]).to(torch.int64)
+ latent_model_input = latents.to(transformer_dtype)
+
+ components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {k: getattr(guider_state_batch, k) for k in guider_inputs.keys()}
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep,
+ return_dict=False,
+ **cond_kwargs,
+ **shared_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ noise_pred = components.guider(guider_state)[0]
+
+ # Scheduler step
+ latents = components.scheduler.step(
+ noise_pred,
+ t,
+ latents,
+ generator=block_state.generator,
+ return_dict=False,
+ )[0]
+
+ if i == len(timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ block_state.latents = latents
+ return components, block_state
+
+
+class HeliosPyramidChunkDenoiseInner(ModularPipelineBlocks):
+ """Nested pyramid stage loop with inner timestep denoising.
+
+ For each pyramid stage (small -> full resolution):
+ 1. Upsample latents + block noise correction (stages > 0)
+ 2. Compute mu from current resolution, set scheduler timesteps
+ 3. Run timestep denoising loop (same logic as HeliosChunkDenoiseInner)
+ """
+
+ model_name = "helios-pyramid"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pyramid denoising inner block: loops over pyramid stages from smallest to full resolution. "
+ "Each stage upsamples latents (with block noise correction), recomputes scheduler parameters, "
+ "and runs the timestep denoising loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HeliosTransformer3DModel),
+ ComponentSpec("scheduler", HeliosScheduler),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeZeroStarGuidance,
+ config=FrozenDict({"guidance_scale": 5.0, "zero_init_steps": 2}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam("prompt_embeds", type_hint=torch.Tensor),
+ InputParam("negative_prompt_embeds", type_hint=torch.Tensor),
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ "pyramid_num_inference_steps_list",
+ default=[10, 10, 10],
+ type_hint=list,
+ description="Number of denoising steps per pyramid stage.",
+ ),
+ InputParam.template("attention_kwargs"),
+ InputParam.template("generator"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ device = components._execution_device
+ transformer_dtype = components.transformer.dtype
+ latents = block_state.latents
+ pyramid_num_stages = len(block_state.pyramid_num_inference_steps_list)
+
+ # Guider inputs: only encoder_hidden_states differs between cond/uncond
+ guider_inputs = {
+ "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds),
+ }
+
+ # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones)
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+ shared_kwargs = {}
+ for field_name, field_value in block_state.denoiser_input_fields.items():
+ if field_name in transformer_args and field_name not in guider_inputs:
+ shared_kwargs[field_name] = field_value
+
+ # Add loop-internal history latents with dtype casting
+ shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype)
+ shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype)
+ shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype)
+ shared_kwargs["attention_kwargs"] = block_state.attention_kwargs
+
+ # Save original zero_init_steps if the guider supports it (e.g. ClassifierFreeZeroStarGuidance).
+ # Helios only applies zero init in pyramid stage 0 (lowest resolution), so we disable it
+ # for subsequent stages by temporarily setting zero_init_steps=0.
+ orig_zero_init_steps = getattr(components.guider, "zero_init_steps", None)
+
+ for i_s in range(pyramid_num_stages):
+ # --- Stage setup ---
+
+ # Disable zero init for stages > 0 (only stage 0 should have zero init)
+ if orig_zero_init_steps is not None and i_s > 0:
+ components.guider.zero_init_steps = 0
+
+ # a. Compute mu from current resolution (before upsample, matching standard pipeline)
+ patch_size = components.transformer.config.patch_size
+ image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // (
+ patch_size[0] * patch_size[1] * patch_size[2]
+ )
+ mu = calculate_shift(
+ image_seq_len,
+ components.scheduler.config.get("base_image_seq_len", 256),
+ components.scheduler.config.get("max_image_seq_len", 4096),
+ components.scheduler.config.get("base_shift", 0.5),
+ components.scheduler.config.get("max_shift", 1.15),
+ )
+
+ # b. Set scheduler timesteps for this stage
+ num_inference_steps = block_state.pyramid_num_inference_steps_list[i_s]
+ components.scheduler.set_timesteps(
+ num_inference_steps,
+ i_s,
+ device=device,
+ mu=mu,
+ )
+ timesteps = components.scheduler.timesteps
+
+ # c. Upsample + block noise correction for stages > 0
+ if i_s > 0:
+ batch_size, num_channels_latents, num_frames, current_h, current_w = latents.shape
+ new_h = current_h * 2
+ new_w = current_w * 2
+
+ latents = latents.permute(0, 2, 1, 3, 4).reshape(
+ batch_size * num_frames, num_channels_latents, current_h, current_w
+ )
+ latents = F.interpolate(latents, size=(new_h, new_w), mode="nearest")
+ latents = latents.reshape(batch_size, num_frames, num_channels_latents, new_h, new_w).permute(
+ 0, 2, 1, 3, 4
+ )
+
+ # Block noise correction
+ ori_sigma = 1 - components.scheduler.ori_start_sigmas[i_s]
+ gamma = components.scheduler.config.gamma
+ alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
+ beta = alpha * (1 - ori_sigma) / math.sqrt(gamma)
+
+ batch_size, num_channels_latents, num_frames, h, w = latents.shape
+ noise = sample_block_noise(
+ batch_size,
+ num_channels_latents,
+ num_frames,
+ h,
+ w,
+ gamma,
+ patch_size,
+ device=device,
+ generator=block_state.generator,
+ )
+ noise = noise.to(dtype=transformer_dtype)
+ latents = alpha * latents + beta * noise
+
+ # --- Timestep denoising loop ---
+ num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order
+
+ with tqdm(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ timestep = t.expand(latents.shape[0]).to(torch.int64)
+ latent_model_input = latents.to(transformer_dtype)
+
+ components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {kk: getattr(guider_state_batch, kk) for kk in guider_inputs.keys()}
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep,
+ return_dict=False,
+ **cond_kwargs,
+ **shared_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ noise_pred = components.guider(guider_state)[0]
+
+ # Scheduler step
+ latents = components.scheduler.step(
+ noise_pred,
+ t,
+ latents,
+ generator=block_state.generator,
+ return_dict=False,
+ )[0]
+
+ if i == len(timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ # Restore original zero_init_steps
+ if orig_zero_init_steps is not None:
+ components.guider.zero_init_steps = orig_zero_init_steps
+
+ block_state.latents = latents
+ return components, block_state
+
+
+# ========================================
+# Post-Denoise Update
+# ========================================
+
+
+class HeliosChunkUpdateStep(ModularPipelineBlocks):
+ """Updates chunk collection and history after denoising a single chunk."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Post-denoising update step: appends the denoised latents to the chunk list, "
+ "captures image_latents from the first chunk if needed, and extends history_latents."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latents", type_hint=torch.Tensor),
+ InputParam("history_latents", type_hint=torch.Tensor),
+ InputParam("keep_first_frame", default=True, type_hint=bool),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ # e. Collect denoised latents for this chunk
+ block_state.latent_chunks.append(block_state.latents)
+
+ # f. Update history
+ if block_state.keep_first_frame and k == 0 and block_state.image_latents is None:
+ block_state.image_latents = block_state.latents[:, :, 0:1, :, :]
+
+ block_state.history_latents = torch.cat([block_state.history_latents, block_state.latents], dim=2)
+
+ return components, block_state
+
+
+# ========================================
+# Chunk Loop Wrapper
+# ========================================
+
+
+class HeliosChunkLoopWrapper(LoopSequentialPipelineBlocks):
+ """Outer chunk loop that iterates over temporal chunks.
+
+ History indices, scheduler params, and history state are prepared by HeliosPrepareHistoryStep and
+ HeliosSetTimestepsStep before this block runs. Sub-blocks handle per-chunk preparation, denoising, and history
+ updates.
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iterates over temporal chunks for progressive video generation. "
+ "At each chunk iteration, it runs sub-blocks for preparation, denoising, and history updates."
+ )
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_latent_chunk", required=True, type_hint=int),
+ ]
+
+ @property
+ def loop_intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.latent_chunks = []
+
+ if not hasattr(block_state, "image_latents"):
+ block_state.image_latents = None
+
+ for k in range(block_state.num_latent_chunk):
+ components, block_state = self.loop_step(components, block_state, k=k)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# ========================================
+# Composed Chunk Denoise Steps
+# ========================================
+
+
+class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """T2V chunk-based denoising: history slice -> noise gen -> scheduler reset -> denoise -> update."""
+
+ block_classes = [
+ HeliosChunkHistorySliceStep,
+ HeliosChunkNoiseGenStep,
+ HeliosChunkSchedulerResetStep,
+ HeliosChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "scheduler_reset", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "T2V chunk denoise step that iterates over temporal chunks.\n"
+ "At each chunk: history_slice -> noise_gen -> scheduler_reset -> denoise_inner -> update_chunk."
+ )
+
+
+class HeliosI2VChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """I2V chunk-based denoising: I2V history slice -> noise gen -> scheduler reset -> denoise -> update."""
+
+ block_classes = [
+ HeliosI2VChunkHistorySliceStep,
+ HeliosChunkNoiseGenStep,
+ HeliosChunkSchedulerResetStep,
+ HeliosChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "scheduler_reset", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "I2V chunk denoise step that iterates over temporal chunks.\n"
+ "At each chunk: history_slice (I2V) -> noise_gen -> scheduler_reset -> denoise_inner -> update_chunk."
+ )
+
+
+class HeliosPyramidDistilledChunkDenoiseInner(ModularPipelineBlocks):
+ """Nested pyramid stage loop with DMD denoising for distilled checkpoints.
+
+ Same progressive multi-resolution strategy as HeliosPyramidChunkDenoiseInner, but:
+ - Guidance is disabled (guidance_scale=1.0, no unconditional pass)
+ - Supports is_amplify_first_chunk (doubles first chunk's timesteps via scheduler)
+ - Tracks start_point_list and passes DMD-specific args to scheduler.step()
+ """
+
+ model_name = "helios-pyramid"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Distilled pyramid denoising inner block for DMD checkpoints. Loops over pyramid stages "
+ "from smallest to full resolution with guidance disabled and DMD scheduler support."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HeliosTransformer3DModel),
+ ComponentSpec("scheduler", HeliosScheduler),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 1.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam("prompt_embeds", type_hint=torch.Tensor),
+ InputParam("negative_prompt_embeds", type_hint=torch.Tensor),
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ "pyramid_num_inference_steps_list",
+ default=[2, 2, 2],
+ type_hint=list,
+ description="Number of denoising steps per pyramid stage.",
+ ),
+ InputParam(
+ "is_amplify_first_chunk",
+ default=True,
+ type_hint=bool,
+ description="Whether to double the first chunk's timesteps via the scheduler for amplified generation.",
+ ),
+ InputParam.template("attention_kwargs"),
+ InputParam.template("generator"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int):
+ device = components._execution_device
+ transformer_dtype = components.transformer.dtype
+ latents = block_state.latents
+ pyramid_num_stages = len(block_state.pyramid_num_inference_steps_list)
+ is_first_chunk = k == 0
+
+ # Track start points for DMD scheduler
+ start_point_list = [latents]
+
+ # Guider inputs: only encoder_hidden_states differs between cond/uncond
+ guider_inputs = {
+ "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds),
+ }
+
+ # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones)
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+ shared_kwargs = {}
+ for field_name, field_value in block_state.denoiser_input_fields.items():
+ if field_name in transformer_args and field_name not in guider_inputs:
+ shared_kwargs[field_name] = field_value
+
+ # Add loop-internal history latents with dtype casting
+ shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype)
+ shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype)
+ shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype)
+ shared_kwargs["attention_kwargs"] = block_state.attention_kwargs
+
+ for i_s in range(pyramid_num_stages):
+ # --- Stage setup ---
+ patch_size = components.transformer.config.patch_size
+
+ # a. Compute mu from current resolution (before upsample, matching standard pipeline)
+ image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // (
+ patch_size[0] * patch_size[1] * patch_size[2]
+ )
+ mu = calculate_shift(
+ image_seq_len,
+ components.scheduler.config.get("base_image_seq_len", 256),
+ components.scheduler.config.get("max_image_seq_len", 4096),
+ components.scheduler.config.get("base_shift", 0.5),
+ components.scheduler.config.get("max_shift", 1.15),
+ )
+
+ # b. Set scheduler timesteps for this stage (with DMD amplification)
+ num_inference_steps = block_state.pyramid_num_inference_steps_list[i_s]
+ components.scheduler.set_timesteps(
+ num_inference_steps,
+ i_s,
+ device=device,
+ mu=mu,
+ is_amplify_first_chunk=block_state.is_amplify_first_chunk and is_first_chunk,
+ )
+ timesteps = components.scheduler.timesteps
+
+ # c. Upsample + block noise correction for stages > 0
+ if i_s > 0:
+ batch_size, num_channels_latents, num_frames, current_h, current_w = latents.shape
+ new_h = current_h * 2
+ new_w = current_w * 2
+
+ latents = latents.permute(0, 2, 1, 3, 4).reshape(
+ batch_size * num_frames, num_channels_latents, current_h, current_w
+ )
+ latents = F.interpolate(latents, size=(new_h, new_w), mode="nearest")
+ latents = latents.reshape(batch_size, num_frames, num_channels_latents, new_h, new_w).permute(
+ 0, 2, 1, 3, 4
+ )
+
+ # Block noise correction
+ ori_sigma = 1 - components.scheduler.ori_start_sigmas[i_s]
+ gamma = components.scheduler.config.gamma
+ alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
+ beta = alpha * (1 - ori_sigma) / math.sqrt(gamma)
+
+ batch_size, num_channels_latents, num_frames, h, w = latents.shape
+ noise = sample_block_noise(
+ batch_size,
+ num_channels_latents,
+ num_frames,
+ h,
+ w,
+ gamma,
+ patch_size,
+ device=device,
+ generator=block_state.generator,
+ )
+ noise = noise.to(dtype=transformer_dtype)
+ latents = alpha * latents + beta * noise
+
+ start_point_list.append(latents)
+
+ # --- Timestep denoising loop ---
+ num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order
+
+ with tqdm(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ timestep = t.expand(latents.shape[0]).to(torch.int64)
+ latent_model_input = latents.to(transformer_dtype)
+
+ components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {k: getattr(guider_state_batch, k) for k in guider_inputs.keys()}
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep,
+ return_dict=False,
+ **cond_kwargs,
+ **shared_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ noise_pred = components.guider(guider_state)[0]
+
+ # Scheduler step with DMD args
+ latents = components.scheduler.step(
+ noise_pred,
+ t,
+ latents,
+ generator=block_state.generator,
+ return_dict=False,
+ cur_sampling_step=i,
+ dmd_noisy_tensor=start_point_list[i_s],
+ dmd_sigmas=components.scheduler.sigmas,
+ dmd_timesteps=components.scheduler.timesteps,
+ all_timesteps=timesteps,
+ )[0]
+
+ if i == len(timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ block_state.latents = latents
+ return components, block_state
+
+
+class HeliosPyramidChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """T2V pyramid chunk denoising: history slice -> pyramid noise gen -> pyramid denoise inner -> update."""
+
+ block_classes = [
+ HeliosChunkHistorySliceStep,
+ HeliosPyramidChunkNoiseGenStep,
+ HeliosPyramidChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "T2V pyramid chunk denoise step that iterates over temporal chunks.\n"
+ "At each chunk: history_slice -> noise_gen (pyramid) -> denoise_inner (pyramid stages) -> update_chunk.\n"
+ "Denoising starts at the smallest resolution and progressively upsamples."
+ )
+
+
+class HeliosPyramidI2VChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """I2V pyramid chunk denoising: I2V history slice -> pyramid noise gen -> pyramid denoise inner -> update."""
+
+ block_classes = [
+ HeliosI2VChunkHistorySliceStep,
+ HeliosPyramidChunkNoiseGenStep,
+ HeliosPyramidChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "I2V pyramid chunk denoise step that iterates over temporal chunks.\n"
+ "At each chunk: history_slice (I2V) -> noise_gen (pyramid) -> denoise_inner (pyramid stages) -> update_chunk.\n"
+ "Denoising starts at the smallest resolution and progressively upsamples."
+ )
+
+
+class HeliosPyramidDistilledChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """T2V distilled pyramid chunk denoising with DMD scheduler and no CFG."""
+
+ block_classes = [
+ HeliosChunkHistorySliceStep,
+ HeliosPyramidChunkNoiseGenStep,
+ HeliosPyramidDistilledChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "T2V distilled pyramid chunk denoise step with DMD scheduler.\n"
+ "At each chunk: history_slice -> noise_gen (pyramid) -> denoise_inner (distilled/DMD) -> update_chunk."
+ )
+
+
+class HeliosPyramidDistilledI2VChunkDenoiseStep(HeliosChunkLoopWrapper):
+ """I2V distilled pyramid chunk denoising with DMD scheduler and no CFG."""
+
+ block_classes = [
+ HeliosI2VChunkHistorySliceStep,
+ HeliosPyramidChunkNoiseGenStep,
+ HeliosPyramidDistilledChunkDenoiseInner,
+ HeliosChunkUpdateStep,
+ ]
+ block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "I2V distilled pyramid chunk denoise step with DMD scheduler.\n"
+ "At each chunk: history_slice (I2V) -> noise_gen (pyramid) -> denoise_inner (distilled/DMD) -> update_chunk."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..5644279715630921e59dc6c3abd9ba815d163ffc
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/encoders.py
@@ -0,0 +1,392 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import html
+
+import regex as re
+import torch
+from transformers import AutoTokenizer, UMT5EncoderModel
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import AutoencoderKLWan
+from ...utils import is_ftfy_available, logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import HeliosModularPipeline
+
+
+if is_ftfy_available():
+ import ftfy
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def basic_clean(text):
+ text = ftfy.fix_text(text)
+ text = html.unescape(html.unescape(text))
+ return text.strip()
+
+
+def whitespace_clean(text):
+ text = re.sub(r"\s+", " ", text)
+ text = text.strip()
+ return text
+
+
+def prompt_clean(text):
+ text = whitespace_clean(basic_clean(text))
+ return text
+
+
+def get_t5_prompt_embeds(
+ text_encoder: UMT5EncoderModel,
+ tokenizer: AutoTokenizer,
+ prompt: str | list[str],
+ max_sequence_length: int,
+ device: torch.device,
+ dtype: torch.dtype | None = None,
+):
+ """Encode text prompts into T5 embeddings for Helios.
+
+ Args:
+ text_encoder: The T5 text encoder model.
+ tokenizer: The tokenizer for the text encoder.
+ prompt: The prompt or prompts to encode.
+ max_sequence_length: Maximum sequence length for tokenization.
+ device: Device to place tensors on.
+ dtype: Optional dtype override. Defaults to `text_encoder.dtype`.
+
+ Returns:
+ A tuple of `(prompt_embeds, attention_mask)` where `prompt_embeds` is the encoded text embeddings and
+ `attention_mask` is a boolean mask.
+ """
+ dtype = dtype or text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ prompt = [prompt_clean(u) for u in prompt]
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+ seq_lens = mask.gt(0).sum(dim=1).long()
+
+ prompt_embeds = text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+ prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
+ )
+
+ return prompt_embeds, text_inputs.attention_mask.bool()
+
+
+class HeliosTextEncoderStep(ModularPipelineBlocks):
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", UMT5EncoderModel),
+ ComponentSpec("tokenizer", AutoTokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam.template("max_sequence_length"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds"),
+ ]
+
+ @staticmethod
+ def check_inputs(prompt, negative_prompt):
+ if prompt is not None and not isinstance(prompt, (str, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if negative_prompt is not None and not isinstance(negative_prompt, (str, list)):
+ raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+ if prompt is not None and negative_prompt is not None:
+ prompt_list = [prompt] if isinstance(prompt, str) else prompt
+ neg_list = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+ if type(prompt_list) is not type(neg_list):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ if len(prompt_list) != len(neg_list):
+ raise ValueError(
+ f"`negative_prompt` has batch size {len(neg_list)}, but `prompt` has batch size"
+ f" {len(prompt_list)}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt = block_state.prompt
+ negative_prompt = block_state.negative_prompt
+ max_sequence_length = block_state.max_sequence_length
+ device = components._execution_device
+
+ self.check_inputs(prompt, negative_prompt)
+
+ # Encode prompt
+ block_state.prompt_embeds, _ = get_t5_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ # Encode negative prompt
+ block_state.negative_prompt_embeds = None
+ if components.requires_unconditional_embeds:
+ negative_prompt = negative_prompt or ""
+ if isinstance(prompt, list) and isinstance(negative_prompt, str):
+ negative_prompt = len(prompt) * [negative_prompt]
+
+ block_state.negative_prompt_embeds, _ = get_t5_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosImageVaeEncoderStep(ModularPipelineBlocks):
+ """Encodes an input image into VAE latent space for image-to-video generation."""
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Image Encoder step that encodes an input image into VAE latent space, "
+ "producing image_latents (first frame prefix) and fake_image_latents (history seed) "
+ "for image-to-video generation."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image"),
+ InputParam.template("height", default=384),
+ InputParam.template("width", default=640),
+ InputParam(
+ "num_latent_frames_per_chunk",
+ default=9,
+ type_hint=int,
+ description="Number of latent frames per temporal chunk.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("image_latents"),
+ OutputParam(
+ "fake_image_latents", type_hint=torch.Tensor, description="Fake image latents for history seeding"
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae = components.vae
+ device = components._execution_device
+
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(
+ vae.device, vae.dtype
+ )
+
+ # Preprocess image to 4D tensor (B, C, H, W)
+ image = components.video_processor.preprocess(
+ block_state.image, height=block_state.height, width=block_state.width
+ )
+ image_5d = image.unsqueeze(2).to(device=device, dtype=vae.dtype) # (B, C, 1, H, W)
+
+ # Encode image to get image_latents
+ image_latents = vae.encode(image_5d).latent_dist.sample(generator=block_state.generator)
+ image_latents = (image_latents - latents_mean) * latents_std
+
+ # Encode fake video to get fake_image_latents
+ min_frames = (block_state.num_latent_frames_per_chunk - 1) * components.vae_scale_factor_temporal + 1
+ fake_video = image_5d.repeat(1, 1, min_frames, 1, 1) # (B, C, min_frames, H, W)
+ fake_latents_full = vae.encode(fake_video).latent_dist.sample(generator=block_state.generator)
+ fake_latents_full = (fake_latents_full - latents_mean) * latents_std
+ fake_image_latents = fake_latents_full[:, :, -1:, :, :]
+
+ block_state.image_latents = image_latents.to(device=device, dtype=torch.float32)
+ block_state.fake_image_latents = fake_image_latents.to(device=device, dtype=torch.float32)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HeliosVideoVaeEncoderStep(ModularPipelineBlocks):
+ """Encodes an input video into VAE latent space for video-to-video generation.
+
+ Produces `image_latents` (first frame) and `video_latents` (remaining frames encoded in chunks).
+ """
+
+ model_name = "helios"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Video Encoder step that encodes an input video into VAE latent space, "
+ "producing image_latents (first frame) and video_latents (chunked video frames) "
+ "for video-to-video generation."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("video", required=True, description="Input video for video-to-video generation"),
+ InputParam.template("height", default=384),
+ InputParam.template("width", default=640),
+ InputParam(
+ "num_latent_frames_per_chunk",
+ default=9,
+ type_hint=int,
+ description="Number of latent frames per temporal chunk.",
+ ),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("image_latents"),
+ OutputParam("video_latents", type_hint=torch.Tensor, description="Encoded video latents (chunked)"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae = components.vae
+ device = components._execution_device
+ num_latent_frames_per_chunk = block_state.num_latent_frames_per_chunk
+
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(
+ vae.device, vae.dtype
+ )
+
+ # Preprocess video
+ video = components.video_processor.preprocess_video(
+ block_state.video, height=block_state.height, width=block_state.width
+ )
+ video = video.to(device=device, dtype=vae.dtype)
+
+ # Encode video into latents
+ num_frames = video.shape[2]
+ min_frames = (num_latent_frames_per_chunk - 1) * 4 + 1
+ num_chunks = num_frames // min_frames
+ if num_chunks == 0:
+ raise ValueError(
+ f"Video must have at least {min_frames} frames "
+ f"(got {num_frames} frames). "
+ f"Required: (num_latent_frames_per_chunk - 1) * 4 + 1 = ({num_latent_frames_per_chunk} - 1) * 4 + 1 = {min_frames}"
+ )
+ total_valid_frames = num_chunks * min_frames
+ start_frame = num_frames - total_valid_frames
+
+ # Encode first frame
+ first_frame = video[:, :, 0:1, :, :]
+ image_latents = vae.encode(first_frame).latent_dist.sample(generator=block_state.generator)
+ image_latents = (image_latents - latents_mean) * latents_std
+
+ # Encode remaining frames in chunks
+ latents_chunks = []
+ for i in range(num_chunks):
+ chunk_start = start_frame + i * min_frames
+ chunk_end = chunk_start + min_frames
+ video_chunk = video[:, :, chunk_start:chunk_end, :, :]
+ chunk_latents = vae.encode(video_chunk).latent_dist.sample(generator=block_state.generator)
+ chunk_latents = (chunk_latents - latents_mean) * latents_std
+ latents_chunks.append(chunk_latents)
+ video_latents = torch.cat(latents_chunks, dim=2)
+
+ block_state.image_latents = image_latents.to(device=device, dtype=torch.float32)
+ block_state.video_latents = video_latents.to(device=device, dtype=torch.float32)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py
new file mode 100644
index 0000000000000000000000000000000000000000..976af2cb38560c0ec48938547eb8b1372ef4eb2e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios.py
@@ -0,0 +1,542 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+ HeliosAdditionalInputsStep,
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosPrepareHistoryStep,
+ HeliosSetTimestepsStep,
+ HeliosTextInputStep,
+ HeliosV2VSeedHistoryStep,
+)
+from .decoders import HeliosDecodeStep
+from .denoise import HeliosChunkDenoiseStep, HeliosI2VChunkDenoiseStep
+from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. Vae Encoder
+# ====================
+
+
+# auto_docstring
+class HeliosAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Encoder step that encodes video or image inputs. This is an auto pipeline block.
+ - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.
+ - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.
+ - If neither is provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ video_latents (`Tensor`):
+ Encoded video latents (chunked)
+ fake_image_latents (`Tensor`):
+ Fake image latents for history seeding
+ """
+
+ block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep]
+ block_names = ["video_encoder", "image_encoder"]
+ block_trigger_inputs = ["video", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Encoder step that encodes video or image inputs. This is an auto pipeline block.\n"
+ " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n"
+ " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n"
+ " - If neither is provided, step will be skipped."
+ )
+
+
+# ====================
+# 2. DENOISE
+# ====================
+
+
+# DENOISE (T2V)
+# auto_docstring
+class HeliosCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block that takes encoded conditions and runs the chunk-based denoising process.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosPrepareHistoryStep,
+ HeliosSetTimestepsStep,
+ HeliosChunkDenoiseStep,
+ ]
+ block_names = ["input", "prepare_history", "set_timesteps", "chunk_denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block that takes encoded conditions and runs the chunk-based denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (I2V)
+# auto_docstring
+class HeliosI2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ I2V denoise block that seeds history with image latents and uses I2V-aware chunk preparation.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video/fake-image latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video/fake-image latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "fake_image_latents",
+ type_hint=torch.Tensor,
+ description="Fake image latents used as history seed for I2V generation.",
+ ),
+ ],
+ ),
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosSetTimestepsStep,
+ HeliosI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_image",
+ "prepare_history",
+ "seed_history",
+ "set_timesteps",
+ "chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "I2V denoise block that seeds history with image latents and uses I2V-aware chunk preparation."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (V2V)
+# auto_docstring
+class HeliosV2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ V2V denoise block that seeds history with video latents and uses I2V-aware chunk preparation.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation."
+ ),
+ ],
+ ),
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosV2VSeedHistoryStep,
+ HeliosSetTimestepsStep,
+ HeliosI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_video",
+ "prepare_history",
+ "seed_history",
+ "set_timesteps",
+ "chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "V2V denoise block that seeds history with video latents and uses I2V-aware chunk preparation."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# AUTO DENOISE
+# auto_docstring
+class HeliosAutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ """
+ Core denoise step that selects the appropriate denoising block.
+ - `HeliosV2VCoreDenoiseStep` (video2video) for video-to-video tasks.
+ - `HeliosI2VCoreDenoiseStep` (image2video) for image-to-video tasks.
+ - `HeliosCoreDenoiseStep` (text2video) for text-to-video tasks.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ block_classes = [HeliosV2VCoreDenoiseStep, HeliosI2VCoreDenoiseStep, HeliosCoreDenoiseStep]
+ block_names = ["video2video", "image2video", "text2video"]
+ block_trigger_inputs = ["video_latents", "fake_image_latents"]
+ default_block_name = "text2video"
+
+ def select_block(self, video_latents=None, fake_image_latents=None):
+ if video_latents is not None:
+ return "video2video"
+ elif fake_image_latents is not None:
+ return "image2video"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Core denoise step that selects the appropriate denoising block.\n"
+ " - `HeliosV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n"
+ " - `HeliosI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n"
+ " - `HeliosCoreDenoiseStep` (text2video) for text-to-video tasks."
+ )
+
+
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", HeliosTextEncoderStep()),
+ ("vae_encoder", HeliosAutoVaeEncoderStep()),
+ ("denoise", HeliosAutoCoreDenoiseStep()),
+ ("decode", HeliosDecodeStep()),
+ ]
+)
+
+# ====================
+# 3. Auto Blocks
+# ====================
+
+
+# auto_docstring
+class HeliosAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-video, image-to-video, and video-to-video tasks using Helios.
+
+ Supported workflows:
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `prompt`, `image`
+ - `video2video`: requires `prompt`, `video`
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler
+ (`HeliosScheduler`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "helios"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2video": {"prompt": True},
+ "image2video": {"prompt": True, "image": True},
+ "video2video": {"prompt": True, "video": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-video, image-to-video, and video-to-video tasks using Helios."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py
new file mode 100644
index 0000000000000000000000000000000000000000..c14f159132984d89fa92ba2f5ad2c2758fed0bff
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py
@@ -0,0 +1,520 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+ HeliosAdditionalInputsStep,
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosPrepareHistoryStep,
+ HeliosTextInputStep,
+ HeliosV2VSeedHistoryStep,
+)
+from .decoders import HeliosDecodeStep
+from .denoise import HeliosPyramidChunkDenoiseStep, HeliosPyramidI2VChunkDenoiseStep
+from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. Vae Encoder
+# ====================
+
+
+# auto_docstring
+class HeliosPyramidAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Encoder step that encodes video or image inputs. This is an auto pipeline block.
+ - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.
+ - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.
+ - If neither is provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ video_latents (`Tensor`):
+ Encoded video latents (chunked)
+ fake_image_latents (`Tensor`):
+ Fake image latents for history seeding
+ """
+
+ block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep]
+ block_names = ["video_encoder", "image_encoder"]
+ block_trigger_inputs = ["video", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Encoder step that encodes video or image inputs. This is an auto pipeline block.\n"
+ " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n"
+ " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n"
+ " - If neither is provided, step will be skipped."
+ )
+
+
+# ====================
+# 2. DENOISE
+# ====================
+
+
+# DENOISE (T2V)
+# auto_docstring
+class HeliosPyramidCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ T2V pyramid denoise block with progressive multi-resolution denoising.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider
+ (`ClassifierFreeZeroStarGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosPrepareHistoryStep,
+ HeliosPyramidChunkDenoiseStep,
+ ]
+ block_names = ["input", "prepare_history", "pyramid_chunk_denoise"]
+
+ @property
+ def description(self):
+ return "T2V pyramid denoise block with progressive multi-resolution denoising."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (I2V)
+# auto_docstring
+class HeliosPyramidI2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ I2V pyramid denoise block with progressive multi-resolution denoising.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider
+ (`ClassifierFreeZeroStarGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video/fake-image latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video/fake-image latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "fake_image_latents",
+ type_hint=torch.Tensor,
+ description="Fake image latents used as history seed for I2V generation.",
+ ),
+ ],
+ ),
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosPyramidI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_image",
+ "prepare_history",
+ "seed_history",
+ "pyramid_chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "I2V pyramid denoise block with progressive multi-resolution denoising."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (V2V)
+# auto_docstring
+class HeliosPyramidV2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ V2V pyramid denoise block with progressive multi-resolution denoising.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider
+ (`ClassifierFreeZeroStarGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation."
+ ),
+ ],
+ ),
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosV2VSeedHistoryStep,
+ HeliosPyramidI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_video",
+ "prepare_history",
+ "seed_history",
+ "pyramid_chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "V2V pyramid denoise block with progressive multi-resolution denoising."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# AUTO DENOISE
+# auto_docstring
+class HeliosPyramidAutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ """
+ Pyramid core denoise step that selects the appropriate denoising block.
+ - `HeliosPyramidV2VCoreDenoiseStep` (video2video) for video-to-video tasks.
+ - `HeliosPyramidI2VCoreDenoiseStep` (image2video) for image-to-video tasks.
+ - `HeliosPyramidCoreDenoiseStep` (text2video) for text-to-video tasks.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider
+ (`ClassifierFreeZeroStarGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ block_classes = [HeliosPyramidV2VCoreDenoiseStep, HeliosPyramidI2VCoreDenoiseStep, HeliosPyramidCoreDenoiseStep]
+ block_names = ["video2video", "image2video", "text2video"]
+ block_trigger_inputs = ["video_latents", "fake_image_latents"]
+ default_block_name = "text2video"
+
+ def select_block(self, video_latents=None, fake_image_latents=None):
+ if video_latents is not None:
+ return "video2video"
+ elif fake_image_latents is not None:
+ return "image2video"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Pyramid core denoise step that selects the appropriate denoising block.\n"
+ " - `HeliosPyramidV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n"
+ " - `HeliosPyramidI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n"
+ " - `HeliosPyramidCoreDenoiseStep` (text2video) for text-to-video tasks."
+ )
+
+
+# ====================
+# 3. Auto Blocks
+# ====================
+
+PYRAMID_AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", HeliosTextEncoderStep()),
+ ("vae_encoder", HeliosPyramidAutoVaeEncoderStep()),
+ ("denoise", HeliosPyramidAutoCoreDenoiseStep()),
+ ("decode", HeliosDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class HeliosPyramidAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for pyramid progressive generation (T2V/I2V/V2V) using Helios.
+
+ Supported workflows:
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `prompt`, `image`
+ - `video2video`: requires `prompt`, `video`
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler
+ (`HeliosScheduler`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "helios-pyramid"
+
+ block_classes = PYRAMID_AUTO_BLOCKS.values()
+ block_names = PYRAMID_AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2video": {"prompt": True},
+ "image2video": {"prompt": True, "image": True},
+ "video2video": {"prompt": True, "video": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for pyramid progressive generation (T2V/I2V/V2V) using Helios."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py
new file mode 100644
index 0000000000000000000000000000000000000000..171cbc2b318d63be73f069acc0c994be9955625d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py
@@ -0,0 +1,530 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+ HeliosAdditionalInputsStep,
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosPrepareHistoryStep,
+ HeliosTextInputStep,
+ HeliosV2VSeedHistoryStep,
+)
+from .decoders import HeliosDecodeStep
+from .denoise import HeliosPyramidDistilledChunkDenoiseStep, HeliosPyramidDistilledI2VChunkDenoiseStep
+from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. Vae Encoder
+# ====================
+
+
+# auto_docstring
+class HeliosPyramidDistilledAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Encoder step for distilled pyramid pipeline.
+ - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.
+ - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.
+ - If neither is provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ video_latents (`Tensor`):
+ Encoded video latents (chunked)
+ fake_image_latents (`Tensor`):
+ Fake image latents for history seeding
+ """
+
+ block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep]
+ block_names = ["video_encoder", "image_encoder"]
+ block_trigger_inputs = ["video", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Encoder step for distilled pyramid pipeline.\n"
+ " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n"
+ " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n"
+ " - If neither is provided, step will be skipped."
+ )
+
+
+# ====================
+# 2. DENOISE
+# ====================
+
+
+# DENOISE (T2V)
+# auto_docstring
+class HeliosPyramidDistilledCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ T2V distilled pyramid denoise block with DMD scheduler and no CFG.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ is_amplify_first_chunk (`bool`, *optional*, defaults to True):
+ Whether to double the first chunk's timesteps via the scheduler for amplified generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosPrepareHistoryStep,
+ HeliosPyramidDistilledChunkDenoiseStep,
+ ]
+ block_names = ["input", "prepare_history", "pyramid_chunk_denoise"]
+
+ @property
+ def description(self):
+ return "T2V distilled pyramid denoise block with DMD scheduler and no CFG."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (I2V)
+# auto_docstring
+class HeliosPyramidDistilledI2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ I2V distilled pyramid denoise block with DMD scheduler and no CFG.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video/fake-image latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video/fake-image latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ is_amplify_first_chunk (`bool`, *optional*, defaults to True):
+ Whether to double the first chunk's timesteps via the scheduler for amplified generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "fake_image_latents",
+ type_hint=torch.Tensor,
+ description="Fake image latents used as history seed for I2V generation.",
+ ),
+ ],
+ ),
+ HeliosAddNoiseToImageLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosI2VSeedHistoryStep,
+ HeliosPyramidDistilledI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_image",
+ "prepare_history",
+ "seed_history",
+ "pyramid_chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "I2V distilled pyramid denoise block with DMD scheduler and no CFG."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# DENOISE (V2V)
+# auto_docstring
+class HeliosPyramidDistilledV2VCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ V2V distilled pyramid denoise block with DMD scheduler and no CFG.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`, *optional*, defaults to [16, 2, 1]):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ is_amplify_first_chunk (`bool`, *optional*, defaults to True):
+ Whether to double the first chunk's timesteps via the scheduler for amplified generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ model_name = "helios-pyramid"
+ block_classes = [
+ HeliosTextInputStep,
+ HeliosAdditionalInputsStep(
+ image_latent_inputs=[InputParam.template("image_latents")],
+ additional_batch_inputs=[
+ InputParam(
+ "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation."
+ ),
+ ],
+ ),
+ HeliosAddNoiseToVideoLatentsStep,
+ HeliosPrepareHistoryStep,
+ HeliosV2VSeedHistoryStep,
+ HeliosPyramidDistilledI2VChunkDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "add_noise_video",
+ "prepare_history",
+ "seed_history",
+ "pyramid_chunk_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "V2V distilled pyramid denoise block with DMD scheduler and no CFG."
+
+ @property
+ def outputs(self):
+ return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")]
+
+
+# AUTO DENOISE
+# auto_docstring
+class HeliosPyramidDistilledAutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ """
+ Distilled pyramid core denoise step that selects the appropriate denoising block.
+ - `HeliosPyramidDistilledV2VCoreDenoiseStep` (video2video) for video-to-video tasks.
+ - `HeliosPyramidDistilledI2VCoreDenoiseStep` (image2video) for image-to-video tasks.
+ - `HeliosPyramidDistilledCoreDenoiseStep` (text2video) for text-to-video tasks.
+
+ Components:
+ transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ is_amplify_first_chunk (`bool`, *optional*, defaults to True):
+ Whether to double the first chunk's timesteps via the scheduler for amplified generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+
+ Outputs:
+ latent_chunks (`list`):
+ List of per-chunk denoised latent tensors
+ """
+
+ block_classes = [
+ HeliosPyramidDistilledV2VCoreDenoiseStep,
+ HeliosPyramidDistilledI2VCoreDenoiseStep,
+ HeliosPyramidDistilledCoreDenoiseStep,
+ ]
+ block_names = ["video2video", "image2video", "text2video"]
+ block_trigger_inputs = ["video_latents", "fake_image_latents"]
+ default_block_name = "text2video"
+
+ def select_block(self, video_latents=None, fake_image_latents=None):
+ if video_latents is not None:
+ return "video2video"
+ elif fake_image_latents is not None:
+ return "image2video"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Distilled pyramid core denoise step that selects the appropriate denoising block.\n"
+ " - `HeliosPyramidDistilledV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n"
+ " - `HeliosPyramidDistilledI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n"
+ " - `HeliosPyramidDistilledCoreDenoiseStep` (text2video) for text-to-video tasks."
+ )
+
+
+# ====================
+# 3. Auto Blocks
+# ====================
+
+DISTILLED_PYRAMID_AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", HeliosTextEncoderStep()),
+ ("vae_encoder", HeliosPyramidDistilledAutoVaeEncoderStep()),
+ ("denoise", HeliosPyramidDistilledAutoCoreDenoiseStep()),
+ ("decode", HeliosDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class HeliosPyramidDistilledAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for distilled pyramid progressive generation (T2V/I2V/V2V) using Helios.
+
+ Supported workflows:
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `prompt`, `image`
+ - `video2video`: requires `prompt`, `video`
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler
+ (`HeliosScheduler`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ video (`None`, *optional*):
+ Input video for video-to-video generation
+ height (`int`, *optional*, defaults to 384):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 640):
+ The width in pixels of the generated image.
+ num_latent_frames_per_chunk (`int`, *optional*, defaults to 9):
+ Number of latent frames per temporal chunk.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos to generate per prompt.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ video_latents (`Tensor`, *optional*):
+ Encoded video latents for V2V generation.
+ image_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for image latent noise.
+ image_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for image latent noise.
+ video_noise_sigma_min (`float`, *optional*, defaults to 0.111):
+ Minimum sigma for video latent noise.
+ video_noise_sigma_max (`float`, *optional*, defaults to 0.135):
+ Maximum sigma for video latent noise.
+ num_frames (`int`, *optional*, defaults to 132):
+ Total number of video frames to generate.
+ history_sizes (`list`):
+ Sizes of long/mid/short history buffers for temporal context.
+ keep_first_frame (`bool`, *optional*, defaults to True):
+ Whether to keep the first frame as a prefix in history.
+ pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]):
+ Number of denoising steps per pyramid stage.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ is_amplify_first_chunk (`bool`, *optional*, defaults to True):
+ Whether to double the first chunk's timesteps via the scheduler for amplified generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ fake_image_latents (`Tensor`, *optional*):
+ Fake image latents used as history seed for I2V generation.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "helios-pyramid"
+
+ block_classes = DISTILLED_PYRAMID_AUTO_BLOCKS.values()
+ block_names = DISTILLED_PYRAMID_AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2video": {"prompt": True},
+ "image2video": {"prompt": True, "image": True},
+ "video2video": {"prompt": True, "video": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for distilled pyramid progressive generation (T2V/I2V/V2V) using Helios."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..d30e49d78f764f2088fc33b7ff4d72a03a5ea127
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/helios/modular_pipeline.py
@@ -0,0 +1,87 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...loaders import HeliosLoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class HeliosModularPipeline(
+ ModularPipeline,
+ HeliosLoraLoaderMixin,
+):
+ """
+ A ModularPipeline for Helios text-to-video generation.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "HeliosAutoBlocks"
+
+ @property
+ def vae_scale_factor_spatial(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = self.vae.config.scale_factor_spatial
+ return vae_scale_factor
+
+ @property
+ def vae_scale_factor_temporal(self):
+ vae_scale_factor = 4
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = self.vae.config.scale_factor_temporal
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ # YiYi TODO: find out default value
+ num_channels_latents = 16
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels
+ return num_channels_latents
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
+
+
+class HeliosPyramidModularPipeline(HeliosModularPipeline):
+ """
+ A ModularPipeline for Helios pyramid (progressive resolution) video generation.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "HeliosPyramidAutoBlocks"
+
+
+class HeliosPyramidDistilledModularPipeline(HeliosModularPipeline):
+ """
+ A ModularPipeline for Helios distilled pyramid video generation using DMD scheduler.
+
+ Uses guidance_scale=1.0 (no CFG) and supports is_amplify_first_chunk for the DMD scheduler.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "HeliosPyramidDistilledAutoBlocks"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..caefc5256495ae7e9ab85ee70cfc8ca89f1b9715
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py
@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_hunyuan_video1_5"] = [
+ "HunyuanVideo15AutoBlocks",
+ ]
+ _import_structure["modular_pipeline"] = ["HunyuanVideo15ModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_hunyuan_video1_5 import HunyuanVideo15AutoBlocks
+ from .modular_pipeline import HunyuanVideo15ModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e10b4d34aabfba52ea24db8315fa4c03046b2ec
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py
@@ -0,0 +1,324 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import HunyuanVideo15Transformer3DModel
+from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import HunyuanVideo15ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class HunyuanVideo15TextInputStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Input processing step that determines batch_size"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt_embeds"),
+ InputParam.template("batch_size", default=None),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("batch_size", type_hint=int),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.batch_size = getattr(block_state, "batch_size", None) or block_state.prompt_embeds.shape[0]
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HunyuanVideo15SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor),
+ OutputParam("num_inference_steps", type_hint=int),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ sigmas = block_state.sigmas
+ if sigmas is None:
+ sigmas = np.linspace(1.0, 0.0, block_state.num_inference_steps + 1)[:-1]
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler, block_state.num_inference_steps, device, sigmas=sigmas
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HunyuanVideo15PrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents, conditioning latents, mask, and image_embeds for T2V"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", HunyuanVideo15Transformer3DModel),
+ ComponentSpec(
+ "video_processor",
+ HunyuanVideo15ImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam("num_frames", type_hint=int, default=121, description="Number of video frames to generate."),
+ InputParam.template("latents"),
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ InputParam.template("generator"),
+ InputParam.template("batch_size", required=True, default=None),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor, description="Pure noise latents"),
+ OutputParam("cond_latents_concat", type_hint=torch.Tensor),
+ OutputParam("mask_concat", type_hint=torch.Tensor),
+ OutputParam("image_embeds", type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ height = block_state.height
+ width = block_state.width
+ if height is None and width is None:
+ height, width = components.video_processor.calculate_default_height_width(
+ components.default_aspect_ratio[1], components.default_aspect_ratio[0], components.target_size
+ )
+
+ batch_size = block_state.batch_size * block_state.num_videos_per_prompt
+ num_frames = block_state.num_frames
+
+ latents = block_state.latents
+ if latents is not None:
+ latents = latents.to(device=device, dtype=dtype)
+ else:
+ shape = (
+ batch_size,
+ components.num_channels_latents,
+ (num_frames - 1) // components.vae_scale_factor_temporal + 1,
+ int(height) // components.vae_scale_factor_spatial,
+ int(width) // components.vae_scale_factor_spatial,
+ )
+ if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+ latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=dtype)
+
+ block_state.latents = latents
+
+ b, c, f, h, w = latents.shape
+ block_state.cond_latents_concat = torch.zeros(b, c, f, h, w, dtype=dtype, device=device)
+ block_state.mask_concat = torch.zeros(b, 1, f, h, w, dtype=dtype, device=device)
+
+ block_state.image_embeds = torch.zeros(
+ block_state.batch_size,
+ components.vision_num_semantic_tokens,
+ components.vision_states_dim,
+ dtype=dtype,
+ device=device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HunyuanVideo15Image2VideoPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepare I2V conditioning from image_latents and image_embeds. "
+ "Expects pure noise `latents` from HunyuanVideo15PrepareLatentsStep. "
+ "Builds cond_latents_concat and mask_concat for the denoiser."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", HunyuanVideo15Transformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.",
+ ),
+ InputParam(
+ "image_embeds",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Siglip image embeddings from the image encoder step, used as extra conditioning for I2V.",
+ ),
+ InputParam.template("latents", required=True),
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ InputParam.template("batch_size", required=True, default=None),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("cond_latents_concat", type_hint=torch.Tensor),
+ OutputParam("mask_concat", type_hint=torch.Tensor),
+ OutputParam("image_embeds", type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ batch_size = block_state.batch_size * block_state.num_videos_per_prompt
+
+ b, c, f, h, w = block_state.latents.shape
+
+ latent_condition = block_state.image_latents.to(device=device, dtype=dtype)
+ latent_condition = latent_condition.repeat(batch_size, 1, f, 1, 1)
+ latent_condition[:, :, 1:, :, :] = 0
+ block_state.cond_latents_concat = latent_condition
+
+ latent_mask = torch.zeros(b, 1, f, h, w, dtype=dtype, device=device)
+ latent_mask[:, :, 0, :, :] = 1.0
+ block_state.mask_concat = latent_mask
+
+ image_embeds = block_state.image_embeds.to(device=device, dtype=dtype)
+ if image_embeds.shape[0] == 1 and batch_size > 1:
+ image_embeds = image_embeds.repeat(batch_size, 1, 1)
+ block_state.image_embeds = image_embeds
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..e347a6b50b596b07c35e3a54dcb0f7a05cefdd17
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py
@@ -0,0 +1,70 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLHunyuanVideo15
+from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)
+
+
+class HunyuanVideo15VaeDecoderStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLHunyuanVideo15),
+ ComponentSpec(
+ "video_processor",
+ HunyuanVideo15ImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into videos"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True),
+ InputParam.template("output_type", default="np"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("videos"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latents = block_state.latents.to(components.vae.dtype) / components.vae.config.scaling_factor
+ video = components.vae.decode(latents, return_dict=False)[0]
+ block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d6335043fd15ef1345b320151a2f32e42a702ba
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py
@@ -0,0 +1,401 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import HunyuanVideo15Transformer3DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam
+from .modular_pipeline import HunyuanVideo15ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class HunyuanVideo15LoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that prepares the latent input"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True),
+ InputParam("cond_latents_concat", required=True, type_hint=torch.Tensor),
+ InputParam("mask_concat", required=True, type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = torch.cat(
+ [block_state.latents, block_state.cond_latents_concat, block_state.mask_concat], dim=1
+ )
+ return components, block_state
+
+
+class HunyuanVideo15LoopDenoiser(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ def __init__(self, guider_input_fields=None):
+ if guider_input_fields is None:
+ guider_input_fields = {
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"),
+ "encoder_hidden_states_2": ("prompt_embeds_2", "negative_prompt_embeds_2"),
+ "encoder_attention_mask_2": ("prompt_embeds_mask_2", "negative_prompt_embeds_mask_2"),
+ }
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", HunyuanVideo15Transformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that denoises the latents with guidance"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("num_inference_steps", required=True, default=None),
+ InputParam(
+ "image_embeds",
+ type_hint=torch.Tensor,
+ description="Siglip image embeddings used as extra conditioning for I2V. Zero-filled for T2V.",
+ ),
+ ]
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ inputs.append(
+ InputParam(
+ name=value[0],
+ required=True,
+ type_hint=torch.Tensor,
+ description=f"Positive branch of the {value[0]!r} field fed into the guider.",
+ )
+ )
+ for neg_name in value[1:]:
+ inputs.append(
+ InputParam(
+ name=neg_name,
+ type_hint=torch.Tensor,
+ description=f"Negative branch of the {neg_name!r} field fed into the guider.",
+ )
+ )
+ else:
+ inputs.append(
+ InputParam(
+ name=value,
+ required=True,
+ type_hint=torch.Tensor,
+ description=f"{value!r} field fed into the guider.",
+ )
+ )
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ timestep = t.expand(block_state.latent_model_input.shape[0]).to(block_state.latent_model_input.dtype)
+
+ # Step 1: Collect model inputs
+ guider_inputs = {
+ input_name: tuple(getattr(block_state, v) for v in value)
+ if isinstance(value, tuple)
+ else getattr(block_state, value)
+ for input_name, value in self._guider_input_fields.items()
+ }
+
+ # Step 2: Update guider state
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # Step 3: Prepare batched inputs
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ # Step 4: Run denoiser for each batch
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ image_embeds=block_state.image_embeds,
+ timestep=timestep,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+
+ components.guider.cleanup_models(components.transformer)
+
+ # Step 5: Combine predictions
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class HunyuanVideo15LoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that updates the latents"
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred, t, block_state.latents, return_dict=False
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class HunyuanVideo15DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Pipeline block that iteratively denoises the latents over timesteps"
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", HunyuanVideo15Transformer3DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("timesteps", required=True),
+ InputParam.template("num_inference_steps", required=True, default=None),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HunyuanVideo15DenoiseStep(HunyuanVideo15DenoiseLoopWrapper):
+ block_classes = [
+ HunyuanVideo15LoopBeforeDenoiser,
+ HunyuanVideo15LoopDenoiser(),
+ HunyuanVideo15LoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents.\n"
+ "At each iteration:\n"
+ " - `HunyuanVideo15LoopBeforeDenoiser`\n"
+ " - `HunyuanVideo15LoopDenoiser`\n"
+ " - `HunyuanVideo15LoopAfterDenoiser`\n"
+ "This block supports text-to-video tasks."
+ )
+
+
+class HunyuanVideo15Image2VideoLoopDenoiser(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ def __init__(self, guider_input_fields=None):
+ if guider_input_fields is None:
+ guider_input_fields = {
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"),
+ "encoder_hidden_states_2": ("prompt_embeds_2", "negative_prompt_embeds_2"),
+ "encoder_attention_mask_2": ("prompt_embeds_mask_2", "negative_prompt_embeds_mask_2"),
+ }
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", HunyuanVideo15Transformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "I2V denoiser with MeanFlow timestep_r support"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("num_inference_steps", required=True, default=None),
+ InputParam(
+ "image_embeds",
+ type_hint=torch.Tensor,
+ description="Siglip image embeddings used as extra conditioning for I2V. Zero-filled for T2V.",
+ ),
+ InputParam.template("timesteps", required=True),
+ ]
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ inputs.append(
+ InputParam(
+ name=value[0],
+ required=True,
+ type_hint=torch.Tensor,
+ description=f"Positive branch of the {value[0]!r} field fed into the guider.",
+ )
+ )
+ for neg_name in value[1:]:
+ inputs.append(
+ InputParam(
+ name=neg_name,
+ type_hint=torch.Tensor,
+ description=f"Negative branch of the {neg_name!r} field fed into the guider.",
+ )
+ )
+ else:
+ inputs.append(
+ InputParam(
+ name=value,
+ required=True,
+ type_hint=torch.Tensor,
+ description=f"{value!r} field fed into the guider.",
+ )
+ )
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ timestep = t.expand(block_state.latent_model_input.shape[0]).to(block_state.latent_model_input.dtype)
+
+ # MeanFlow timestep_r (lines 855-862)
+ if components.transformer.config.use_meanflow:
+ if i == len(block_state.timesteps) - 1:
+ timestep_r = torch.tensor([0.0], device=timestep.device)
+ else:
+ timestep_r = block_state.timesteps[i + 1]
+ timestep_r = timestep_r.expand(block_state.latents.shape[0]).to(block_state.latents.dtype)
+ else:
+ timestep_r = None
+
+ guider_inputs = {
+ input_name: tuple(getattr(block_state, v) for v in value)
+ if isinstance(value, tuple)
+ else getattr(block_state, value)
+ for input_name, value in self._guider_input_fields.items()
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ image_embeds=block_state.image_embeds,
+ timestep=timestep,
+ timestep_r=timestep_r,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class HunyuanVideo15Image2VideoDenoiseStep(HunyuanVideo15DenoiseLoopWrapper):
+ block_classes = [
+ HunyuanVideo15LoopBeforeDenoiser,
+ HunyuanVideo15Image2VideoLoopDenoiser(),
+ HunyuanVideo15LoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step for image-to-video with MeanFlow support.\n"
+ "At each iteration:\n"
+ " - `HunyuanVideo15LoopBeforeDenoiser`\n"
+ " - `HunyuanVideo15Image2VideoLoopDenoiser`\n"
+ " - `HunyuanVideo15LoopAfterDenoiser`"
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..a984bfdb94096b34297a1928114d9bb8e7ee4e13
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py
@@ -0,0 +1,441 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+import torch
+from transformers import (
+ ByT5Tokenizer,
+ Qwen2_5_VLTextModel,
+ Qwen2TokenizerFast,
+ SiglipImageProcessor,
+ SiglipVisionModel,
+ T5EncoderModel,
+)
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import AutoencoderKLHunyuanVideo15
+from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import HunyuanVideo15ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+def format_text_input(prompt, system_message):
+ return [
+ [{"role": "system", "content": system_message}, {"role": "user", "content": p if p else " "}] for p in prompt
+ ]
+
+
+def extract_glyph_texts(prompt):
+ pattern = r"\"(.*?)\"|\"(.*?)\""
+ matches = re.findall(pattern, prompt)
+ result = [match[0] or match[1] for match in matches]
+ result = list(dict.fromkeys(result)) if len(result) > 1 else result
+ if result:
+ formatted_result = ". ".join([f'Text "{text}"' for text in result]) + ". "
+ else:
+ formatted_result = None
+ return formatted_result
+
+
+def _get_mllm_prompt_embeds(
+ text_encoder,
+ tokenizer,
+ prompt,
+ device,
+ tokenizer_max_length=1000,
+ num_hidden_layers_to_skip=2,
+ # fmt: off
+ system_message="You are a helpful assistant. Describe the video by detailing the following aspects: \
+ 1. The main content and theme of the video. \
+ 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \
+ 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \
+ 4. background environment, light, style and atmosphere. \
+ 5. camera angles, movements, and transitions used in the video.",
+ # fmt: on
+ crop_start=108,
+):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ prompt = format_text_input(prompt, system_message)
+
+ text_inputs = tokenizer.apply_chat_template(
+ prompt,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=True,
+ padding="max_length",
+ max_length=tokenizer_max_length + crop_start,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids.to(device=device)
+ prompt_attention_mask = text_inputs.attention_mask.to(device=device)
+
+ prompt_embeds = text_encoder(
+ input_ids=text_input_ids,
+ attention_mask=prompt_attention_mask,
+ output_hidden_states=True,
+ ).hidden_states[-(num_hidden_layers_to_skip + 1)]
+
+ if crop_start is not None and crop_start > 0:
+ prompt_embeds = prompt_embeds[:, crop_start:]
+ prompt_attention_mask = prompt_attention_mask[:, crop_start:]
+
+ return prompt_embeds, prompt_attention_mask
+
+
+def _get_byt5_prompt_embeds(tokenizer, text_encoder, prompt, device, tokenizer_max_length=256):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ glyph_texts = [extract_glyph_texts(p) for p in prompt]
+
+ prompt_embeds_list = []
+ prompt_embeds_mask_list = []
+
+ for glyph_text in glyph_texts:
+ if glyph_text is None:
+ glyph_text_embeds = torch.zeros(
+ (1, tokenizer_max_length, text_encoder.config.d_model), device=device, dtype=text_encoder.dtype
+ )
+ glyph_text_embeds_mask = torch.zeros((1, tokenizer_max_length), device=device, dtype=torch.int64)
+ else:
+ txt_tokens = tokenizer(
+ glyph_text,
+ padding="max_length",
+ max_length=tokenizer_max_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ ).to(device)
+
+ glyph_text_embeds = text_encoder(
+ input_ids=txt_tokens.input_ids,
+ attention_mask=txt_tokens.attention_mask.float(),
+ )[0]
+ glyph_text_embeds = glyph_text_embeds.to(device=device)
+ glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device)
+
+ prompt_embeds_list.append(glyph_text_embeds)
+ prompt_embeds_mask_list.append(glyph_text_embeds_mask)
+
+ return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0)
+
+
+class HunyuanVideo15TextEncoderStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Dual text encoder step using Qwen2.5-VL (MLLM) and ByT5 (glyph text)"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen2_5_VLTextModel),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ComponentSpec("text_encoder_2", T5EncoderModel),
+ ComponentSpec("tokenizer_2", ByT5Tokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", required=False),
+ InputParam.template("negative_prompt"),
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("prompt_embeds_mask"),
+ OutputParam.template("negative_prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds_mask"),
+ OutputParam(
+ "prompt_embeds_2",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="ByT5 glyph-text embeddings used as a second conditioning stream for the transformer.",
+ ),
+ OutputParam(
+ "prompt_embeds_mask_2",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Attention mask for the ByT5 glyph-text embeddings.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds_2",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="ByT5 glyph-text negative embeddings for classifier-free guidance.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds_mask_2",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Attention mask for the ByT5 glyph-text negative embeddings.",
+ ),
+ ]
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt,
+ device=None,
+ dtype=None,
+ batch_size=1,
+ num_videos_per_prompt=1,
+ ):
+ device = device or components._execution_device
+ dtype = dtype or components.text_encoder.dtype
+
+ if prompt is None:
+ prompt = [""] * batch_size
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ prompt_embeds, prompt_embeds_mask = _get_mllm_prompt_embeds(
+ tokenizer=components.tokenizer,
+ text_encoder=components.text_encoder,
+ prompt=prompt,
+ device=device,
+ tokenizer_max_length=components.tokenizer_max_length,
+ system_message=components.system_message,
+ crop_start=components.prompt_template_encode_start_idx,
+ )
+
+ prompt_embeds_2, prompt_embeds_mask_2 = _get_byt5_prompt_embeds(
+ tokenizer=components.tokenizer_2,
+ text_encoder=components.text_encoder_2,
+ prompt=prompt,
+ device=device,
+ tokenizer_max_length=components.tokenizer_2_max_length,
+ )
+
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1).view(
+ batch_size * num_videos_per_prompt, seq_len, -1
+ )
+ prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_videos_per_prompt, 1).view(
+ batch_size * num_videos_per_prompt, seq_len
+ )
+
+ _, seq_len_2, _ = prompt_embeds_2.shape
+ prompt_embeds_2 = prompt_embeds_2.repeat(1, num_videos_per_prompt, 1).view(
+ batch_size * num_videos_per_prompt, seq_len_2, -1
+ )
+ prompt_embeds_mask_2 = prompt_embeds_mask_2.repeat(1, num_videos_per_prompt, 1).view(
+ batch_size * num_videos_per_prompt, seq_len_2
+ )
+
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+ prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype, device=device)
+ prompt_embeds_2 = prompt_embeds_2.to(dtype=dtype, device=device)
+ prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype, device=device)
+
+ return prompt_embeds, prompt_embeds_mask, prompt_embeds_2, prompt_embeds_mask_2
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ dtype = components.transformer.dtype
+
+ prompt = block_state.prompt
+ negative_prompt = block_state.negative_prompt
+ num_videos_per_prompt = block_state.num_videos_per_prompt
+
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = 1
+
+ (
+ block_state.prompt_embeds,
+ block_state.prompt_embeds_mask,
+ block_state.prompt_embeds_2,
+ block_state.prompt_embeds_mask_2,
+ ) = self.encode_prompt(
+ components,
+ prompt=prompt,
+ device=device,
+ dtype=dtype,
+ batch_size=batch_size,
+ num_videos_per_prompt=num_videos_per_prompt,
+ )
+
+ if components.requires_unconditional_embeds:
+ (
+ block_state.negative_prompt_embeds,
+ block_state.negative_prompt_embeds_mask,
+ block_state.negative_prompt_embeds_2,
+ block_state.negative_prompt_embeds_mask_2,
+ ) = self.encode_prompt(
+ components,
+ prompt=negative_prompt,
+ device=device,
+ dtype=dtype,
+ batch_size=batch_size,
+ num_videos_per_prompt=num_videos_per_prompt,
+ )
+
+ state.set("batch_size", batch_size)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class HunyuanVideo15VaeEncoderStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "VAE Encoder step that encodes an input image into latent space for image-to-video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLHunyuanVideo15),
+ ComponentSpec(
+ "video_processor",
+ HunyuanVideo15ImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image", required=True),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Encoded image latents from the VAE encoder",
+ ),
+ OutputParam("height", type_hint=int, description="Target height resolved from image"),
+ OutputParam("width", type_hint=int, description="Target width resolved from image"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ image = block_state.image
+ height = block_state.height
+ width = block_state.width
+ if height is None or width is None:
+ height, width = components.video_processor.calculate_default_height_width(
+ height=image.size[1], width=image.size[0], target_size=components.target_size
+ )
+ image = components.video_processor.resize(image, height=height, width=width, resize_mode="crop")
+
+ vae_dtype = components.vae.dtype
+ image_tensor = components.video_processor.preprocess(image, height=height, width=width).to(
+ device=device, dtype=vae_dtype
+ )
+ image_tensor = image_tensor.unsqueeze(2)
+ image_latents = retrieve_latents(components.vae.encode(image_tensor), sample_mode="argmax")
+ image_latents = image_latents * components.vae.config.scaling_factor
+
+ block_state.image_latents = image_latents
+ block_state.height = height
+ block_state.width = width
+ state.set("image", image)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class HunyuanVideo15ImageEncoderStep(ModularPipelineBlocks):
+ model_name = "hunyuan-video-1.5"
+
+ @property
+ def description(self) -> str:
+ return "Siglip image encoder step that produces image_embeds for image-to-video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("image_encoder", SiglipVisionModel),
+ ComponentSpec("feature_extractor", SiglipImageProcessor),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_embeds",
+ type_hint=torch.Tensor,
+ description="Image embeddings from the Siglip vision encoder",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ image_encoder_dtype = next(components.image_encoder.parameters()).dtype
+ image_inputs = components.feature_extractor.preprocess(
+ images=block_state.image, do_resize=True, return_tensors="pt", do_convert_rgb=True
+ )
+ image_inputs = image_inputs.to(device=device, dtype=image_encoder_dtype)
+ image_embeds = components.image_encoder(**image_inputs).last_hidden_state
+
+ block_state.image_embeds = image_embeds
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b39c5344e98644e965eca5eaae93e27ac241e98
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py
@@ -0,0 +1,535 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ HunyuanVideo15Image2VideoPrepareLatentsStep,
+ HunyuanVideo15PrepareLatentsStep,
+ HunyuanVideo15SetTimestepsStep,
+ HunyuanVideo15TextInputStep,
+)
+from .decoders import HunyuanVideo15VaeDecoderStep
+from .denoise import HunyuanVideo15DenoiseStep, HunyuanVideo15Image2VideoDenoiseStep
+from .encoders import (
+ HunyuanVideo15ImageEncoderStep,
+ HunyuanVideo15TextEncoderStep,
+ HunyuanVideo15VaeEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# auto_docstring
+class HunyuanVideo15CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`)
+ video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ batch_size (`int`, *optional*):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds' field fed into the guider.
+ prompt_embeds_mask (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask' field fed into the guider.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider.
+ prompt_embeds_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_2' field fed into the guider.
+ negative_prompt_embeds_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_2' field fed into the guider.
+ prompt_embeds_mask_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask_2' field fed into the guider.
+ negative_prompt_embeds_mask_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [
+ HunyuanVideo15TextInputStep,
+ HunyuanVideo15SetTimestepsStep,
+ HunyuanVideo15PrepareLatentsStep,
+ HunyuanVideo15DenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class HunyuanVideo15Blocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for HunyuanVideo 1.5 text-to-video.
+
+ Components:
+ text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`)
+ tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ transformer (`HunyuanVideo15Transformer3DModel`) video_processor (`HunyuanVideo15ImageProcessor`) vae
+ (`AutoencoderKLHunyuanVideo15`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ batch_size (`int`, *optional*):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [
+ HunyuanVideo15TextEncoderStep,
+ HunyuanVideo15CoreDenoiseStep,
+ HunyuanVideo15VaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "denoise", "decode"]
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for HunyuanVideo 1.5 text-to-video."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
+
+
+# auto_docstring
+class HunyuanVideo15Image2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block for image-to-video that takes encoded conditions and runs the denoising process.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`)
+ video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ batch_size (`int`, *optional*):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`):
+ Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds' field fed into the guider.
+ prompt_embeds_mask (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask' field fed into the guider.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider.
+ prompt_embeds_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_2' field fed into the guider.
+ negative_prompt_embeds_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_2' field fed into the guider.
+ prompt_embeds_mask_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask_2' field fed into the guider.
+ negative_prompt_embeds_mask_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [
+ HunyuanVideo15TextInputStep,
+ HunyuanVideo15SetTimestepsStep,
+ HunyuanVideo15PrepareLatentsStep,
+ HunyuanVideo15Image2VideoPrepareLatentsStep,
+ HunyuanVideo15Image2VideoDenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block for image-to-video that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class HunyuanVideo15AutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image input into its latent representation.
+ This is an auto pipeline block that works for image-to-video tasks.
+ - `HunyuanVideo15VaeEncoderStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLHunyuanVideo15`) video_processor (`HunyuanVideo15ImageProcessor`)
+
+ Inputs:
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+
+ Outputs:
+ image_latents (`Tensor`):
+ Encoded image latents from the VAE encoder
+ height (`int`):
+ Target height resolved from image
+ width (`int`):
+ Target width resolved from image
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [HunyuanVideo15VaeEncoderStep]
+ block_names = ["vae_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image input into its latent representation.\n"
+ "This is an auto pipeline block that works for image-to-video tasks.\n"
+ " - `HunyuanVideo15VaeEncoderStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+# auto_docstring
+class HunyuanVideo15AutoImageEncoderStep(AutoPipelineBlocks):
+ """
+ Siglip image encoder step that produces image_embeds.
+ This is an auto pipeline block that works for image-to-video tasks.
+ - `HunyuanVideo15ImageEncoderStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ image_encoder (`SiglipVisionModel`) feature_extractor (`SiglipImageProcessor`)
+
+ Inputs:
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ image_embeds (`Tensor`):
+ Image embeddings from the Siglip vision encoder
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [HunyuanVideo15ImageEncoderStep]
+ block_names = ["image_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "Siglip image encoder step that produces image_embeds.\n"
+ "This is an auto pipeline block that works for image-to-video tasks.\n"
+ " - `HunyuanVideo15ImageEncoderStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+# auto_docstring
+class HunyuanVideo15AutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto denoise block that selects the appropriate denoise pipeline based on inputs.
+ - `HunyuanVideo15Image2VideoCoreDenoiseStep` is used when `image_latents` is provided.
+ - `HunyuanVideo15CoreDenoiseStep` is used otherwise (text-to-video).
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`)
+ video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`, *optional*):
+ Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds' field fed into the guider.
+ prompt_embeds_mask (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask' field fed into the guider.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider.
+ prompt_embeds_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_2' field fed into the guider.
+ negative_prompt_embeds_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_2' field fed into the guider.
+ prompt_embeds_mask_2 (`Tensor`):
+ Positive branch of the 'prompt_embeds_mask_2' field fed into the guider.
+ negative_prompt_embeds_mask_2 (`Tensor`, *optional*):
+ Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [HunyuanVideo15Image2VideoCoreDenoiseStep, HunyuanVideo15CoreDenoiseStep]
+ block_names = ["image2video", "text2video"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto denoise block that selects the appropriate denoise pipeline based on inputs.\n"
+ " - `HunyuanVideo15Image2VideoCoreDenoiseStep` is used when `image_latents` is provided.\n"
+ " - `HunyuanVideo15CoreDenoiseStep` is used otherwise (text-to-video)."
+ )
+
+
+# auto_docstring
+class HunyuanVideo15AutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks for HunyuanVideo 1.5 that support both text-to-video and image-to-video workflows.
+
+ Supported workflows:
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`)
+ tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo15`)
+ video_processor (`HunyuanVideo15ImageProcessor`) image_encoder (`SiglipVisionModel`) feature_extractor
+ (`SiglipImageProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
+ (`HunyuanVideo15Transformer3DModel`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`, *optional*):
+ Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [
+ HunyuanVideo15TextEncoderStep,
+ HunyuanVideo15AutoVaeEncoderStep,
+ HunyuanVideo15AutoImageEncoderStep,
+ HunyuanVideo15AutoCoreDenoiseStep,
+ HunyuanVideo15VaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "image_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2video": {"prompt": True},
+ "image2video": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto blocks for HunyuanVideo 1.5 that support both text-to-video and image-to-video workflows."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
+
+
+# auto_docstring
+class HunyuanVideo15Image2VideoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for HunyuanVideo 1.5 image-to-video.
+
+ Components:
+ text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`)
+ tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo15`)
+ video_processor (`HunyuanVideo15ImageProcessor`) image_encoder (`SiglipVisionModel`) feature_extractor
+ (`SiglipImageProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
+ (`HunyuanVideo15Transformer3DModel`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ batch_size (`int`, *optional*):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ num_frames (`int`, *optional*, defaults to 121):
+ Number of video frames to generate.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`):
+ Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "hunyuan-video-1.5"
+ block_classes = [
+ HunyuanVideo15TextEncoderStep,
+ HunyuanVideo15AutoVaeEncoderStep,
+ HunyuanVideo15AutoImageEncoderStep,
+ HunyuanVideo15Image2VideoCoreDenoiseStep,
+ HunyuanVideo15VaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "image_encoder", "denoise", "decode"]
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for HunyuanVideo 1.5 image-to-video."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..61cbc709a164fbc8ecc39407a9a024cf480f2fc5
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py
@@ -0,0 +1,90 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...loaders import HunyuanVideoLoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class HunyuanVideo15ModularPipeline(
+ ModularPipeline,
+ HunyuanVideoLoraLoaderMixin,
+):
+ """
+ A ModularPipeline for HunyuanVideo 1.5.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "HunyuanVideo15AutoBlocks"
+
+ @property
+ def vae_scale_factor_spatial(self):
+ return self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16
+
+ @property
+ def vae_scale_factor_temporal(self):
+ return self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4
+
+ @property
+ def num_channels_latents(self):
+ return self.vae.config.latent_channels if getattr(self, "vae", None) else 32
+
+ @property
+ def target_size(self):
+ return self.transformer.config.target_size if getattr(self, "transformer", None) else 640
+
+ @property
+ def default_aspect_ratio(self):
+ return (16, 9)
+
+ @property
+ def vision_num_semantic_tokens(self):
+ return 729
+
+ @property
+ def vision_states_dim(self):
+ return self.transformer.config.image_embed_dim if getattr(self, "transformer", None) else 1152
+
+ @property
+ def tokenizer_max_length(self):
+ return 1000
+
+ @property
+ def tokenizer_2_max_length(self):
+ return 256
+
+ # fmt: off
+ @property
+ def system_message(self):
+ return "You are a helpful assistant. Describe the video by detailing the following aspects: \
+ 1. The main content and theme of the video. \
+ 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \
+ 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \
+ 4. background environment, light, style and atmosphere. \
+ 5. camera angles, movements, and transitions used in the video."
+ # fmt: on
+
+ @property
+ def prompt_template_encode_start_idx(self):
+ return 108
+
+ @property
+ def requires_unconditional_embeds(self):
+ if hasattr(self, "guider") and self.guider is not None:
+ return self.guider._enabled and self.guider.num_conditions > 1
+ return False
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d52f92167c99318b3356f8586e4ad53d0087c6ae
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_ideogram4"] = ["Ideogram4AutoBlocks"]
+ _import_structure["modular_pipeline"] = ["Ideogram4ModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_ideogram4 import Ideogram4AutoBlocks
+ from .modular_pipeline import Ideogram4ModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..32e7d46bcd2884be5d3774a4b2f3955e23f36a0b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/before_denoise.py
@@ -0,0 +1,558 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+
+import torch
+
+from ...models.transformers.transformer_ideogram4 import (
+ IMAGE_POSITION_OFFSET,
+ LLM_TOKEN_INDICATOR,
+ OUTPUT_IMAGE_INDICATOR,
+ SEQUENCE_PADDING_INDICATOR,
+ Ideogram4Transformer2DModel,
+)
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Ideogram4ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+# Default per-step guidance schedule (length must equal `num_inference_steps`): 7.0 for the main steps,
+# dropping to 3.0 for the final 3 "polish" steps.
+DEFAULT_GUIDANCE_SCHEDULE = (7.0,) * 45 + (3.0,) * 3
+
+
+# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._logit_normal_sigmas
+def _logit_normal_sigmas(
+ num_inference_steps: int,
+ mu: float,
+ std: float = 1.0,
+ logsnr_min: float = -15.0,
+ logsnr_max: float = 18.0,
+ device: torch.device | None = None,
+) -> torch.Tensor:
+ r"""
+ Build a length-`num_inference_steps` sigma schedule using the Ideogram4 logit-normal flow-matching schedule.
+
+ Sigmas are returned in `[0, 1]` in decreasing order (sigma close to 1 corresponds to pure noise, sigma close to 0
+ to clean data), matching diffusers conventions.
+
+ The Ideogram4 schedule applies `sigma(s) = 1 - logit_normal_cdf_inverse(1 - s)` to `s = linspace(0, 1, N + 1)` and
+ keeps the first `N` entries; a terminal zero is appended downstream by the scheduler.
+ """
+ intervals = torch.linspace(0.0, 1.0, num_inference_steps + 1, dtype=torch.float64)
+ # Apply the inverse CDF of a normal then push through the logistic to obtain a logit-normal CDF inverse.
+ z = torch.special.ndtri(intervals)
+ y = mu + std * z
+ t = 1.0 - torch.special.expit(y)
+ t_min = 1.0 / (1.0 + math.exp(0.5 * logsnr_max))
+ t_max = 1.0 / (1.0 + math.exp(0.5 * logsnr_min))
+ t = t.clamp(t_min, t_max)
+ # Convert from model time (0 = noise, 1 = data) to diffusers sigma (1 = noise, 0 = data) and reverse.
+ sigmas = (1.0 - t).flip(0)
+ # Drop the trailing 0; FlowMatchEulerDiscreteScheduler.set_timesteps appends one back internally.
+ sigmas = sigmas[:-1].to(dtype=torch.float32, device=device)
+ return sigmas
+
+
+# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._resolution_aware_mu
+def _resolution_aware_mu(
+ height: int,
+ width: int,
+ base_mu: float,
+ base_resolution: tuple[int, int] = (512, 512),
+) -> float:
+ """Shift the schedule mean as a function of image resolution."""
+ num_pixels = height * width
+ base_pixels = base_resolution[0] * base_resolution[1]
+ return base_mu + 0.5 * math.log(num_pixels / base_pixels)
+
+
+# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._expand_tensor_to_effective_batch
+def _expand_tensor_to_effective_batch(
+ tensor: torch.Tensor,
+ batch_size: int,
+ num_per_prompt: int,
+ tensor_name: str | None = None,
+) -> torch.Tensor:
+ """Replicate `tensor` along dim 0 from `batch_size` (or 1) to `batch_size * num_per_prompt`."""
+ target_batch_size = batch_size * num_per_prompt
+
+ if tensor.shape[0] == target_batch_size:
+ return tensor
+
+ if tensor.shape[0] == 1:
+ repeat_by = target_batch_size
+ elif tensor.shape[0] == batch_size:
+ repeat_by = num_per_prompt
+ else:
+ tensor_name = f"`{tensor_name}`" if tensor_name is not None else "Tensor"
+ raise ValueError(
+ f"{tensor_name} batch size must be 1, `batch_size` ({batch_size}), or "
+ f"`batch_size * num_*_per_prompt` ({target_batch_size}), but got {tensor.shape[0]}."
+ )
+
+ return torch.repeat_interleave(tensor, repeats=repeat_by, dim=0, output_size=tensor.shape[0] * repeat_by)
+
+
+# auto_docstring
+class Ideogram4TextInputsStep(ModularPipelineBlocks):
+ """
+ Input step that determines `batch_size`/`dtype` from the per-prompt `text_features` and replicates the text outputs
+ to `batch_size * num_images_per_prompt`. Place after the text encoder.
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ text_features (`Tensor`):
+ Per-prompt text features from the encoder.
+ text_lengths (`list`):
+ Per-prompt text-token counts from the encoder.
+
+ Outputs:
+ batch_size (`int`):
+ Effective batch size (num prompts * num_images_per_prompt).
+ dtype (`dtype`):
+ The dtype of the text features.
+ text_features (`Tensor`):
+ Text features, batch-expanded.
+ text_lengths (`list`):
+ Text-token counts, batch-expanded.
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input step that determines `batch_size`/`dtype` from the per-prompt `text_features` and replicates the "
+ "text outputs to `batch_size * num_images_per_prompt`. Place after the text encoder."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt", default=1),
+ InputParam(
+ name="text_features",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt text features from the encoder.",
+ ),
+ InputParam(
+ name="text_lengths",
+ required=True,
+ type_hint=list,
+ description="Per-prompt text-token counts from the encoder.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="batch_size",
+ type_hint=int,
+ description="Effective batch size (num prompts * num_images_per_prompt).",
+ ),
+ OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."),
+ OutputParam(name="text_features", type_hint=torch.Tensor, description="Text features, batch-expanded."),
+ OutputParam(name="text_lengths", type_hint=list, description="Text-token counts, batch-expanded."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_batch = block_state.text_features.shape[0]
+ num_per_prompt = block_state.num_images_per_prompt
+
+ block_state.dtype = block_state.text_features.dtype
+ block_state.text_features = _expand_tensor_to_effective_batch(
+ block_state.text_features, prompt_batch, num_per_prompt, "text_features"
+ )
+ block_state.text_lengths = [n for n in block_state.text_lengths for _ in range(num_per_prompt)]
+ block_state.batch_size = prompt_batch * num_per_prompt
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Ideogram4PrepareLatentsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the packed image latents (B, num_image_tokens, latent_dim) for the denoising loop.
+
+ Components:
+ transformer (`Ideogram4Transformer2DModel`)
+
+ Inputs:
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ batch_size (`int`):
+ Effective batch size.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial packed image latents (B, num_image_tokens, latent_dim).
+ num_image_tokens (`int`):
+ Number of image tokens (grid_h * grid_w).
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the packed image latents (B, num_image_tokens, latent_dim) for the denoising loop."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Ideogram4Transformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("generator"),
+ InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="The initial packed image latents (B, num_image_tokens, latent_dim).",
+ ),
+ OutputParam(
+ name="num_image_tokens", type_hint=int, description="Number of image tokens (grid_h * grid_w)."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ patch = components.patch_size
+ grid_h = block_state.height // (components.vae_scale_factor * patch)
+ grid_w = block_state.width // (components.vae_scale_factor * patch)
+ num_image_tokens = grid_h * grid_w
+ latent_dim = components.transformer.config.in_channels
+
+ shape = (block_state.batch_size, num_image_tokens, latent_dim)
+ if block_state.latents is None:
+ block_state.latents = randn_tensor(
+ shape, generator=block_state.generator, device=device, dtype=torch.float32
+ )
+ else:
+ block_state.latents = block_state.latents.to(device=device, dtype=torch.float32)
+
+ block_state.num_image_tokens = num_image_tokens
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Ideogram4SetTimestepsStep(ModularPipelineBlocks):
+ """
+ Step that sets the resolution-aware logit-normal sigma schedule on the scheduler and resolves the per-step guidance
+ weights.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 48):
+ The number of denoising steps.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ mu (`float`, *optional*, defaults to 0.0):
+ Base mean of the logit-normal schedule.
+ std (`float`, *optional*, defaults to 1.5):
+ Std of the logit-normal schedule.
+ guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)):
+ Per-step guidance scale schedule (length num_inference_steps).
+
+ Outputs:
+ timesteps (`Tensor`):
+ The denoising timesteps.
+ gw (`Tensor`):
+ Per-step guidance weights (num_inference_steps,).
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that sets the resolution-aware logit-normal sigma schedule on the scheduler and resolves the "
+ "per-step guidance weights."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", default=48),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam(name="mu", default=0.0, type_hint=float, description="Base mean of the logit-normal schedule."),
+ InputParam(name="std", default=1.5, type_hint=float, description="Std of the logit-normal schedule."),
+ InputParam(
+ name="guidance_schedule",
+ default=DEFAULT_GUIDANCE_SCHEDULE,
+ type_hint=list,
+ description="Per-step guidance scale schedule (length num_inference_steps).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps."),
+ OutputParam(
+ name="gw", type_hint=torch.Tensor, description="Per-step guidance weights (num_inference_steps,)."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ if len(block_state.guidance_schedule) != block_state.num_inference_steps:
+ raise ValueError(
+ f"`guidance_schedule` must have length `num_inference_steps` ({block_state.num_inference_steps}), "
+ f"got {len(block_state.guidance_schedule)}."
+ )
+
+ schedule_mu = _resolution_aware_mu(height=block_state.height, width=block_state.width, base_mu=block_state.mu)
+ sigmas = _logit_normal_sigmas(block_state.num_inference_steps, schedule_mu, std=block_state.std, device=device)
+ components.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device)
+
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.gw = torch.as_tensor(block_state.guidance_schedule, dtype=torch.float32, device=device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Ideogram4PrepareAdditionalInputsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the additional denoiser inputs from the packed-sequence layout: the conditional
+ encoder_hidden_states (text features packed with image padding) and the position_ids/segment_ids/indicator, plus
+ the unconditional (image-only) counterparts. Place after prepare_latents.
+
+ Inputs:
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ text_features (`Tensor`):
+ Batch-expanded text features.
+ text_lengths (`list`):
+ Batch-expanded text-token counts.
+ batch_size (`int`):
+ Effective batch size.
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ Packed conditional encoder_hidden_states (B, total_seq, dim).
+ position_ids (`Tensor`):
+ Conditional 3-axis MRoPE position ids.
+ segment_ids (`Tensor`):
+ Conditional block-diagonal segment ids.
+ indicator (`Tensor`):
+ Conditional per-token text/image/pad role.
+ negative_prompt_embeds (`Tensor`):
+ Unconditional (zeroed) text features (B, num_image_tokens, dim).
+ negative_position_ids (`Tensor`):
+ Unconditional position ids (image region).
+ negative_segment_ids (`Tensor`):
+ Unconditional segment ids (image region).
+ negative_indicator (`Tensor`):
+ Unconditional indicator (image region).
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that prepares the additional denoiser inputs from the packed-sequence layout: the conditional "
+ "encoder_hidden_states (text features packed with image padding) and the position_ids/segment_ids/"
+ "indicator, plus the unconditional (image-only) counterparts. Place after prepare_latents."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam(
+ name="text_features",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Batch-expanded text features.",
+ ),
+ InputParam(
+ name="text_lengths", required=True, type_hint=list, description="Batch-expanded text-token counts."
+ ),
+ InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Packed conditional encoder_hidden_states (B, total_seq, dim).",
+ ),
+ OutputParam(
+ name="position_ids", type_hint=torch.Tensor, description="Conditional 3-axis MRoPE position ids."
+ ),
+ OutputParam(
+ name="segment_ids", type_hint=torch.Tensor, description="Conditional block-diagonal segment ids."
+ ),
+ OutputParam(
+ name="indicator", type_hint=torch.Tensor, description="Conditional per-token text/image/pad role."
+ ),
+ OutputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Unconditional (zeroed) text features (B, num_image_tokens, dim).",
+ ),
+ OutputParam(
+ name="negative_position_ids",
+ type_hint=torch.Tensor,
+ description="Unconditional position ids (image region).",
+ ),
+ OutputParam(
+ name="negative_segment_ids",
+ type_hint=torch.Tensor,
+ description="Unconditional segment ids (image region).",
+ ),
+ OutputParam(
+ name="negative_indicator",
+ type_hint=torch.Tensor,
+ description="Unconditional indicator (image region).",
+ ),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4.Ideogram4Pipeline._prepare_ids
+ def _prepare_ids(
+ text_lengths: list[int],
+ grid_h: int,
+ grid_w: int,
+ max_text_tokens: int,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Build the packed `[left-pad][text][image]` layout from the per-prompt text lengths and the image grid.
+
+ Returns `position_ids` (3-axis MRoPE), `segment_ids` (block-diagonal attention) and `indicator` (per-token
+ text/image/pad role).
+ """
+ batch_size = len(text_lengths)
+ num_image_tokens = grid_h * grid_w
+ total_seq_len = max_text_tokens + num_image_tokens
+
+ # Image position ids (t=0, h, w); offset keeps them disjoint from text positions.
+ h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
+ w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
+ t_idx = torch.zeros_like(h_idx)
+ image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
+
+ position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
+ segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
+ indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
+
+ for b, num_text in enumerate(text_lengths):
+ offset = max_text_tokens - num_text
+
+ text_pos = torch.arange(num_text)
+ text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1)
+ position_ids[b, offset : offset + num_text] = text_pos_3d
+ position_ids[b, offset + num_text :] = image_pos
+
+ indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR
+ indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR
+
+ segment_ids[b, offset : offset + num_text + num_image_tokens] = 1
+
+ return position_ids.to(device), segment_ids.to(device), indicator.to(device)
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ patch = components.patch_size
+ grid_h = block_state.height // (components.vae_scale_factor * patch)
+ grid_w = block_state.width // (components.vae_scale_factor * patch)
+ num_image_tokens = grid_h * grid_w
+
+ text_features = block_state.text_features
+ max_text_tokens = text_features.shape[1]
+ feature_dim = text_features.shape[-1]
+
+ position_ids, segment_ids, indicator = self._prepare_ids(
+ block_state.text_lengths, grid_h, grid_w, max_text_tokens, device
+ )
+
+ # Pack the text features into the full sequence; image positions carry no text features.
+ image_feature_padding = torch.zeros(
+ block_state.batch_size, num_image_tokens, feature_dim, dtype=text_features.dtype, device=device
+ )
+ block_state.prompt_embeds = torch.cat([text_features, image_feature_padding], dim=1)
+
+ # Unconditional (image-only) branch, derived from the conditioning.
+ block_state.negative_prompt_embeds = torch.zeros(
+ block_state.batch_size, num_image_tokens, feature_dim, dtype=text_features.dtype, device=device
+ )
+ block_state.position_ids = position_ids
+ block_state.segment_ids = segment_ids
+ block_state.indicator = indicator
+ block_state.negative_position_ids = position_ids[:, max_text_tokens:]
+ block_state.negative_segment_ids = segment_ids[:, max_text_tokens:]
+ block_state.negative_indicator = indicator[:, max_text_tokens:]
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab0977a18addc5f6737121dbad22e5f4ff39d7b6
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/decoders.py
@@ -0,0 +1,112 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKLFlux2
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Ideogram4ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# auto_docstring
+class Ideogram4DecodeStep(ModularPipelineBlocks):
+ """
+ Step that decodes the unpatchified (B, ae_channels, H, W) latents into images: de-normalizes with the VAE
+ batch-norm statistics and decodes through the VAE.
+
+ Components:
+ vae (`AutoencoderKLFlux2`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ latents (`Tensor`):
+ The unpatchified (B, ae_channels, H, W) latents to decode, from the after-denoise step.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that decodes the unpatchified (B, ae_channels, H, W) latents into images: de-normalizes with the "
+ "VAE batch-norm statistics and decodes through the VAE."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLFlux2),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("output_type", default="pil"),
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The unpatchified (B, ae_channels, H, W) latents to decode, from the after-denoise step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ z = block_state.latents
+ patch = components.patch_size
+ ae_channels = z.shape[1]
+ grid_h, grid_w = z.shape[2] // patch, z.shape[3] // patch
+
+ # VAE bn stores per-channel statistics over the packed channels, laid out as (patch_row, patch_col,
+ # ae_channel). Reshape them into an (ae_channels, patch, patch) tile and repeat across the grid so the
+ # denormalization on the unpatchified latents matches the packed-space statistics.
+ bn_mean = components.vae.bn.running_mean.view(patch, patch, ae_channels).permute(2, 0, 1)
+ bn_std = torch.sqrt(components.vae.bn.running_var + components.vae.config.batch_norm_eps)
+ bn_std = bn_std.view(patch, patch, ae_channels).permute(2, 0, 1)
+ bn_mean = bn_mean.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype)
+ bn_std = bn_std.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype)
+ z = z * bn_std + bn_mean
+
+ decoded = components.vae.decode(z.to(components.vae.dtype), return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(
+ decoded.float(), output_type=block_state.output_type
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd9852b93b7ada9e9c74911d49098eb928730875
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/denoise.py
@@ -0,0 +1,363 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...models.transformers.transformer_ideogram4 import Ideogram4Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Ideogram4ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class Ideogram4LoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Within the denoising loop: build the conditional packed input `[text-padding][image latents]` and the "
+ "model timestep. Compose into the `sub_blocks` of `Ideogram4DenoiseLoopWrapper`."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+ InputParam(
+ name="position_ids", required=True, type_hint=torch.Tensor, description="Conditional position ids."
+ ),
+ InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # Conditional packed sequence is [text-padding][image latents]; text region length = total - image tokens.
+ max_text_tokens = block_state.position_ids.shape[1] - block_state.latents.shape[1]
+ text_z_padding = torch.zeros(
+ block_state.latents.shape[0],
+ max_text_tokens,
+ block_state.latents.shape[-1],
+ dtype=block_state.latents.dtype,
+ device=block_state.latents.device,
+ )
+ block_state.pos_z = torch.cat([text_z_padding, block_state.latents], dim=1)
+ block_state.max_text_tokens = max_text_tokens
+
+ # Map sigma-domain timestep to model time t in [0, 1] (0 = noise, 1 = clean data).
+ num_train_timesteps = components.scheduler.config.num_train_timesteps
+ t_model = 1.0 - (t.float() / num_train_timesteps)
+ block_state.t_model = t_model.expand(block_state.batch_size)
+ return components, block_state
+
+
+class Ideogram4LoopDenoiser(ModularPipelineBlocks):
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Within the denoising loop: run the conditional `transformer` on the full packed sequence and the "
+ "`unconditional_transformer` on the image-only sequence, then blend with the per-step guidance weight "
+ "(asymmetric CFG, no guider). Compose into `Ideogram4DenoiseLoopWrapper`."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", Ideogram4Transformer2DModel),
+ ComponentSpec("unconditional_transformer", Ideogram4Transformer2DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Packed conditional encoder_hidden_states.",
+ ),
+ InputParam(
+ name="position_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditional 3-axis MRoPE position ids.",
+ ),
+ InputParam(
+ name="segment_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditional block-diagonal segment ids.",
+ ),
+ InputParam(
+ name="indicator",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditional per-token text/image/pad role.",
+ ),
+ InputParam(
+ name="negative_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Unconditional (zeroed) text features.",
+ ),
+ InputParam(
+ name="negative_position_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Unconditional position ids (image region).",
+ ),
+ InputParam(
+ name="negative_segment_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Unconditional segment ids (image region).",
+ ),
+ InputParam(
+ name="negative_indicator",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Unconditional indicator (image region).",
+ ),
+ InputParam(name="gw", required=True, type_hint=torch.Tensor, description="Per-step guidance weights."),
+ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ transformer = components.transformer
+ unconditional_transformer = components.unconditional_transformer
+
+ # Conditional pass operates on the full packed sequence; the velocity is the image-token region.
+ pos_out = transformer(
+ hidden_states=block_state.pos_z.to(transformer.dtype),
+ timestep=block_state.t_model.to(transformer.dtype),
+ encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype),
+ position_ids=block_state.position_ids,
+ segment_ids=block_state.segment_ids,
+ indicator=block_state.indicator,
+ return_dict=False,
+ )[0]
+ pos_v = pos_out[:, block_state.max_text_tokens :].to(torch.float32)
+
+ # Unconditional pass uses the image-only positions with zeroed text features.
+ neg_v = unconditional_transformer(
+ hidden_states=block_state.latents.to(unconditional_transformer.dtype),
+ timestep=block_state.t_model.to(unconditional_transformer.dtype),
+ encoder_hidden_states=block_state.negative_prompt_embeds.to(unconditional_transformer.dtype),
+ position_ids=block_state.negative_position_ids,
+ segment_ids=block_state.negative_segment_ids,
+ indicator=block_state.negative_indicator,
+ return_dict=False,
+ )[0].to(torch.float32)
+
+ gw_i = block_state.gw[i]
+ v = gw_i * pos_v + (1.0 - gw_i) * neg_v
+ # The scheduler integrates `-v` (Ideogram predicts velocity v = x0 - noise).
+ block_state.noise_pred = -v
+ return components, block_state
+
+
+class Ideogram4LoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return "Within the denoising loop: scheduler step. Compose into `Ideogram4DenoiseLoopWrapper`."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="latents", type_hint=torch.Tensor, description="The denoised latents.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred, t, block_state.latents, return_dict=False
+ )[0]
+ return components, block_state
+
+
+# auto_docstring
+class Ideogram4DenoiseStep(LoopSequentialPipelineBlocks):
+ """
+ Denoising loop that iteratively denoises the packed image latents over `timesteps`, running both the conditional
+ and unconditional transformers and blending with the per-step guidance schedule.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Ideogram4Transformer2DModel`)
+ unconditional_transformer (`Ideogram4Transformer2DModel`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ Denoising timesteps from set_timesteps.
+ num_inference_steps (`int`, *optional*, defaults to 48):
+ The number of denoising steps.
+ latents (`Tensor`):
+ Packed image latents.
+ position_ids (`Tensor`):
+ Conditional position ids.
+ batch_size (`int`):
+ Effective batch size.
+ prompt_embeds (`Tensor`):
+ Packed conditional encoder_hidden_states.
+ position_ids (`Tensor`):
+ Conditional 3-axis MRoPE position ids.
+ segment_ids (`Tensor`):
+ Conditional block-diagonal segment ids.
+ indicator (`Tensor`):
+ Conditional per-token text/image/pad role.
+ negative_prompt_embeds (`Tensor`):
+ Unconditional (zeroed) text features.
+ negative_position_ids (`Tensor`):
+ Unconditional position ids (image region).
+ negative_segment_ids (`Tensor`):
+ Unconditional segment ids (image region).
+ negative_indicator (`Tensor`):
+ Unconditional indicator (image region).
+ gw (`Tensor`):
+ Per-step guidance weights.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoised latents.
+ """
+
+ model_name = "ideogram4"
+ block_classes = [Ideogram4LoopBeforeDenoiser, Ideogram4LoopDenoiser, Ideogram4LoopAfterDenoiser]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoising loop that iteratively denoises the packed image latents over `timesteps`, running both the "
+ "conditional and unconditional transformers and blending with the per-step guidance schedule."
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Denoising timesteps from set_timesteps.",
+ ),
+ InputParam.template("num_inference_steps", default=48),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Ideogram4AfterDenoiseStep(ModularPipelineBlocks):
+ """
+ Step that runs after the denoising loop: unpatchifies the packed image latents (B, num_image_tokens, ae_channels *
+ patch ** 2) into a (B, ae_channels, H, W) latent for the decoder.
+
+ Inputs:
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ latents (`Tensor`):
+ The denoised packed image latents (B, num_image_tokens, latent_dim).
+
+ Outputs:
+ latents (`Tensor`):
+ Unpatchified latents (B, ae_channels, H, W) ready for the VAE decoder.
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that runs after the denoising loop: unpatchifies the packed image latents "
+ "(B, num_image_tokens, ae_channels * patch ** 2) into a (B, ae_channels, H, W) latent for the decoder."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised packed image latents (B, num_image_tokens, latent_dim).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="Unpatchified latents (B, ae_channels, H, W) ready for the VAE decoder.",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ z = block_state.latents
+ patch = components.patch_size
+ grid_h = block_state.height // (components.vae_scale_factor * patch)
+ grid_w = block_state.width // (components.vae_scale_factor * patch)
+
+ ae_channels = z.shape[-1] // (patch * patch)
+ z = z.view(z.shape[0], grid_h, grid_w, patch, patch, ae_channels)
+ z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
+ z = z.view(z.shape[0], ae_channels, grid_h * patch, grid_w * patch)
+
+ block_state.latents = z
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..41b4541214464740e29302792d3417cb57b97f13
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/encoders.py
@@ -0,0 +1,327 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from transformers import Qwen2Tokenizer, Qwen3VLModel
+from transformers.masking_utils import create_causal_mask
+
+from ...pipelines.ideogram4.prompt_enhancer import (
+ PROMPT_UPSAMPLE_TEMPERATURE,
+ Ideogram4PromptEnhancerHead,
+ build_caption_logits_processor,
+ build_prompt_enhancer,
+ generate_captions,
+)
+from ...utils import is_outlines_available, logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Ideogram4ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Hidden states of these Qwen3-VL decoder layers are concatenated to form the per-token
+# text conditioning consumed by the Ideogram4 transformer.
+QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+
+
+# auto_docstring
+class Ideogram4PromptUpsampleStep(ModularPipelineBlocks):
+ """
+ Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption when
+ `prompt_upsampling=True` (the format the model is trained on). Requires a generative `text_encoder` (a
+ `Qwen3VLForConditionalGeneration`); install `outlines` for schema-constrained captions.
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired
+ with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): LM head grafted onto the text
+ encoder for prompt upsampling.
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ prompt_upsampling (`bool`, *optional*, defaults to False):
+ If True, rewrite the prompt into Ideogram4's native JSON caption before encoding.
+ prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0):
+ Sampling temperature for prompt upsampling.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ max_sequence_length (`int`, *optional*, defaults to 2048):
+ Maximum sequence length for prompt encoding.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ prompt (`list`):
+ The (possibly upsampled) prompt forwarded to the text encoder.
+ """
+
+ model_name = "ideogram4"
+
+ def __init__(self):
+ # Built lazily on first upsample: the head-less encoder body + `prompt_enhancer_head`, combined.
+ self._prompt_enhancer = None
+ # Outlines logits processor for schema-constrained captions; built lazily on first upsample.
+ self._caption_logits_processor = None
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption when "
+ "`prompt_upsampling=True` (the format the model is trained on). Requires a generative `text_encoder` "
+ "(a `Qwen3VLForConditionalGeneration`); install `outlines` for schema-constrained captions."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."),
+ ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer paired with the text encoder."),
+ ComponentSpec(
+ "prompt_enhancer_head",
+ Ideogram4PromptEnhancerHead,
+ description="LM head grafted onto the text encoder for prompt upsampling.",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", required=True),
+ InputParam(
+ name="prompt_upsampling",
+ type_hint=bool,
+ default=False,
+ description="If True, rewrite the prompt into Ideogram4's native JSON caption before encoding.",
+ ),
+ InputParam(
+ name="prompt_upsampling_temperature",
+ type_hint=float,
+ default=PROMPT_UPSAMPLE_TEMPERATURE,
+ description="Sampling temperature for prompt upsampling.",
+ ),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("max_sequence_length", default=2048),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="prompt",
+ type_hint=list,
+ description="The (possibly upsampled) prompt forwarded to the text encoder.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ if block_state.prompt_upsampling:
+ if components.prompt_enhancer_head is None:
+ raise ValueError(
+ "Prompt upsampling requires the `prompt_enhancer_head` component, which is not loaded. Load an "
+ "`Ideogram4PromptEnhancerHead` and add it to the pipeline."
+ )
+ if self._prompt_enhancer is None:
+ self._prompt_enhancer = build_prompt_enhancer(components.text_encoder, components.prompt_enhancer_head)
+ if self._caption_logits_processor is None and is_outlines_available():
+ self._caption_logits_processor = build_caption_logits_processor(
+ self._prompt_enhancer, components.tokenizer
+ )
+ if self._caption_logits_processor is None:
+ logger.warning_once(
+ "`outlines` is not installed; prompt upsampling runs unconstrained and may not return "
+ "schema-valid JSON. Install with `pip install outlines` for structured captions."
+ )
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ block_state.prompt = generate_captions(
+ self._prompt_enhancer,
+ components.tokenizer,
+ self._caption_logits_processor,
+ block_state.prompt,
+ height,
+ width,
+ temperature=block_state.prompt_upsampling_temperature,
+ max_new_tokens=block_state.max_sequence_length,
+ generator=block_state.generator,
+ device=components._execution_device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Ideogram4TextEncoderStep(ModularPipelineBlocks):
+ """
+ Text encoder step that tokenizes the prompt(s) and runs the Qwen3-VL text encoder, returning the per-token text
+ features (concatenated from a fixed set of activation layers). Only the text tokens are encoded; the packed image
+ tokens are appended later (the encoder is causal with image after text, so they never affect the text features).
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired
+ with the text encoder.
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ max_sequence_length (`int`, *optional*, defaults to 2048):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ text_features (`Tensor`):
+ Per-prompt text features (B, max_sequence_length, llm_features_dim), padding zeroed.
+ text_lengths (`list`):
+ Per-prompt real text-token counts, used to lay out the packed sequence.
+ """
+
+ model_name = "ideogram4"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text encoder step that tokenizes the prompt(s) and runs the Qwen3-VL text encoder, returning the "
+ "per-token text features (concatenated from a fixed set of activation layers). Only the text tokens are "
+ "encoded; the packed image tokens are appended later (the encoder is causal with image after text, so "
+ "they never affect the text features)."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."),
+ ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer paired with the text encoder."),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", required=True),
+ InputParam.template("max_sequence_length", default=2048),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="text_features",
+ type_hint=torch.Tensor,
+ description="Per-prompt text features (B, max_sequence_length, llm_features_dim), padding zeroed.",
+ ),
+ OutputParam(
+ name="text_lengths",
+ type_hint=list,
+ description="Per-prompt real text-token counts, used to lay out the packed sequence.",
+ ),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4.Ideogram4Pipeline._get_text_encoder_hidden_states
+ def _get_text_encoder_hidden_states(
+ text_encoder,
+ token_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ pos_2d: torch.Tensor,
+ ) -> list[torch.Tensor]:
+ """Run the text encoder's decoder layers, returning the hidden states tapped at each activation layer."""
+
+ language_model = text_encoder.language_model
+
+ inputs_embeds = language_model.embed_tokens(token_ids)
+
+ position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
+ text_position_ids = position_ids_4d[0]
+ mrope_position_ids = position_ids_4d[1:]
+
+ causal_mask = create_causal_mask(
+ config=language_model.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=None,
+ position_ids=text_position_ids,
+ )
+ position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
+
+ tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
+ captured: dict[int, torch.Tensor] = {}
+ hidden_states = inputs_embeds
+ for layer_idx, decoder_layer in enumerate(language_model.layers):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=text_position_ids,
+ past_key_values=None,
+ position_embeddings=position_embeddings,
+ )
+ if layer_idx in tap_set:
+ captured[layer_idx] = hidden_states
+
+ return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
+
+ @torch.no_grad()
+ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ tokenizer = components.tokenizer
+ max_text_tokens = block_state.max_sequence_length
+
+ prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt)
+ batch_size = len(prompts)
+
+ # Tokenize each chat-formatted prompt and left-pad to `max_sequence_length`.
+ token_ids = torch.zeros(batch_size, max_text_tokens, dtype=torch.long)
+ attention_mask = torch.zeros(batch_size, max_text_tokens, dtype=torch.long)
+ text_position_ids = torch.zeros(batch_size, max_text_tokens, dtype=torch.long)
+ text_lengths = []
+ for b, text_prompt in enumerate(prompts):
+ messages = [{"role": "user", "content": [{"type": "text", "text": text_prompt}]}]
+ text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+ toks = tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+ n = int(toks.shape[0])
+ if n > max_text_tokens:
+ raise ValueError(f"prompt has {n} tokens, exceeds max_sequence_length={max_text_tokens}")
+ text_lengths.append(n)
+ offset = max_text_tokens - n
+ token_ids[b, offset:] = toks
+ attention_mask[b, offset:] = 1
+ text_position_ids[b, offset:] = torch.arange(n)
+
+ token_ids = token_ids.to(device)
+ attention_mask = attention_mask.to(device)
+ text_position_ids = text_position_ids.to(device)
+
+ # Run the text encoder, tapping the activation-layer hidden states, then concatenate them into per-token
+ # text features (padding zeroed).
+ selected = self._get_text_encoder_hidden_states(
+ components.text_encoder, token_ids, attention_mask, text_position_ids
+ )
+ text_features = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(batch_size, max_text_tokens, -1)
+ text_features = (text_features * attention_mask.to(text_features.dtype).unsqueeze(-1)).to(torch.float32)
+
+ block_state.text_features = text_features
+ block_state.text_lengths = text_lengths
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py
new file mode 100644
index 0000000000000000000000000000000000000000..7075108d6678161ab454a432daf3d87fcdac2dae
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py
@@ -0,0 +1,185 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Ideogram4PrepareAdditionalInputsStep,
+ Ideogram4PrepareLatentsStep,
+ Ideogram4SetTimestepsStep,
+ Ideogram4TextInputsStep,
+)
+from .decoders import Ideogram4DecodeStep
+from .denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep
+from .encoders import Ideogram4PromptUpsampleStep, Ideogram4TextEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Core denoise: consumes the per-prompt text features and produces the unpatchified latents
+# (batch/latents/timesteps/ids inputs -> denoising loop -> unpatchify).
+CORE_DENOISE_BLOCKS = InsertableDict(
+ [
+ ("input", Ideogram4TextInputsStep()),
+ ("prepare_latents", Ideogram4PrepareLatentsStep()),
+ ("set_timesteps", Ideogram4SetTimestepsStep()),
+ ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()),
+ ("denoise", Ideogram4DenoiseStep()),
+ ("after_denoise", Ideogram4AfterDenoiseStep()),
+ ]
+)
+
+
+# auto_docstring
+class Ideogram4CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed denoiser
+ inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional transformers, and
+ unpatchifies the result for the decoder.
+
+ Components:
+ transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ unconditional_transformer (`Ideogram4Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ text_features (`Tensor`):
+ Per-prompt text features from the encoder.
+ text_lengths (`list`):
+ Per-prompt text-token counts from the encoder.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 48):
+ The number of denoising steps.
+ mu (`float`, *optional*, defaults to 0.0):
+ Base mean of the logit-normal schedule.
+ std (`float`, *optional*, defaults to 1.5):
+ Std of the logit-normal schedule.
+ guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)):
+ Per-step guidance scale schedule (length num_inference_steps).
+
+ Outputs:
+ latents (`Tensor`):
+ Unpatchified (B, ae_channels, H, W) latents.
+ """
+
+ model_name = "ideogram4"
+ block_classes = list(CORE_DENOISE_BLOCKS.values())
+ block_names = list(CORE_DENOISE_BLOCKS.keys())
+
+ @property
+ def description(self) -> str:
+ return (
+ "Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed "
+ "denoiser inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional "
+ "transformers, and unpatchifies the result for the decoder."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ # The only meaningful product of the core step is the unpatchified latents; the batch/timesteps/packed-sequence
+ # inputs prepared along the way are consumed within the loop and are not updated by it.
+ return [OutputParam.template("latents", description="Unpatchified (B, ae_channels, H, W) latents.")]
+
+
+# auto_docstring
+class Ideogram4AutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> encode text ->
+ core denoise (asymmetric CFG over two transformers) -> decode.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired
+ with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): LM head grafted onto the text
+ encoder for prompt upsampling. transformer (`Ideogram4Transformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) unconditional_transformer (`Ideogram4Transformer2DModel`) vae
+ (`AutoencoderKLFlux2`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ prompt_upsampling (`bool`, *optional*, defaults to False):
+ If True, rewrite the prompt into Ideogram4's native JSON caption before encoding.
+ prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0):
+ Sampling temperature for prompt upsampling.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ max_sequence_length (`int`, *optional*, defaults to 2048):
+ Maximum sequence length for prompt encoding.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ num_inference_steps (`int`, *optional*, defaults to 48):
+ The number of denoising steps.
+ mu (`float`, *optional*, defaults to 0.0):
+ Base mean of the logit-normal schedule.
+ std (`float`, *optional*, defaults to 1.5):
+ Std of the logit-normal schedule.
+ guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0,
+ 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)):
+ Per-step guidance scale schedule (length num_inference_steps).
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "ideogram4"
+ block_classes = [
+ Ideogram4PromptUpsampleStep(),
+ Ideogram4TextEncoderStep(),
+ Ideogram4CoreDenoiseStep(),
+ Ideogram4DecodeStep(),
+ ]
+ block_names = ["prompt_upsample", "text_encoder", "denoise", "decode"]
+
+ # Workflow map declaring the trigger conditions for each supported workflow.
+ # `True` means the workflow triggers when the input is not None.
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return (
+ "Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> "
+ "encode text -> core denoise (asymmetric CFG over two transformers) -> decode."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..1fc566161b8102b31d2caddfb36e5d2550d358e3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ideogram4/modular_pipeline.py
@@ -0,0 +1,46 @@
+# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...loaders import Ideogram4LoraLoaderMixin
+from ..modular_pipeline import ModularPipeline
+
+
+class Ideogram4ModularPipeline(ModularPipeline, Ideogram4LoraLoaderMixin):
+ """
+ A ModularPipeline for Ideogram4.
+
+ > [!WARNING] > This is an experimental feature!
+ """
+
+ default_blocks_name = "Ideogram4AutoBlocks"
+
+ # Ideogram4 patchifies the VAE output by a factor of 2 before feeding the transformer.
+ @property
+ def patch_size(self):
+ return 2
+
+ @property
+ def default_height(self):
+ return 2048
+
+ @property
+ def default_width(self):
+ return 2048
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if getattr(self, "vae", None) is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee9c920619b3dc24aa91a32cc3a6a494c6d8f071
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/__init__.py
@@ -0,0 +1,49 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_krea2"] = ["Krea2AutoBlocks"]
+ _import_structure["modular_blocks_krea2_turbo"] = ["Krea2TurboAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["Krea2ModularPipeline", "Krea2TurboModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_krea2 import Krea2AutoBlocks
+ from .modular_blocks_krea2_turbo import Krea2TurboAutoBlocks
+ from .modular_pipeline import Krea2ModularPipeline, Krea2TurboModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..54f2b2f77018a6eb216af448b10a58bb716f8b3e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/before_denoise.py
@@ -0,0 +1,590 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import numpy as np
+import torch
+
+from ...models.transformers.transformer_krea2 import Krea2Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Krea2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Copied from diffusers.pipelines.krea2.pipeline_krea2.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# auto_docstring
+class Krea2TextInputsStep(ModularPipelineBlocks):
+ """
+ Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the text
+ conditioning (and the optional negative branch) to `batch_size * num_images_per_prompt`. Place after the text
+ encoder.
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Per-prompt negative text features.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Per-prompt negative text mask.
+
+ Outputs:
+ batch_size (`int`):
+ Effective batch size (num prompts * num_images_per_prompt).
+ dtype (`dtype`):
+ The dtype of the text features.
+ prompt_embeds (`Tensor`):
+ Text features, batch-expanded.
+ prompt_embeds_mask (`Tensor`):
+ Text mask, batch-expanded.
+ negative_prompt_embeds (`Tensor`):
+ Negative text features, batch-expanded.
+ negative_prompt_embeds_mask (`Tensor`):
+ Negative text mask, batch-expanded.
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the "
+ "text conditioning (and the optional negative branch) to `batch_size * num_images_per_prompt`. Place after "
+ "the text encoder."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt", default=1),
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).",
+ ),
+ InputParam(
+ name="prompt_embeds_mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt boolean text mask (B, text_seq_len).",
+ ),
+ InputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Per-prompt negative text features.",
+ ),
+ InputParam(
+ name="negative_prompt_embeds_mask",
+ type_hint=torch.Tensor,
+ description="Per-prompt negative text mask.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="batch_size",
+ type_hint=int,
+ description="Effective batch size (num prompts * num_images_per_prompt).",
+ ),
+ OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."),
+ OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."),
+ OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."),
+ OutputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Negative text features, batch-expanded.",
+ ),
+ OutputParam(
+ name="negative_prompt_embeds_mask",
+ type_hint=torch.Tensor,
+ description="Negative text mask, batch-expanded.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_batch, seq_len, num_layers, dim = block_state.prompt_embeds.shape
+ n = block_state.num_images_per_prompt
+
+ block_state.dtype = block_state.prompt_embeds.dtype
+ block_state.batch_size = prompt_batch * n
+
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, n, 1, 1).view(
+ prompt_batch * n, seq_len, num_layers, dim
+ )
+ block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, n).view(prompt_batch * n, seq_len)
+
+ if block_state.negative_prompt_embeds is not None:
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(1, n, 1, 1).view(
+ prompt_batch * n, seq_len, num_layers, dim
+ )
+ block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat(1, n).view(
+ prompt_batch * n, seq_len
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2TurboTextInputsStep(ModularPipelineBlocks):
+ """
+ Input step for the distilled Krea 2 turbo checkpoint that determines `batch_size`/`dtype` from the per-prompt
+ `prompt_embeds` and replicates the text conditioning to `batch_size * num_images_per_prompt`. The distilled
+ checkpoint runs without classifier-free guidance, so there is no negative branch. Place after the text encoder.
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+
+ Outputs:
+ batch_size (`int`):
+ Effective batch size (num prompts * num_images_per_prompt).
+ dtype (`dtype`):
+ The dtype of the text features.
+ prompt_embeds (`Tensor`):
+ Text features, batch-expanded.
+ prompt_embeds_mask (`Tensor`):
+ Text mask, batch-expanded.
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input step for the distilled Krea 2 turbo checkpoint that determines `batch_size`/`dtype` from the "
+ "per-prompt `prompt_embeds` and replicates the text conditioning to `batch_size * num_images_per_prompt`. "
+ "The distilled checkpoint runs without classifier-free guidance, so there is no negative branch. Place "
+ "after the text encoder."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt", default=1),
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).",
+ ),
+ InputParam(
+ name="prompt_embeds_mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Per-prompt boolean text mask (B, text_seq_len).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="batch_size",
+ type_hint=int,
+ description="Effective batch size (num prompts * num_images_per_prompt).",
+ ),
+ OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."),
+ OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."),
+ OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ prompt_batch, seq_len, num_layers, dim = block_state.prompt_embeds.shape
+ n = block_state.num_images_per_prompt
+
+ block_state.dtype = block_state.prompt_embeds.dtype
+ block_state.batch_size = prompt_batch * n
+
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, n, 1, 1).view(
+ prompt_batch * n, seq_len, num_layers, dim
+ )
+ block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, n).view(prompt_batch * n, seq_len)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2PrepareLatentsStep(ModularPipelineBlocks):
+ """
+ Step that samples the spatial image latents and patch-packs them into (B, image_seq_len, in_channels) for the
+ denoising loop.
+
+ Components:
+ transformer (`Krea2Transformer2DModel`)
+
+ Inputs:
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ batch_size (`int`):
+ Effective batch size.
+ dtype (`dtype`):
+ The working dtype.
+
+ Outputs:
+ latents (`Tensor`):
+ The initial packed image latents (B, image_seq_len, in_channels).
+ image_seq_len (`int`):
+ Number of image tokens (grid_h * grid_w).
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that samples the spatial image latents and patch-packs them into (B, image_seq_len, in_channels) "
+ "for the denoising loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Krea2Transformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam.template("height", default=1024),
+ InputParam.template("width", default=1024),
+ InputParam.template("generator"),
+ InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
+ InputParam(name="dtype", required=True, type_hint=torch.dtype, description="The working dtype."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="The initial packed image latents (B, image_seq_len, in_channels).",
+ ),
+ OutputParam(name="image_seq_len", type_hint=int, description="Number of image tokens (grid_h * grid_w)."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ p = components.patch_size
+ num_channels_latents = components.transformer.config.in_channels // (p**2)
+
+ multiple = components.vae_scale_factor * components.patch_size
+ if block_state.height % multiple != 0 or block_state.width % multiple != 0:
+ rounded_height = ((block_state.height + multiple - 1) // multiple) * multiple
+ rounded_width = ((block_state.width + multiple - 1) // multiple) * multiple
+ logger.warning(
+ f"`height` and `width` must be multiples of {multiple}; rounding up from {block_state.height}x{block_state.width} to"
+ f" {rounded_height}x{rounded_width}."
+ )
+ block_state.height, block_state.width = rounded_height, rounded_width
+
+ latent_height = block_state.height // components.vae_scale_factor
+ latent_width = block_state.width // components.vae_scale_factor
+
+ if block_state.latents is not None:
+ block_state.latents = block_state.latents.to(device=device, dtype=block_state.dtype)
+ else:
+ latents = randn_tensor(
+ (block_state.batch_size, num_channels_latents, latent_height, latent_width),
+ generator=block_state.generator,
+ device=device,
+ dtype=block_state.dtype,
+ )
+ latents = latents.view(
+ block_state.batch_size, num_channels_latents, latent_height // p, p, latent_width // p, p
+ )
+ latents = latents.permute(0, 2, 4, 1, 3, 5)
+ block_state.latents = latents.reshape(
+ block_state.batch_size, (latent_height // p) * (latent_width // p), num_channels_latents * p * p
+ )
+
+ block_state.image_seq_len = block_state.latents.shape[1]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2SetTimestepsStep(ModularPipelineBlocks):
+ """
+ Step that sets the Krea 2 flow-matching schedule on the scheduler: a linear sigma schedule with a resolution-aware
+ dynamic time shift `mu`.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 28):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+ image_seq_len (`int`):
+ Number of image tokens, used to compute the resolution-aware shift.
+
+ Outputs:
+ timesteps (`Tensor`):
+ The denoising timesteps.
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that sets the Krea 2 flow-matching schedule on the scheduler: a linear sigma schedule with a "
+ "resolution-aware dynamic time shift `mu`."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", default=28),
+ InputParam(
+ name="sigmas", type_hint=list, description="Custom sigma schedule (defaults to a linear ramp)."
+ ),
+ InputParam(
+ name="image_seq_len",
+ required=True,
+ type_hint=int,
+ description="Number of image tokens, used to compute the resolution-aware shift.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ num_inference_steps = block_state.num_inference_steps
+
+ sigmas = block_state.sigmas
+ if sigmas is None:
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+ else:
+ block_state.num_inference_steps = len(sigmas)
+
+ config = components.scheduler.config
+ mu = calculate_shift(
+ block_state.image_seq_len,
+ config.get("base_image_seq_len", 256),
+ config.get("max_image_seq_len", 6400),
+ config.get("base_shift", 0.5),
+ config.get("max_shift", 1.15),
+ )
+
+ components.scheduler.set_timesteps(sigmas=sigmas, mu=mu, device=device)
+ components.scheduler.set_begin_index(0)
+ block_state.timesteps = components.scheduler.timesteps
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2TurboSetTimestepsStep(ModularPipelineBlocks):
+ """
+ Step that sets the flow-matching schedule for the distilled Krea 2 turbo checkpoint on the scheduler: a linear
+ sigma schedule with the fixed time shift `mu=1.15` the checkpoint was distilled with.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 8):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+
+ Outputs:
+ timesteps (`Tensor`):
+ The denoising timesteps.
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that sets the flow-matching schedule for the distilled Krea 2 turbo checkpoint on the scheduler: a "
+ "linear sigma schedule with the fixed time shift `mu=1.15` the checkpoint was distilled with."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", default=8),
+ InputParam(
+ name="sigmas", type_hint=list, description="Custom sigma schedule (defaults to a linear ramp)."
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ num_inference_steps = block_state.num_inference_steps
+
+ sigmas = block_state.sigmas
+ if sigmas is None:
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+ else:
+ block_state.num_inference_steps = len(sigmas)
+
+ components.scheduler.set_timesteps(sigmas=sigmas, mu=1.15, device=device)
+ components.scheduler.set_begin_index(0)
+ block_state.timesteps = components.scheduler.timesteps
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2PreparePositionIdsStep(ModularPipelineBlocks):
+ """
+ Step that builds the shared rotary position ids for the combined [text | image] sequence: text at the origin, image
+ tokens at their (0, h, w) latent-grid coordinates. Place after prepare_latents.
+
+ Inputs:
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ prompt_embeds (`Tensor`):
+ Batch-expanded text features (only text_seq_len is used).
+
+ Outputs:
+ position_ids (`Tensor`):
+ Shared rotary coordinates (text_seq_len + grid_h * grid_w, 3).
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that builds the shared rotary position ids for the combined [text | image] sequence: text at the "
+ "origin, image tokens at their (0, h, w) latent-grid coordinates. Place after prepare_latents."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", default=1024),
+ InputParam.template("width", default=1024),
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Batch-expanded text features (only text_seq_len is used).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="position_ids",
+ type_hint=torch.Tensor,
+ description="Shared rotary coordinates (text_seq_len + grid_h * grid_w, 3).",
+ )
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.krea2.pipeline_krea2.Krea2Pipeline.prepare_position_ids
+ def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device):
+ """Build the `(text_seq_len + grid_height * grid_width, 3)` rotary coordinates for the combined sequence:
+ text tokens sit at the origin, image tokens carry their `(0, h, w)` latent-grid coordinates."""
+ text_ids = torch.zeros(text_seq_len, 3, device=device)
+ image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
+ image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
+ image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
+ image_ids = image_ids.reshape(grid_height * grid_width, 3)
+ return torch.cat([text_ids, image_ids], dim=0)
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ p = components.patch_size
+ grid_h = block_state.height // (components.vae_scale_factor * p)
+ grid_w = block_state.width // (components.vae_scale_factor * p)
+ text_seq_len = block_state.prompt_embeds.shape[1]
+
+ block_state.position_ids = self.prepare_position_ids(text_seq_len, grid_h, grid_w, device)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..190225bc18487f1a5f37f4e529a45737e6fb19d2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/decoders.py
@@ -0,0 +1,121 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKLQwenImage
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Krea2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# auto_docstring
+class Krea2DecodeStep(ModularPipelineBlocks):
+ """
+ Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the VAE's
+ per-channel statistics, and decodes them through the Qwen-Image VAE into images.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ latents (`Tensor`):
+ The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the "
+ "VAE's per-channel statistics, and decodes them through the Qwen-Image VAE into images."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLQwenImage),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ # Effective pixel-to-token downsampling factor: vae_scale_factor (8) * patch_size (2).
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("output_type", default="pil"),
+ InputParam.template("height", default=1024),
+ InputParam.template("width", default=1024),
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae = components.vae
+ p = components.patch_size
+ latents = block_state.latents
+
+ batch_size, _, channels = latents.shape
+ height = p * (int(block_state.height) // (components.vae_scale_factor * p))
+ width = p * (int(block_state.width) // (components.vae_scale_factor * p))
+ latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+ latents = latents.reshape(batch_size, channels // (p * p), 1, height, width)
+
+ latents = latents.to(vae.dtype)
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(
+ latents.device, latents.dtype
+ )
+ latents = latents / latents_std + latents_mean
+ image = vae.decode(latents, return_dict=False)[0][:, :, 0]
+ block_state.images = components.image_processor.postprocess(image, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..b45a419ff38b88e79f900140c74cbc7c98081a48
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/denoise.py
@@ -0,0 +1,369 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models.transformers.transformer_krea2 import Krea2Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Krea2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class Krea2LoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Within the denoising loop: normalize the scheduler timestep into the model's flow time and broadcast it "
+ "across the batch. Compose into the `sub_blocks` of a `Krea2DenoiseLoopWrapper`-based step."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+ InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ num_train_timesteps = components.scheduler.config.num_train_timesteps
+ block_state.timestep = (t / num_train_timesteps).expand(block_state.batch_size)
+ return components, block_state
+
+
+class Krea2LoopDenoiser(ModularPipelineBlocks):
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Within the denoising loop: run the `transformer` on the conditional (and, when the guider enables CFG, "
+ "the negative) text features and combine them through the `guider`. Compose into `Krea2DenoiseStep`."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ # Krea 2 uses cond-anchored CFG (`cond + scale * (cond - uncond)`), which is the
+ # `use_original_formulation` branch of ClassifierFreeGuidance; scale 0 disables it (distilled TDM).
+ config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", Krea2Transformer2DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+ InputParam.template("num_inference_steps", required=True),
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditional stacked text features.",
+ ),
+ InputParam(
+ name="prompt_embeds_mask", required=True, type_hint=torch.Tensor, description="Conditional text mask."
+ ),
+ InputParam(
+ name="position_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Shared rotary coordinates for the [text | image] sequence.",
+ ),
+ InputParam(
+ name="negative_prompt_embeds", type_hint=torch.Tensor, description="Negative stacked text features."
+ ),
+ InputParam(name="negative_prompt_embeds_mask", type_hint=torch.Tensor, description="Negative text mask."),
+ InputParam.template("attention_kwargs"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ transformer = components.transformer
+
+ latents = block_state.latents.to(transformer.dtype)
+ timestep = block_state.timestep.to(transformer.dtype)
+
+ guider_inputs = {
+ "encoder_hidden_states": (
+ block_state.prompt_embeds.to(transformer.dtype),
+ block_state.negative_prompt_embeds.to(transformer.dtype)
+ if block_state.negative_prompt_embeds is not None
+ else None,
+ ),
+ "encoder_attention_mask": (
+ block_state.prompt_embeds_mask,
+ block_state.negative_prompt_embeds_mask,
+ ),
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs}
+ guider_state_batch.noise_pred = transformer(
+ hidden_states=latents,
+ timestep=timestep,
+ position_ids=block_state.position_ids,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state).pred
+ return components, block_state
+
+
+class Krea2TurboLoopDenoiser(ModularPipelineBlocks):
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Within the denoising loop: run the `transformer` on the conditional text features. The distilled Krea 2 "
+ "turbo checkpoint runs without classifier-free guidance, so there is no negative branch or guider. Compose "
+ "into the `sub_blocks` of `Krea2TurboDenoiseStep`."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("transformer", Krea2Transformer2DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+ InputParam(
+ name="prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Conditional stacked text features.",
+ ),
+ InputParam(
+ name="prompt_embeds_mask", required=True, type_hint=torch.Tensor, description="Conditional text mask."
+ ),
+ InputParam(
+ name="position_ids",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Shared rotary coordinates for the [text | image] sequence.",
+ ),
+ InputParam.template("attention_kwargs"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ transformer = components.transformer
+
+ latents = block_state.latents.to(transformer.dtype)
+ timestep = block_state.timestep.to(transformer.dtype)
+
+ block_state.noise_pred = transformer(
+ hidden_states=latents,
+ timestep=timestep,
+ position_ids=block_state.position_ids,
+ attention_kwargs=block_state.attention_kwargs,
+ encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype),
+ encoder_attention_mask=block_state.prompt_embeds_mask,
+ return_dict=False,
+ )[0]
+ return components, block_state
+
+
+class Krea2LoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return "Within the denoising loop: scheduler step. Compose into a `Krea2DenoiseLoopWrapper`-based step."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="latents", type_hint=torch.Tensor, description="The denoised latents.")]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred, t, block_state.latents, return_dict=False
+ )[0]
+ block_state.latents = block_state.latents.to(latents_dtype)
+ return components, block_state
+
+
+class Krea2DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoises the packed image latents over `timesteps`. "
+ "The specific steps within each iteration can be customized with the `sub_blocks` attribute."
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Denoising timesteps from set_timesteps.",
+ ),
+ InputParam.template("num_inference_steps", required=True),
+ InputParam.template("attention_kwargs"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2DenoiseStep(Krea2DenoiseLoopWrapper):
+ """
+ Denoising loop that iteratively denoises the packed image latents over `timesteps`, running the transformer on the
+ conditional (and, when the guider enables CFG, the negative) text features and combining them through the `guider`.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer
+ (`Krea2Transformer2DModel`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ Denoising timesteps from set_timesteps.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ latents (`Tensor`):
+ Packed image latents.
+ batch_size (`int`):
+ Effective batch size.
+ prompt_embeds (`Tensor`):
+ Conditional stacked text features.
+ prompt_embeds_mask (`Tensor`):
+ Conditional text mask.
+ position_ids (`Tensor`):
+ Shared rotary coordinates for the [text | image] sequence.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Negative stacked text features.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Negative text mask.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoised latents.
+ """
+
+ model_name = "krea2"
+ block_classes = [Krea2LoopBeforeDenoiser, Krea2LoopDenoiser, Krea2LoopAfterDenoiser]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoising loop that iteratively denoises the packed image latents over `timesteps`, running the "
+ "transformer on the conditional (and, when the guider enables CFG, the negative) text features and "
+ "combining them through the `guider`."
+ )
+
+
+# auto_docstring
+class Krea2TurboDenoiseStep(Krea2DenoiseLoopWrapper):
+ """
+ Denoising loop for the distilled Krea 2 turbo checkpoint that iteratively denoises the packed image latents over
+ `timesteps`, running the transformer on the conditional text features. The distilled checkpoint runs without
+ classifier-free guidance.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Krea2Transformer2DModel`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ Denoising timesteps from set_timesteps.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ latents (`Tensor`):
+ Packed image latents.
+ batch_size (`int`):
+ Effective batch size.
+ prompt_embeds (`Tensor`):
+ Conditional stacked text features.
+ prompt_embeds_mask (`Tensor`):
+ Conditional text mask.
+ position_ids (`Tensor`):
+ Shared rotary coordinates for the [text | image] sequence.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoised latents.
+ """
+
+ model_name = "krea2"
+ block_classes = [Krea2LoopBeforeDenoiser, Krea2TurboLoopDenoiser, Krea2LoopAfterDenoiser]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoising loop for the distilled Krea 2 turbo checkpoint that iteratively denoises the packed image "
+ "latents over `timesteps`, running the transformer on the conditional text features. The distilled "
+ "checkpoint runs without classifier-free guidance."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec36106a9aaa28c60b518eb965eb4fdb4e62bd1d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/encoders.py
@@ -0,0 +1,276 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from transformers import AutoTokenizer, Qwen3VLModel
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Krea2ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Indices into the Qwen3-VL `hidden_states` tuple (0 is the embedding output) whose states are stacked per token as the
+# transformer's text conditioning. Must have `transformer.config.num_text_layers` entries.
+KREA2_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+
+# Krea 2 wraps the prompt in this Qwen-Image chat template before encoding. The prompt is padded to a fixed length
+# first and the assistant suffix is appended *after* the padding (matching how the model was sampled at training time);
+# the first `_PROMPT_TEMPLATE_ENCODE_START_IDX` (system prefix) tokens are dropped from the encoder outputs.
+_PROMPT_TEMPLATE_ENCODE_PREFIX = (
+ "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, "
+ "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
+)
+_PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
+_PROMPT_TEMPLATE_ENCODE_START_IDX = 34
+_PROMPT_TEMPLATE_ENCODE_NUM_SUFFIX_TOKENS = 5
+
+
+# auto_docstring
+class Krea2TextEncoderStep(ModularPipelineBlocks):
+ """
+ Text encoder step that tokenizes the prompt(s) with the Krea 2 chat template, runs the Qwen3-VL text encoder, and
+ stacks a fixed set of decoder-layer hidden states per token as the transformer's text conditioning. The negative
+ prompt is encoded the same way when the guider enables CFG.
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired
+ with the text encoder. guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The negative prompt(s) for CFG.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+ negative_prompt_embeds (`Tensor`):
+ Per-prompt negative text features (only when guidance is enabled).
+ negative_prompt_embeds_mask (`Tensor`):
+ Per-prompt negative text mask (only when guidance is enabled).
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text encoder step that tokenizes the prompt(s) with the Krea 2 chat template, runs the Qwen3-VL text "
+ "encoder, and stacks a fixed set of decoder-layer hidden states per token as the transformer's text "
+ "conditioning. The negative prompt is encoded the same way when the guider enables CFG."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."),
+ ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", required=True),
+ InputParam(name="negative_prompt", type_hint=str, description="The negative prompt(s) for CFG."),
+ InputParam.template("max_sequence_length", default=512),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).",
+ ),
+ OutputParam(
+ name="prompt_embeds_mask",
+ type_hint=torch.Tensor,
+ description="Per-prompt boolean text mask (B, text_seq_len).",
+ ),
+ OutputParam(
+ name="negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Per-prompt negative text features (only when guidance is enabled).",
+ ),
+ OutputParam(
+ name="negative_prompt_embeds_mask",
+ type_hint=torch.Tensor,
+ description="Per-prompt negative text mask (only when guidance is enabled).",
+ ),
+ ]
+
+ def _encode_prompt(self, components, prompt, max_sequence_length, device):
+ """Tokenize `prompt` into the fixed-length Krea 2 layout and tap the selected encoder hidden states.
+
+ Mirrors `Krea2Pipeline.get_text_hidden_states`. Returns a `(hidden_states, attention_mask)` tuple of shapes
+ `(batch_size, text_seq_len, num_text_layers, text_hidden_dim)` and `(batch_size, text_seq_len)` (bool).
+ """
+ tokenizer = components.tokenizer
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ prefix_idx = _PROMPT_TEMPLATE_ENCODE_START_IDX
+ text = [_PROMPT_TEMPLATE_ENCODE_PREFIX + e for e in prompt]
+ text_tokens = tokenizer(
+ text,
+ truncation=True,
+ padding="max_length",
+ max_length=max_sequence_length + prefix_idx - _PROMPT_TEMPLATE_ENCODE_NUM_SUFFIX_TOKENS,
+ return_tensors="pt",
+ ).to(device)
+ suffix_tokens = tokenizer([_PROMPT_TEMPLATE_ENCODE_SUFFIX] * len(text), return_tensors="pt").to(device)
+
+ input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1)
+ attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool()
+
+ # Krea 2 pads in the middle of the template (`[prefix | prompt | PAD | suffix]`), so the suffix tokens sit
+ # downstream of the padding. The text features must use positions that count only real tokens (padding does
+ # not consume a position) to match how the model was trained; otherwise the suffix gets a shifted mRoPE phase.
+ position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = components.text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ output_hidden_states=True,
+ )
+ hidden_states = torch.stack([outputs.hidden_states[i] for i in KREA2_TEXT_ENCODER_SELECT_LAYERS], dim=2)
+
+ hidden_states = hidden_states[:, prefix_idx:]
+ attention_mask = attention_mask[:, prefix_idx:]
+ return hidden_states, attention_mask
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt)
+
+ block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt(
+ components, prompts, block_state.max_sequence_length, device
+ )
+
+ block_state.negative_prompt_embeds = None
+ block_state.negative_prompt_embeds_mask = None
+ if components.requires_unconditional_embeds:
+ negative_prompt = block_state.negative_prompt
+ if negative_prompt is None:
+ negative_prompt = ""
+ if isinstance(negative_prompt, str):
+ negative_prompt = [negative_prompt] * len(prompts)
+ block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._encode_prompt(
+ components, negative_prompt, block_state.max_sequence_length, device
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class Krea2TurboTextEncoderStep(Krea2TextEncoderStep):
+ """
+ Text encoder step for the distilled Krea 2 turbo checkpoint that tokenizes the prompt(s) with the Krea 2 chat
+ template, runs the Qwen3-VL text encoder, and stacks a fixed set of decoder-layer hidden states per token as the
+ transformer's text conditioning. The distilled checkpoint runs without classifier-free guidance, so it takes no
+ negative prompt and has no guider.
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired
+ with the text encoder.
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+ """
+
+ model_name = "krea2"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text encoder step for the distilled Krea 2 turbo checkpoint that tokenizes the prompt(s) with the Krea 2 "
+ "chat template, runs the Qwen3-VL text encoder, and stacks a fixed set of decoder-layer hidden states per "
+ "token as the transformer's text conditioning. The distilled checkpoint runs without classifier-free "
+ "guidance, so it takes no negative prompt and has no guider."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."),
+ ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", required=True),
+ InputParam.template("max_sequence_length", default=512),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).",
+ ),
+ OutputParam(
+ name="prompt_embeds_mask",
+ type_hint=torch.Tensor,
+ description="Per-prompt boolean text mask (B, text_seq_len).",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt)
+
+ block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt(
+ components, prompts, block_state.max_sequence_length, device
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py
new file mode 100644
index 0000000000000000000000000000000000000000..658133121dd6f5f1e07dec7dea3e686b9a04c675
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py
@@ -0,0 +1,170 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Krea2PrepareLatentsStep,
+ Krea2PreparePositionIdsStep,
+ Krea2SetTimestepsStep,
+ Krea2TextInputsStep,
+)
+from .decoders import Krea2DecodeStep
+from .denoise import Krea2DenoiseStep
+from .encoders import Krea2TextEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+CORE_DENOISE_BLOCKS = InsertableDict(
+ [
+ ("input", Krea2TextInputsStep()),
+ ("prepare_latents", Krea2PrepareLatentsStep()),
+ ("set_timesteps", Krea2SetTimestepsStep()),
+ ("prepare_position_ids", Krea2PreparePositionIdsStep()),
+ ("denoise", Krea2DenoiseStep()),
+ ]
+)
+
+
+# auto_docstring
+class Krea2CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for Krea 2 text-to-image: prepares the batch/latents/timesteps and the shared position ids,
+ then runs the symmetric-CFG denoising loop, producing the denoised packed latents for the decoder.
+
+ Components:
+ transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Per-prompt negative text features.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ Per-prompt negative text mask.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 28):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoised packed latents (B, image_seq_len, in_channels).
+ """
+
+ model_name = "krea2"
+ block_classes = list(CORE_DENOISE_BLOCKS.values())
+ block_names = list(CORE_DENOISE_BLOCKS.keys())
+
+ @property
+ def description(self) -> str:
+ return (
+ "Core denoising workflow for Krea 2 text-to-image: prepares the batch/latents/timesteps and the shared "
+ "position ids, then runs the symmetric-CFG denoising loop, producing the denoised packed latents for the "
+ "decoder."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("latents", description="The denoised packed latents (B, image_seq_len, in_channels).")
+ ]
+
+
+# auto_docstring
+class Krea2AutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise (symmetric CFG) ->
+ decode.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired
+ with the text encoder. guider (`ClassifierFreeGuidance`) transformer (`Krea2Transformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The negative prompt(s) for CFG.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 28):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "krea2"
+ block_classes = [
+ Krea2TextEncoderStep,
+ Krea2CoreDenoiseStep,
+ Krea2DecodeStep,
+ ]
+ block_names = ["text_encoder", "denoise", "decode"]
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return (
+ "Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise "
+ "(symmetric CFG) -> decode."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py
new file mode 100644
index 0000000000000000000000000000000000000000..3455f5d253b4094dc2409ce16726540599202f3b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py
@@ -0,0 +1,164 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ Krea2PrepareLatentsStep,
+ Krea2PreparePositionIdsStep,
+ Krea2TurboSetTimestepsStep,
+ Krea2TurboTextInputsStep,
+)
+from .decoders import Krea2DecodeStep
+from .denoise import Krea2TurboDenoiseStep
+from .encoders import Krea2TurboTextEncoderStep
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+CORE_DENOISE_BLOCKS = InsertableDict(
+ [
+ ("input", Krea2TurboTextInputsStep()),
+ ("prepare_latents", Krea2PrepareLatentsStep()),
+ ("set_timesteps", Krea2TurboSetTimestepsStep()),
+ ("prepare_position_ids", Krea2PreparePositionIdsStep()),
+ ("denoise", Krea2TurboDenoiseStep()),
+ ]
+)
+
+
+# auto_docstring
+class Krea2TurboCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for the distilled Krea 2 turbo text-to-image checkpoint: prepares the
+ batch/latents/timesteps and the shared position ids, then runs the guidance-free denoising loop, producing the
+ denoised packed latents for the decoder.
+
+ Components:
+ transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).
+ prompt_embeds_mask (`Tensor`):
+ Per-prompt boolean text mask (B, text_seq_len).
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 8):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoised packed latents (B, image_seq_len, in_channels).
+ """
+
+ model_name = "krea2"
+ block_classes = list(CORE_DENOISE_BLOCKS.values())
+ block_names = list(CORE_DENOISE_BLOCKS.keys())
+
+ @property
+ def description(self) -> str:
+ return (
+ "Core denoising workflow for the distilled Krea 2 turbo text-to-image checkpoint: prepares the "
+ "batch/latents/timesteps and the shared position ids, then runs the guidance-free denoising loop, "
+ "producing the denoised packed latents for the decoder."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("latents", description="The denoised packed latents (B, image_seq_len, in_channels).")
+ ]
+
+
+# auto_docstring
+class Krea2TurboAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode text -> core
+ denoise (guidance-free) -> decode.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+
+ Components:
+ text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired
+ with the text encoder. transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ max_sequence_length (`int`, *optional*, defaults to 512):
+ Maximum sequence length for prompt encoding.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 8):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigma schedule (defaults to a linear ramp).
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "krea2"
+ block_classes = [
+ Krea2TurboTextEncoderStep,
+ Krea2TurboCoreDenoiseStep,
+ Krea2DecodeStep,
+ ]
+ block_names = ["text_encoder", "denoise", "decode"]
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return (
+ "Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode "
+ "text -> core denoise (guidance-free) -> decode."
+ )
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c251c82a5f9883a8b529210be73ff856cb7a1fd
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/krea2/modular_pipeline.py
@@ -0,0 +1,67 @@
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...loaders import Krea2LoraLoaderMixin
+from ..modular_pipeline import ModularPipeline
+
+
+class Krea2ModularPipeline(ModularPipeline, Krea2LoraLoaderMixin):
+ """
+ A ModularPipeline for Krea 2.
+
+ > [!WARNING] > This is an experimental feature!
+ """
+
+ default_blocks_name = "Krea2AutoBlocks"
+
+ @property
+ def patch_size(self):
+ return 2
+
+ @property
+ def default_height(self):
+ return 1024
+
+ @property
+ def default_width(self):
+ return 1024
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if getattr(self, "vae", None) is not None:
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+ return requires_unconditional_embeds
+
+
+class Krea2TurboModularPipeline(Krea2ModularPipeline):
+ """
+ A ModularPipeline for the distilled Krea 2 turbo (TDM) checkpoint. It runs without classifier-free guidance, so it
+ takes no negative prompt and has no guider.
+
+ > [!WARNING] > This is an experimental feature!
+ """
+
+ default_blocks_name = "Krea2TurboAutoBlocks"
+
+ @property
+ def requires_unconditional_embeds(self):
+ return False
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1de33ffee874dab8f45702b42c257d0aa770ea6
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_ltx"] = ["LTXAutoBlocks", "LTXBlocks", "LTXImage2VideoBlocks"]
+ _import_structure["modular_pipeline"] = ["LTXModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_ltx import LTXAutoBlocks, LTXBlocks, LTXImage2VideoBlocks
+ from .modular_pipeline import LTXModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c1716580cd931a2171892a5c3e15df37b9b7ac7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/before_denoise.py
@@ -0,0 +1,392 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import LTXModularPipeline, LTXVideoPachifier
+
+
+logger = logging.get_logger(__name__)
+
+
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class LTXTextInputStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Adjusts input tensor shapes based on `batch_size` and `num_videos_per_prompt`"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ InputParam.template("prompt_embeds", required=True),
+ InputParam.template("prompt_embeds_mask", name="prompt_attention_mask"),
+ InputParam.template("negative_prompt_embeds"),
+ InputParam.template("negative_prompt_embeds_mask", name="negative_prompt_attention_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("batch_size", type_hint=int),
+ OutputParam("dtype", type_hint=torch.dtype),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+ num_videos = block_state.num_videos_per_prompt
+
+ # Repeat prompt_embeds for num_videos_per_prompt
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, num_videos, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(block_state.batch_size * num_videos, seq_len, -1)
+
+ if block_state.prompt_attention_mask is not None:
+ block_state.prompt_attention_mask = block_state.prompt_attention_mask.repeat(num_videos, 1)
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(1, num_videos, 1)
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * num_videos, seq_len, -1
+ )
+
+ if block_state.negative_prompt_attention_mask is not None:
+ block_state.negative_prompt_attention_mask = block_state.negative_prompt_attention_mask.repeat(
+ num_videos, 1
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class LTXSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("timesteps"),
+ InputParam.template("sigmas"),
+ InputParam.template("height", default=512),
+ InputParam.template("width", default=704),
+ InputParam("num_frames", type_hint=int, default=161),
+ InputParam("frame_rate", type_hint=int, default=25),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor),
+ OutputParam("num_inference_steps", type_hint=int),
+ OutputParam("rope_interpolation_scale", type_hint=tuple),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ height = block_state.height
+ width = block_state.width
+ num_frames = block_state.num_frames
+ frame_rate = block_state.frame_rate
+
+ latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1
+ latent_height = height // components.vae_spatial_compression_ratio
+ latent_width = width // components.vae_spatial_compression_ratio
+ video_sequence_length = latent_num_frames * latent_height * latent_width
+
+ custom_timesteps = block_state.timesteps
+ sigmas = block_state.sigmas
+
+ if custom_timesteps is not None:
+ # User provided custom timesteps, don't compute sigmas
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ device,
+ custom_timesteps,
+ )
+ else:
+ if sigmas is None:
+ sigmas = np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+
+ mu = calculate_shift(
+ video_sequence_length,
+ components.scheduler.config.get("base_image_seq_len", 256),
+ components.scheduler.config.get("max_image_seq_len", 4096),
+ components.scheduler.config.get("base_shift", 0.5),
+ components.scheduler.config.get("max_shift", 1.15),
+ )
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+
+ block_state.rope_interpolation_scale = (
+ components.vae_temporal_compression_ratio / frame_rate,
+ components.vae_spatial_compression_ratio,
+ components.vae_spatial_compression_ratio,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class LTXPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the latents for the text-to-video generation process"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "pachifier",
+ LTXVideoPachifier,
+ config=FrozenDict({"patch_size": 1, "patch_size_t": 1}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", default=512),
+ InputParam.template("width", default=704),
+ InputParam("num_frames", type_hint=int, default=161),
+ InputParam.template("latents"),
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ InputParam.template("generator"),
+ InputParam.template("batch_size", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ batch_size = block_state.batch_size * block_state.num_videos_per_prompt
+ num_channels_latents = components.transformer.config.in_channels
+
+ if block_state.latents is not None:
+ block_state.latents = block_state.latents.to(device=device, dtype=torch.float32)
+ else:
+ height = block_state.height // components.vae_spatial_compression_ratio
+ width = block_state.width // components.vae_spatial_compression_ratio
+ num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
+
+ shape = (batch_size, num_channels_latents, num_frames, height, width)
+ block_state.latents = randn_tensor(
+ shape, generator=block_state.generator, device=device, dtype=torch.float32
+ )
+ block_state.latents = components.pachifier.pack_latents(block_state.latents)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class LTXImage2VideoPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Prepare image-to-video latents: adds noise to pre-encoded image latents and creates a conditioning mask. "
+ "Expects pure noise `latents` from LTXPrepareLatentsStep."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "pachifier",
+ LTXVideoPachifier,
+ config=FrozenDict({"patch_size": 1, "patch_size_t": 1}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image_latents", type_hint=torch.Tensor, required=True),
+ InputParam.template("latents", required=True),
+ InputParam.template("height", default=512),
+ InputParam.template("width", default=704),
+ InputParam("num_frames", type_hint=int, default=161),
+ InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
+ InputParam.template("batch_size", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("latents", type_hint=torch.Tensor),
+ OutputParam("conditioning_mask", type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ batch_size = block_state.batch_size * block_state.num_videos_per_prompt
+
+ height = block_state.height // components.vae_spatial_compression_ratio
+ width = block_state.width // components.vae_spatial_compression_ratio
+ num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
+
+ init_latents = block_state.image_latents.to(device=device, dtype=torch.float32)
+ if init_latents.shape[0] < batch_size:
+ init_latents = init_latents.repeat_interleave(batch_size // init_latents.shape[0], dim=0)
+ init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
+
+ conditioning_mask = torch.zeros(
+ init_latents.shape[0],
+ 1,
+ init_latents.shape[2],
+ init_latents.shape[3],
+ init_latents.shape[4],
+ device=device,
+ dtype=torch.float32,
+ )
+ conditioning_mask[:, :, 0] = 1.0
+
+ noise = components.pachifier.unpack_latents(block_state.latents, num_frames, height, width)
+ latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
+
+ conditioning_mask = components.pachifier.pack_latents(conditioning_mask).squeeze(-1)
+ latents = components.pachifier.pack_latents(latents)
+
+ block_state.latents = latents
+ block_state.conditioning_mask = conditioning_mask
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..37004949ea56473b32167f70c665b9707560eeb7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/decoders.py
@@ -0,0 +1,132 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLLTXVideo
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import LTXVideoPachifier
+
+
+logger = logging.get_logger(__name__)
+
+
+def _denormalize_latents(
+ latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+) -> torch.Tensor:
+ # Denormalize latents across the channel dimension [B, C, F, H, W]
+ latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+ latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+ latents = latents * latents_std / scaling_factor + latents_mean
+ return latents
+
+
+class LTXVaeDecoderStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLLTXVideo),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 32}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec(
+ "pachifier",
+ LTXVideoPachifier,
+ config=FrozenDict({"patch_size": 1, "patch_size_t": 1}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into videos"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam.template("latents", required=True),
+ InputParam.template("output_type", default="np"),
+ InputParam.template("height", default=512),
+ InputParam.template("width", default=704),
+ InputParam("num_frames", type_hint=int, default=161),
+ InputParam("decode_timestep", default=0.0),
+ InputParam("decode_noise_scale", default=None),
+ InputParam.template("generator"),
+ InputParam.template("batch_size"),
+ InputParam.template("dtype", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("videos")]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae = components.vae
+
+ latents = block_state.latents
+
+ height = block_state.height
+ width = block_state.width
+ num_frames = block_state.num_frames
+
+ latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1
+ latent_height = height // components.vae_spatial_compression_ratio
+ latent_width = width // components.vae_spatial_compression_ratio
+
+ latents = components.pachifier.unpack_latents(latents, latent_num_frames, latent_height, latent_width)
+ latents = _denormalize_latents(latents, vae.latents_mean, vae.latents_std, vae.config.scaling_factor)
+ latents = latents.to(block_state.dtype)
+
+ if not vae.config.timestep_conditioning:
+ timestep = None
+ else:
+ device = latents.device
+ batch_size = block_state.batch_size
+ decode_timestep = block_state.decode_timestep
+ decode_noise_scale = block_state.decode_noise_scale
+
+ noise = randn_tensor(latents.shape, generator=block_state.generator, device=device, dtype=latents.dtype)
+ if not isinstance(decode_timestep, list):
+ decode_timestep = [decode_timestep] * batch_size
+ if decode_noise_scale is None:
+ decode_noise_scale = decode_timestep
+ elif not isinstance(decode_noise_scale, list):
+ decode_noise_scale = [decode_noise_scale] * batch_size
+
+ timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
+ decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
+ :, None, None, None, None
+ ]
+ latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
+
+ latents = latents.to(vae.dtype)
+ video = vae.decode(latents, timestep, return_dict=False)[0]
+ block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..79a9f771217c0b577507c7f64b5bfa46968db52c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/denoise.py
@@ -0,0 +1,458 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import LTXVideoTransformer3DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam
+from .modular_pipeline import LTXModularPipeline, LTXVideoPachifier
+
+
+class LTXLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `LTXDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True),
+ InputParam.template("dtype", required=True),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = block_state.latents.to(block_state.dtype)
+ return components, block_state
+
+
+class LTXLoopDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] | None = None,
+ ):
+ if guider_input_fields is None:
+ guider_input_fields = {
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"),
+ }
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 3.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", LTXVideoTransformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoises the latents with guidance. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `LTXDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ inputs = [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("num_inference_steps", required=True),
+ InputParam("rope_interpolation_scale", type_hint=tuple),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam("num_frames", type_hint=int),
+ ]
+ guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.extend(value)
+ else:
+ guider_input_names.append(value)
+
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
+ latent_height = block_state.height // components.vae_spatial_compression_ratio
+ latent_width = block_state.width // components.vae_spatial_compression_ratio
+
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = guider_state_batch.as_dict()
+ cond_kwargs = {
+ k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v
+ for k, v in cond_kwargs.items()
+ if k in self._guider_input_fields.keys()
+ }
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key, None)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=t.expand(block_state.latent_model_input.shape[0]).to(block_state.dtype),
+ num_frames=latent_num_frames,
+ height=latent_height,
+ width=latent_width,
+ rope_interpolation_scale=block_state.rope_interpolation_scale,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class LTXLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that updates the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `LTXDenoiseLoopWrapper`)"
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class LTXDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoises the latents over `timesteps`. "
+ "The specific steps within each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", LTXVideoTransformer3DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("timesteps", required=True),
+ InputParam.template("num_inference_steps", required=True),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class LTXDenoiseStep(LTXDenoiseLoopWrapper):
+ block_classes = [
+ LTXLoopBeforeDenoiser,
+ LTXLoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"),
+ }
+ ),
+ LTXLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoises the latents.\n"
+ "Its loop logic is defined in `LTXDenoiseLoopWrapper.__call__` method.\n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `LTXLoopBeforeDenoiser`\n"
+ " - `LTXLoopDenoiser`\n"
+ " - `LTXLoopAfterDenoiser`\n"
+ "This block supports text-to-video tasks."
+ )
+
+
+class LTXImage2VideoLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the i2v denoising loop that prepares the latent input and modulates "
+ "the timestep with the conditioning mask."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents", required=True),
+ InputParam("conditioning_mask", required=True, type_hint=torch.Tensor),
+ InputParam.template("dtype", required=True),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = block_state.latents.to(block_state.dtype)
+ block_state.timestep_adjusted = t.expand(block_state.latent_model_input.shape[0]).unsqueeze(-1) * (
+ 1 - block_state.conditioning_mask
+ )
+ return components, block_state
+
+
+class LTXImage2VideoLoopDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] | None = None,
+ ):
+ if guider_input_fields is None:
+ guider_input_fields = {
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"),
+ }
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 3.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", LTXVideoTransformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the i2v denoising loop that denoises the latents with guidance "
+ "using timestep modulated by the conditioning mask."
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ inputs = [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("num_inference_steps", required=True),
+ InputParam("rope_interpolation_scale", type_hint=tuple),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam("num_frames", type_hint=int),
+ ]
+ guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.extend(value)
+ else:
+ guider_input_names.append(value)
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
+ latent_height = block_state.height // components.vae_spatial_compression_ratio
+ latent_width = block_state.width // components.vae_spatial_compression_ratio
+
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = guider_state_batch.as_dict()
+ cond_kwargs = {
+ k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v
+ for k, v in cond_kwargs.items()
+ if k in self._guider_input_fields.keys()
+ }
+
+ context_name = getattr(guider_state_batch, components.guider._identifier_key, None)
+ with components.transformer.cache_context(context_name):
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=block_state.timestep_adjusted,
+ num_frames=latent_num_frames,
+ height=latent_height,
+ width=latent_width,
+ rope_interpolation_scale=block_state.rope_interpolation_scale,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class LTXImage2VideoLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec(
+ "pachifier",
+ LTXVideoPachifier,
+ config=FrozenDict({"patch_size": 1, "patch_size_t": 1}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the i2v denoising loop that updates the latents, "
+ "applying the scheduler step only to frames after the first (conditioned) frame."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam("num_frames", type_hint=int),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
+ latent_height = block_state.height // components.vae_spatial_compression_ratio
+ latent_width = block_state.width // components.vae_spatial_compression_ratio
+
+ noise_pred = components.pachifier.unpack_latents(
+ block_state.noise_pred, latent_num_frames, latent_height, latent_width
+ )
+ latents = components.pachifier.unpack_latents(
+ block_state.latents, latent_num_frames, latent_height, latent_width
+ )
+
+ noise_pred = noise_pred[:, :, 1:]
+ noise_latents = latents[:, :, 1:]
+ pred_latents = components.scheduler.step(noise_pred, t, noise_latents, return_dict=False)[0]
+
+ latents = torch.cat([latents[:, :, :1], pred_latents], dim=2)
+ block_state.latents = components.pachifier.pack_latents(latents)
+
+ return components, block_state
+
+
+class LTXImage2VideoDenoiseStep(LTXDenoiseLoopWrapper):
+ block_classes = [
+ LTXImage2VideoLoopBeforeDenoiser,
+ LTXImage2VideoLoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"),
+ }
+ ),
+ LTXImage2VideoLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step for image-to-video that iteratively denoises the latents.\n"
+ "The first frame is kept fixed via a conditioning mask.\n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `LTXImage2VideoLoopBeforeDenoiser`\n"
+ " - `LTXImage2VideoLoopDenoiser`\n"
+ " - `LTXImage2VideoLoopAfterDenoiser`"
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..86080fe024c13909cdfc0ecf85cbf56aebf939f3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/encoders.py
@@ -0,0 +1,273 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from transformers import T5EncoderModel, T5TokenizerFast
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import AutoencoderKLLTXVideo
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import LTXModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_t5_prompt_embeds(
+ components,
+ prompt: str | list[str],
+ max_sequence_length: int,
+ device: torch.device,
+ dtype: torch.dtype,
+):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ text_inputs = components.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ prompt_attention_mask = text_inputs.attention_mask
+ prompt_attention_mask = prompt_attention_mask.bool().to(device)
+
+ prompt_embeds = components.text_encoder(text_input_ids.to(device))[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ return prompt_embeds, prompt_attention_mask
+
+
+class LTXTextEncoderStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", T5EncoderModel),
+ ComponentSpec("tokenizer", T5TokenizerFast),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 3.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam.template("max_sequence_length", default=128),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("prompt_embeds_mask", name="prompt_attention_mask"),
+ OutputParam.template("negative_prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds_mask", name="negative_prompt_attention_mask"),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ if block_state.prompt is not None and (
+ not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list)
+ ):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str,
+ device: torch.device | None = None,
+ prepare_unconditional_embeds: bool = True,
+ negative_prompt: str | None = None,
+ max_sequence_length: int = 128,
+ ):
+ device = device or components._execution_device
+ dtype = components.text_encoder.dtype
+
+ if not isinstance(prompt, list):
+ prompt = [prompt]
+ batch_size = len(prompt)
+
+ prompt_embeds, prompt_attention_mask = _get_t5_prompt_embeds(
+ components=components,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ negative_prompt_embeds = None
+ negative_prompt_attention_mask = None
+
+ if prepare_unconditional_embeds:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds, negative_prompt_attention_mask = _get_t5_prompt_embeds(
+ components=components,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ (
+ block_state.prompt_embeds,
+ block_state.prompt_attention_mask,
+ block_state.negative_prompt_embeds,
+ block_state.negative_prompt_attention_mask,
+ ) = self.encode_prompt(
+ components=components,
+ prompt=block_state.prompt,
+ device=block_state.device,
+ prepare_unconditional_embeds=components.requires_unconditional_embeds,
+ negative_prompt=block_state.negative_prompt,
+ max_sequence_length=block_state.max_sequence_length,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def _normalize_latents(
+ latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+) -> torch.Tensor:
+ # Normalize latents across the channel dimension [B, C, F, H, W]
+ latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+ latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+ latents = (latents - latents_mean) * scaling_factor / latents_std
+ return latents
+
+
+class LTXVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "ltx"
+
+ @property
+ def description(self) -> str:
+ return "VAE Encoder step that encodes an input image into latent space for image-to-video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLLTXVideo),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 32}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image", required=True),
+ InputParam.template("height", default=512),
+ InputParam.template("width", default=704),
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="Encoded image latents from the VAE encoder",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ image = block_state.image
+ if not isinstance(image, torch.Tensor):
+ image = components.video_processor.preprocess(image, height=block_state.height, width=block_state.width)
+ image = image.to(device=device, dtype=torch.float32)
+
+ vae_dtype = components.vae.dtype
+
+ num_images = image.shape[0]
+ if isinstance(block_state.generator, list):
+ init_latents = [
+ retrieve_latents(
+ components.vae.encode(image[i].unsqueeze(0).unsqueeze(2).to(vae_dtype)),
+ block_state.generator[i],
+ )
+ for i in range(num_images)
+ ]
+ else:
+ init_latents = [
+ retrieve_latents(
+ components.vae.encode(img.unsqueeze(0).unsqueeze(2).to(vae_dtype)),
+ block_state.generator,
+ )
+ for img in image
+ ]
+
+ init_latents = torch.cat(init_latents, dim=0).to(torch.float32)
+ block_state.image_latents = _normalize_latents(
+ init_latents, components.vae.latents_mean, components.vae.latents_std
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7ef983eb735a00060ff6c5fcf74024808662c6b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py
@@ -0,0 +1,487 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ LTXImage2VideoPrepareLatentsStep,
+ LTXPrepareLatentsStep,
+ LTXSetTimestepsStep,
+ LTXTextInputStep,
+)
+from .decoders import LTXVaeDecoderStep
+from .denoise import LTXDenoiseStep, LTXImage2VideoDenoiseStep
+from .encoders import LTXTextEncoderStep, LTXVaeEncoderStep
+
+
+logger = logging.get_logger(__name__)
+
+
+# auto_docstring
+class LTXCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider
+ (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_attention_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_attention_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "ltx"
+ block_classes = [
+ LTXTextInputStep,
+ LTXSetTimestepsStep,
+ LTXPrepareLatentsStep,
+ LTXDenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class LTXImage2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Denoise block for image-to-video that takes encoded conditions and image latents, and runs the denoising process.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider
+ (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_attention_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_attention_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`):
+ TODO: Add description.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "ltx"
+ block_classes = [
+ LTXTextInputStep,
+ LTXSetTimestepsStep,
+ LTXPrepareLatentsStep,
+ LTXImage2VideoPrepareLatentsStep,
+ LTXImage2VideoDenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "Denoise block for image-to-video that takes encoded conditions and image latents, and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class LTXBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for LTX Video text-to-video.
+
+ Components:
+ text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) transformer
+ (`LTXVideoTransformer3DModel`) vae (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 128):
+ Maximum sequence length for prompt encoding.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+ decode_timestep (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ decode_noise_scale (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "ltx"
+ block_classes = [
+ LTXTextEncoderStep,
+ LTXCoreDenoiseStep,
+ LTXVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "denoise", "decode"]
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for LTX Video text-to-video."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
+
+
+# auto_docstring
+class LTXAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder step that encodes the image input into its latent representation.
+ This is an auto pipeline block that works for image-to-video tasks.
+ - `LTXVaeEncoderStep` is used when `image` is provided.
+ - If `image` is not provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ image_latents (`Tensor`):
+ Encoded image latents from the VAE encoder
+ """
+
+ model_name = "ltx"
+ block_classes = [LTXVaeEncoderStep]
+ block_names = ["vae_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder step that encodes the image input into its latent representation.\n"
+ "This is an auto pipeline block that works for image-to-video tasks.\n"
+ " - `LTXVaeEncoderStep` is used when `image` is provided.\n"
+ " - If `image` is not provided, step will be skipped."
+ )
+
+
+# auto_docstring
+class LTXAutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Auto denoise block that selects the appropriate denoise pipeline based on inputs.
+ - `LTXImage2VideoCoreDenoiseStep` is used when `image_latents` is provided.
+ - `LTXCoreDenoiseStep` is used otherwise (text-to-video).
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider
+ (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`)
+
+ Inputs:
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_attention_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_attention_mask (`Tensor`):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ timesteps (`Tensor`):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ image_latents (`Tensor`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "ltx"
+ block_classes = [LTXImage2VideoCoreDenoiseStep, LTXCoreDenoiseStep]
+ block_names = ["image2video", "text2video"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Auto denoise block that selects the appropriate denoise pipeline based on inputs.\n"
+ " - `LTXImage2VideoCoreDenoiseStep` is used when `image_latents` is provided.\n"
+ " - `LTXCoreDenoiseStep` is used otherwise (text-to-video)."
+ )
+
+
+# auto_docstring
+class LTXAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto blocks for LTX Video that support both text-to-video and image-to-video workflows.
+
+ Supported workflows:
+ - `text2video`: requires `prompt`
+ - `image2video`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ pachifier (`LTXVideoPachifier`) transformer (`LTXVideoTransformer3DModel`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 128):
+ Maximum sequence length for prompt encoding.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ timesteps (`Tensor`):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ image_latents (`Tensor`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+ decode_timestep (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ decode_noise_scale (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "ltx"
+ block_classes = [
+ LTXTextEncoderStep,
+ LTXAutoVaeEncoderStep,
+ LTXAutoCoreDenoiseStep,
+ LTXVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2video": {"prompt": True},
+ "image2video": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto blocks for LTX Video that support both text-to-video and image-to-video workflows."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
+
+
+# auto_docstring
+class LTXImage2VideoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for LTX Video image-to-video.
+
+ Components:
+ text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ pachifier (`LTXVideoPachifier`) transformer (`LTXVideoTransformer3DModel`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 128):
+ Maximum sequence length for prompt encoding.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*, defaults to 512):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to 704):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`Tensor`, *optional*):
+ Timesteps for the denoising process.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ num_frames (`int`, *optional*, defaults to 161):
+ TODO: Add description.
+ frame_rate (`int`, *optional*, defaults to 25):
+ TODO: Add description.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ image_latents (`Tensor`):
+ TODO: Add description.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to np):
+ Output format: 'pil', 'np', 'pt'.
+ decode_timestep (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ decode_noise_scale (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "ltx"
+ block_classes = [
+ LTXTextEncoderStep,
+ LTXAutoVaeEncoderStep,
+ LTXImage2VideoCoreDenoiseStep,
+ LTXVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for LTX Video image-to-video."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..6af4abcdd7f2a0c4c5375b6651a4dd8d1336c25a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ltx/modular_pipeline.py
@@ -0,0 +1,95 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...loaders import LTXVideoLoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class LTXVideoPachifier(ConfigMixin):
+ """
+ A class to pack and unpack latents for LTX Video.
+ """
+
+ config_name = "config.json"
+
+ @register_to_config
+ def __init__(self, patch_size: int = 1, patch_size_t: int = 1):
+ super().__init__()
+
+ def pack_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ batch_size, _, num_frames, height, width = latents.shape
+ patch_size = self.config.patch_size
+ patch_size_t = self.config.patch_size_t
+ post_patch_num_frames = num_frames // patch_size_t
+ post_patch_height = height // patch_size
+ post_patch_width = width // patch_size
+ latents = latents.reshape(
+ batch_size,
+ -1,
+ post_patch_num_frames,
+ patch_size_t,
+ post_patch_height,
+ patch_size,
+ post_patch_width,
+ patch_size,
+ )
+ latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+ return latents
+
+ def unpack_latents(self, latents: torch.Tensor, num_frames: int, height: int, width: int) -> torch.Tensor:
+ batch_size = latents.size(0)
+ patch_size = self.config.patch_size
+ patch_size_t = self.config.patch_size_t
+ latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
+ latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+ return latents
+
+
+class LTXModularPipeline(
+ ModularPipeline,
+ LTXVideoLoraLoaderMixin,
+):
+ """
+ A ModularPipeline for LTX Video.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "LTXAutoBlocks"
+
+ @property
+ def vae_spatial_compression_ratio(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.spatial_compression_ratio
+ return 32
+
+ @property
+ def vae_temporal_compression_ratio(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.temporal_compression_ratio
+ return 8
+
+ @property
+ def requires_unconditional_embeds(self):
+ if hasattr(self, "guider") and self.guider is not None:
+ return self.guider._enabled and self.guider.num_conditions > 1
+ return False
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/mellon_node_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/mellon_node_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f93d9c1c110b4a814fdd4019af9434d79537eaf
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/mellon_node_utils.py
@@ -0,0 +1,1101 @@
+import copy
+import json
+import logging
+import os
+
+# Simple typed wrapper for parameter overrides
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from huggingface_hub import create_repo, hf_hub_download, upload_file
+from huggingface_hub.utils import (
+ EntryNotFoundError,
+ HfHubHTTPError,
+ RepositoryNotFoundError,
+ RevisionNotFoundError,
+)
+
+from ..utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT
+from .modular_pipeline_utils import InputParam, OutputParam
+
+
+logger = logging.getLogger(__name__)
+
+
+def _name_to_label(name: str) -> str:
+ """Convert snake_case name to Title Case label."""
+ return name.replace("_", " ").title()
+
+
+# Template definitions for standard diffuser pipeline parameters
+MELLON_PARAM_TEMPLATES = {
+ # Image I/O
+ "image": {"label": "Image", "type": "image", "display": "input", "required_block_params": ["image"]},
+ "images": {"label": "Images", "type": "image", "display": "output", "required_block_params": ["images"]},
+ "control_image": {
+ "label": "Control Image",
+ "type": "image",
+ "display": "input",
+ "required_block_params": ["control_image"],
+ },
+ # Latents
+ "latents": {"label": "Latents", "type": "latents", "display": "input", "required_block_params": ["latents"]},
+ "image_latents": {
+ "label": "Image Latents",
+ "type": "latents",
+ "display": "input",
+ "required_block_params": ["image_latents"],
+ },
+ "first_frame_latents": {
+ "label": "First Frame Latents",
+ "type": "latents",
+ "display": "input",
+ "required_block_params": ["first_frame_latents"],
+ },
+ "latents_preview": {"label": "Latents Preview", "type": "latent", "display": "output"},
+ # Image Latents with Strength
+ "image_latents_with_strength": {
+ "name": "image_latents", # name is not same as template key
+ "label": "Image Latents",
+ "type": "latents",
+ "display": "input",
+ "onChange": {"false": ["height", "width"], "true": ["strength"]},
+ "required_block_params": ["image_latents", "strength"],
+ },
+ # Embeddings
+ "embeddings": {"label": "Text Embeddings", "type": "embeddings", "display": "output"},
+ "image_embeds": {
+ "label": "Image Embeddings",
+ "type": "image_embeds",
+ "display": "output",
+ "required_block_params": ["image_embeds"],
+ },
+ # Text inputs
+ "prompt": {
+ "label": "Prompt",
+ "type": "string",
+ "display": "textarea",
+ "default": "",
+ "required_block_params": ["prompt"],
+ },
+ "negative_prompt": {
+ "label": "Negative Prompt",
+ "type": "string",
+ "display": "textarea",
+ "default": "",
+ "required_block_params": ["negative_prompt"],
+ },
+ # Numeric params
+ "guidance_scale": {
+ "label": "Guidance Scale",
+ "type": "float",
+ "display": "slider",
+ "default": 5.0,
+ "min": 1.0,
+ "max": 30.0,
+ "step": 0.1,
+ },
+ "strength": {
+ "label": "Strength",
+ "type": "float",
+ "default": 0.5,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.01,
+ "required_block_params": ["strength"],
+ },
+ "height": {
+ "label": "Height",
+ "type": "int",
+ "default": 1024,
+ "min": 64,
+ "step": 8,
+ "required_block_params": ["height"],
+ },
+ "width": {
+ "label": "Width",
+ "type": "int",
+ "default": 1024,
+ "min": 64,
+ "step": 8,
+ "required_block_params": ["width"],
+ },
+ "seed": {
+ "label": "Seed",
+ "type": "int",
+ "default": 0,
+ "min": 0,
+ "max": 4294967295,
+ "display": "random",
+ "required_block_params": ["generator"],
+ },
+ "num_inference_steps": {
+ "label": "Steps",
+ "type": "int",
+ "default": 25,
+ "min": 1,
+ "max": 100,
+ "display": "slider",
+ "required_block_params": ["num_inference_steps"],
+ },
+ "num_frames": {
+ "label": "Frames",
+ "type": "int",
+ "default": 81,
+ "min": 1,
+ "max": 480,
+ "display": "slider",
+ "required_block_params": ["num_frames"],
+ },
+ "layers": {
+ "label": "Layers",
+ "type": "int",
+ "default": 4,
+ "min": 1,
+ "max": 10,
+ "display": "slider",
+ "required_block_params": ["layers"],
+ },
+ "output_type": {
+ "label": "Output Type",
+ "type": "dropdown",
+ "default": "np",
+ "options": ["np", "pil", "pt"],
+ },
+ # ControlNet
+ "controlnet_conditioning_scale": {
+ "label": "Controlnet Conditioning Scale",
+ "type": "float",
+ "default": 0.5,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.01,
+ "required_block_params": ["controlnet_conditioning_scale"],
+ },
+ "control_guidance_start": {
+ "label": "Control Guidance Start",
+ "type": "float",
+ "default": 0.0,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.01,
+ "required_block_params": ["control_guidance_start"],
+ },
+ "control_guidance_end": {
+ "label": "Control Guidance End",
+ "type": "float",
+ "default": 1.0,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.01,
+ "required_block_params": ["control_guidance_end"],
+ },
+ # Video
+ "videos": {"label": "Videos", "type": "video", "display": "output", "required_block_params": ["videos"]},
+ # Models
+ "vae": {"label": "VAE", "type": "diffusers_auto_model", "display": "input", "required_block_params": ["vae"]},
+ "image_encoder": {
+ "label": "Image Encoder",
+ "type": "diffusers_auto_model",
+ "display": "input",
+ "required_block_params": ["image_encoder"],
+ },
+ "unet": {"label": "Denoise Model", "type": "diffusers_auto_model", "display": "input"},
+ "scheduler": {"label": "Scheduler", "type": "diffusers_auto_model", "display": "input"},
+ "controlnet": {
+ "label": "ControlNet Model",
+ "type": "diffusers_auto_model",
+ "display": "input",
+ "required_block_params": ["controlnet"],
+ },
+ "text_encoders": {
+ "label": "Text Encoders",
+ "type": "diffusers_auto_models",
+ "display": "input",
+ "required_block_params": ["text_encoder"],
+ },
+ # Bundles/Custom
+ "controlnet_bundle": {
+ "label": "ControlNet",
+ "type": "custom_controlnet",
+ "display": "input",
+ "required_block_params": "controlnet_image",
+ },
+ "ip_adapter": {"label": "IP Adapter", "type": "custom_ip_adapter", "display": "input"},
+ "guider": {
+ "label": "Guider",
+ "type": "custom_guider",
+ "display": "input",
+ "onChange": {False: ["guidance_scale"], True: []},
+ },
+ "doc": {"label": "Doc", "type": "string", "display": "output"},
+}
+
+
+class MellonParamMeta(type):
+ """Metaclass that enables MellonParam.template_name(**overrides) syntax."""
+
+ def __getattr__(cls, name: str):
+ if name in MELLON_PARAM_TEMPLATES:
+
+ def factory(default=None, **overrides):
+ template = MELLON_PARAM_TEMPLATES[name]
+ # Use template's name if specified, otherwise use the key
+ params = {"name": template.get("name", name), **template, **overrides}
+ if default is not None:
+ params["default"] = default
+ return cls(**params)
+
+ return factory
+
+ raise AttributeError(f"type object 'MellonParam' has no attribute '{name}'")
+
+
+@dataclass(frozen=True)
+class MellonParam(metaclass=MellonParamMeta):
+ """
+ Parameter definition for Mellon nodes.
+
+ Usage:
+ ```python
+ # From template (standard diffuser params)
+ MellonParam.seed()
+ MellonParam.prompt(default="a cat")
+ MellonParam.latents(display="output")
+
+ # Generic inputs (for custom blocks)
+ MellonParam.Input.slider("my_scale", default=1.0, min=0.0, max=2.0)
+ MellonParam.Input.dropdown("mode", options=["fast", "slow"])
+
+ # Generic outputs
+ MellonParam.Output.image("result_images")
+
+ # Fully custom
+ MellonParam(name="custom", label="Custom", type="float", default=0.5)
+ ```
+ """
+
+ name: str
+ label: str
+ type: str
+ display: str | None = None
+ default: Any = None
+ min: float | None = None
+ max: float | None = None
+ step: float | None = None
+ options: Any = None
+ value: Any = None
+ fieldOptions: dict[str, Any] | None = None
+ onChange: Any = None
+ onSignal: Any = None
+ required_block_params: str | list[str] | None = None
+
+ def to_dict(self) -> dict[str, Any]:
+ """Convert to dict for Mellon schema, excluding None values and internal fields."""
+ data = asdict(self)
+ return {k: v for k, v in data.items() if v is not None and k not in ("name", "required_block_params")}
+
+ # =========================================================================
+ # Input: Generic input parameter factories (for custom blocks)
+ # =========================================================================
+ class Input:
+ """input UI elements for custom blocks."""
+
+ @classmethod
+ def image(cls, name: str) -> "MellonParam":
+ """image input."""
+ return MellonParam(name=name, label=_name_to_label(name), type="image", display="input")
+
+ @classmethod
+ def textbox(cls, name: str, default: str = "") -> "MellonParam":
+ """text input as textarea."""
+ return MellonParam(
+ name=name, label=_name_to_label(name), type="string", display="textarea", default=default
+ )
+
+ @classmethod
+ def dropdown(cls, name: str, options: list[str] = None, default: str = None) -> "MellonParam":
+ """dropdown selection."""
+ if options and not default:
+ default = options[0]
+ if not default:
+ default = ""
+ if not options:
+ options = [default]
+ return MellonParam(name=name, label=_name_to_label(name), type="string", options=options, value=default)
+
+ @classmethod
+ def slider(
+ cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None
+ ) -> "MellonParam":
+ """slider input."""
+ is_float = isinstance(default, float) or (step is not None and isinstance(step, float))
+ param_type = "float" if is_float else "int"
+ if min is None:
+ min = default
+ if max is None:
+ max = default
+ if step is None:
+ step = 0.01 if is_float else 1
+ return MellonParam(
+ name=name,
+ label=_name_to_label(name),
+ type=param_type,
+ display="slider",
+ default=default,
+ min=min,
+ max=max,
+ step=step,
+ )
+
+ @classmethod
+ def number(
+ cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None
+ ) -> "MellonParam":
+ """number input (no slider)."""
+ is_float = isinstance(default, float) or (step is not None and isinstance(step, float))
+ param_type = "float" if is_float else "int"
+ return MellonParam(
+ name=name, label=_name_to_label(name), type=param_type, default=default, min=min, max=max, step=step
+ )
+
+ @classmethod
+ def seed(cls, name: str = "seed", default: int = 0) -> "MellonParam":
+ """seed input with randomize button."""
+ return MellonParam(
+ name=name,
+ label=_name_to_label(name),
+ type="int",
+ display="random",
+ default=default,
+ min=0,
+ max=4294967295,
+ )
+
+ @classmethod
+ def checkbox(cls, name: str, default: bool = False) -> "MellonParam":
+ """boolean checkbox."""
+ return MellonParam(name=name, label=_name_to_label(name), type="boolean", value=default)
+
+ @classmethod
+ def custom_type(cls, name: str, type: str) -> "MellonParam":
+ """custom type input for node connections."""
+ return MellonParam(name=name, label=_name_to_label(name), type=type, display="input")
+
+ @classmethod
+ def model(cls, name: str) -> "MellonParam":
+ """model input for diffusers components."""
+ return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="input")
+
+ # =========================================================================
+ # Output: Generic output parameter factories (for custom blocks)
+ # =========================================================================
+ class Output:
+ """output UI elements for custom blocks."""
+
+ @classmethod
+ def image(cls, name: str) -> "MellonParam":
+ """image output."""
+ return MellonParam(name=name, label=_name_to_label(name), type="image", display="output")
+
+ @classmethod
+ def video(cls, name: str) -> "MellonParam":
+ """video output."""
+ return MellonParam(name=name, label=_name_to_label(name), type="video", display="output")
+
+ @classmethod
+ def text(cls, name: str) -> "MellonParam":
+ """text output."""
+ return MellonParam(name=name, label=_name_to_label(name), type="string", display="output")
+
+ @classmethod
+ def custom_type(cls, name: str, type: str) -> "MellonParam":
+ """custom type output for node connections."""
+ return MellonParam(name=name, label=_name_to_label(name), type=type, display="output")
+
+ @classmethod
+ def model(cls, name: str) -> "MellonParam":
+ """model output for diffusers components."""
+ return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="output")
+
+
+def input_param_to_mellon_param(input_param: "InputParam") -> MellonParam:
+ """
+ Convert an InputParam to a MellonParam using metadata.
+
+ Args:
+ input_param: An InputParam with optional metadata containing either:
+ - {"mellon": ""} for simple types (image, textbox, slider, etc.)
+ - {"mellon": MellonParam(...)} for full control over UI configuration
+
+ Returns:
+ MellonParam instance
+ """
+ name = input_param.name
+ metadata = input_param.metadata
+ mellon_value = metadata.get("mellon") if metadata else None
+ default = input_param.default
+
+ # If it's already a MellonParam, return it directly
+ if isinstance(mellon_value, MellonParam):
+ return mellon_value
+
+ mellon_type = mellon_value
+
+ if mellon_type == "image":
+ return MellonParam.Input.image(name)
+ elif mellon_type == "textbox":
+ return MellonParam.Input.textbox(name, default=default or "")
+ elif mellon_type == "dropdown":
+ return MellonParam.Input.dropdown(name, default=default or "")
+ elif mellon_type == "slider":
+ return MellonParam.Input.slider(name, default=default or 0)
+ elif mellon_type == "number":
+ return MellonParam.Input.number(name, default=default or 0)
+ elif mellon_type == "seed":
+ return MellonParam.Input.seed(name, default=default or 0)
+ elif mellon_type == "checkbox":
+ return MellonParam.Input.checkbox(name, default=default or False)
+ elif mellon_type == "model":
+ return MellonParam.Input.model(name)
+ else:
+ # None or unknown -> custom
+ return MellonParam.Input.custom_type(name, type="custom")
+
+
+def output_param_to_mellon_param(output_param: "OutputParam") -> MellonParam:
+ """
+ Convert an OutputParam to a MellonParam using metadata.
+
+ Args:
+ output_param: An OutputParam with optional metadata={"mellon": ""} where type is one of:
+ image, video, text, model. If metadata is None or unknown, maps to "custom".
+
+ Returns:
+ MellonParam instance
+ """
+ name = output_param.name
+ metadata = output_param.metadata
+ mellon_type = metadata.get("mellon") if metadata else None
+
+ if mellon_type == "image":
+ return MellonParam.Output.image(name)
+ elif mellon_type == "video":
+ return MellonParam.Output.video(name)
+ elif mellon_type == "text":
+ return MellonParam.Output.text(name)
+ elif mellon_type == "model":
+ return MellonParam.Output.model(name)
+ else:
+ # None or unknown -> custom
+ return MellonParam.Output.custom_type(name, type="custom")
+
+
+DEFAULT_NODE_SPECS = {
+ "controlnet": None,
+ "denoise": {
+ "inputs": [
+ MellonParam.embeddings(display="input"),
+ MellonParam.width(),
+ MellonParam.height(),
+ MellonParam.seed(),
+ MellonParam.num_inference_steps(),
+ MellonParam.num_frames(),
+ MellonParam.guidance_scale(),
+ MellonParam.strength(),
+ MellonParam.image_latents_with_strength(),
+ MellonParam.image_latents(),
+ MellonParam.first_frame_latents(),
+ MellonParam.controlnet_bundle(display="input"),
+ ],
+ "model_inputs": [
+ MellonParam.unet(),
+ MellonParam.guider(),
+ MellonParam.scheduler(),
+ ],
+ "outputs": [
+ MellonParam.latents(display="output"),
+ MellonParam.latents_preview(),
+ MellonParam.doc(),
+ ],
+ "required_inputs": ["embeddings"],
+ "required_model_inputs": ["unet", "scheduler"],
+ "block_name": "denoise",
+ },
+ "vae_encoder": {
+ "inputs": [
+ MellonParam.image(),
+ ],
+ "model_inputs": [
+ MellonParam.vae(),
+ ],
+ "outputs": [
+ MellonParam.image_latents(display="output"),
+ MellonParam.doc(),
+ ],
+ "required_inputs": ["image"],
+ "required_model_inputs": ["vae"],
+ "block_name": "vae_encoder",
+ },
+ "text_encoder": {
+ "inputs": [
+ MellonParam.prompt(),
+ MellonParam.negative_prompt(),
+ ],
+ "model_inputs": [
+ MellonParam.text_encoders(),
+ ],
+ "outputs": [
+ MellonParam.embeddings(display="output"),
+ MellonParam.doc(),
+ ],
+ "required_inputs": ["prompt"],
+ "required_model_inputs": ["text_encoders"],
+ "block_name": "text_encoder",
+ },
+ "decoder": {
+ "inputs": [
+ MellonParam.latents(display="input"),
+ ],
+ "model_inputs": [
+ MellonParam.vae(),
+ ],
+ "outputs": [
+ MellonParam.images(),
+ MellonParam.videos(),
+ MellonParam.doc(),
+ ],
+ "required_inputs": ["latents"],
+ "required_model_inputs": ["vae"],
+ "block_name": "decode",
+ },
+}
+
+
+def mark_required(label: str, marker: str = " *") -> str:
+ """Add required marker to label if not already present."""
+ if label.endswith(marker):
+ return label
+ return f"{label}{marker}"
+
+
+def node_spec_to_mellon_dict(node_spec: dict[str, Any], node_type: str) -> dict[str, Any]:
+ """
+ Convert a node spec dict into Mellon format.
+
+ A node spec is how we define a Mellon diffusers node in code. This function converts it into the `params` map
+ format that Mellon UI expects.
+
+ The `params` map is a dict where keys are parameter names and values are UI configuration:
+ ```python
+ {"seed": {"label": "Seed", "type": "int", "default": 0}}
+ ```
+
+ For Modular Mellon nodes, we need to distinguish:
+ - `inputs`: Pipeline inputs (e.g., seed, prompt, image)
+ - `model_inputs`: Model components (e.g., unet, vae, scheduler)
+ - `outputs`: Node outputs (e.g., latents, images)
+
+ The node spec also includes:
+ - `required_inputs` / `required_model_inputs`: Which params are required (marked with *)
+ - `block_name`: The modular pipeline block this node corresponds to on backend
+
+ We provide factory methods for common parameters (e.g., `MellonParam.seed()`, `MellonParam.unet()`) so you don't
+ have to manually specify all the UI configuration.
+
+ Args:
+ node_spec: Dict with `inputs`, `model_inputs`, `outputs` (lists of MellonParam),
+ plus `required_inputs`, `required_model_inputs`, `block_name`.
+ node_type: The node type string (e.g., "denoise", "controlnet")
+
+ Returns:
+ Dict with:
+ - `params`: Flat dict of all params in Mellon UI format
+ - `input_names`: List of input parameter names
+ - `model_input_names`: List of model input parameter names
+ - `output_names`: List of output parameter names
+ - `block_name`: The backend block name
+ - `node_type`: The node type
+
+ Example:
+ ```python
+ node_spec = {
+ "inputs": [MellonParam.seed(), MellonParam.prompt()],
+ "model_inputs": [MellonParam.unet()],
+ "outputs": [MellonParam.latents(display="output")],
+ "required_inputs": ["prompt"],
+ "required_model_inputs": ["unet"],
+ "block_name": "denoise",
+ }
+
+ result = node_spec_to_mellon_dict(node_spec, "denoise")
+ # Returns:
+ # {
+ # "params": {
+ # "seed": {"label": "Seed", "type": "int", "default": 0},
+ # "prompt": {"label": "Prompt *", "type": "string", "default": ""}, # * marks required
+ # "unet": {"label": "Denoise Model *", "type": "diffusers_auto_model", "display": "input"},
+ # "latents": {"label": "Latents", "type": "latents", "display": "output"},
+ # },
+ # "input_names": ["seed", "prompt"],
+ # "model_input_names": ["unet"],
+ # "output_names": ["latents"],
+ # "block_name": "denoise",
+ # "node_type": "denoise",
+ # }
+ ```
+ """
+ params = {}
+ input_names = []
+ model_input_names = []
+ output_names = []
+
+ required_inputs = node_spec.get("required_inputs", [])
+ required_model_inputs = node_spec.get("required_model_inputs", [])
+
+ # Process inputs
+ for p in node_spec.get("inputs", []):
+ param_dict = p.to_dict()
+ if p.name in required_inputs:
+ param_dict["label"] = mark_required(param_dict["label"])
+ params[p.name] = param_dict
+ input_names.append(p.name)
+
+ # Process model_inputs
+ for p in node_spec.get("model_inputs", []):
+ param_dict = p.to_dict()
+ if p.name in required_model_inputs:
+ param_dict["label"] = mark_required(param_dict["label"])
+ params[p.name] = param_dict
+ model_input_names.append(p.name)
+
+ # Process outputs: add a prefix to the output name if it already exists as an input
+ for p in node_spec.get("outputs", []):
+ if p.name in input_names:
+ # rename to out_
+ output_name = f"out_{p.name}"
+ else:
+ output_name = p.name
+ params[output_name] = p.to_dict()
+ output_names.append(output_name)
+
+ return {
+ "params": params,
+ "input_names": input_names,
+ "model_input_names": model_input_names,
+ "output_names": output_names,
+ "block_name": node_spec.get("block_name"),
+ "node_type": node_type,
+ }
+
+
+class MellonPipelineConfig:
+ """
+ Configuration for an entire Mellon pipeline containing multiple nodes.
+
+ Accepts node specs as dicts with inputs/model_inputs/outputs lists of MellonParam, converts them to Mellon-ready
+ format, and handles save/load to Hub.
+
+ Example:
+ ```python
+ config = MellonPipelineConfig(
+ node_specs={
+ "denoise": {
+ "inputs": [MellonParam.seed(), MellonParam.prompt()],
+ "model_inputs": [MellonParam.unet()],
+ "outputs": [MellonParam.latents(display="output")],
+ "required_inputs": ["prompt"],
+ "required_model_inputs": ["unet"],
+ "block_name": "denoise",
+ },
+ "decoder": {
+ "inputs": [MellonParam.latents(display="input")],
+ "outputs": [MellonParam.images()],
+ "block_name": "decoder",
+ },
+ },
+ label="My Pipeline",
+ default_repo="user/my-pipeline",
+ default_dtype="float16",
+ )
+
+ # Access Mellon format dict
+ denoise = config.node_params["denoise"]
+ input_names = denoise["input_names"]
+ params = denoise["params"]
+
+ # Save to Hub
+ config.save("./my_config", push_to_hub=True, repo_id="user/my-pipeline")
+
+ # Load from Hub
+ loaded = MellonPipelineConfig.load("user/my-pipeline")
+ ```
+ """
+
+ config_name = "mellon_pipeline_config.json"
+
+ def __init__(
+ self,
+ node_specs: dict[str, dict[str, Any] | None],
+ label: str = "",
+ default_repo: str = "",
+ default_dtype: str = "",
+ ):
+ """
+ Args:
+ node_specs: Dict mapping node_type to node spec or None.
+ Node spec has: inputs, model_inputs, outputs, required_inputs, required_model_inputs,
+ block_name (all optional)
+ label: Human-readable label for the pipeline
+ default_repo: Default HuggingFace repo for this pipeline
+ default_dtype: Default dtype (e.g., "float16", "bfloat16")
+ """
+ # Convert all node specs to Mellon format immediately
+ self.node_specs = node_specs
+
+ self.label = label
+ self.default_repo = default_repo
+ self.default_dtype = default_dtype
+
+ @property
+ def node_params(self) -> dict[str, Any]:
+ """Lazily compute node_params from node_specs."""
+ if self.node_specs is None:
+ return self._node_params
+
+ params = {}
+ for node_type, spec in self.node_specs.items():
+ if spec is None:
+ params[node_type] = None
+ else:
+ params[node_type] = node_spec_to_mellon_dict(spec, node_type)
+ return params
+
+ def __repr__(self) -> str:
+ lines = [
+ f"MellonPipelineConfig(label={self.label!r}, default_repo={self.default_repo!r}, default_dtype={self.default_dtype!r})"
+ ]
+ for node_type, spec in self.node_specs.items():
+ if spec is None:
+ lines.append(f" {node_type}: None")
+ else:
+ inputs = [p.name for p in spec.get("inputs", [])]
+ model_inputs = [p.name for p in spec.get("model_inputs", [])]
+ outputs = [p.name for p in spec.get("outputs", [])]
+ lines.append(f" {node_type}:")
+ lines.append(f" inputs: {inputs}")
+ lines.append(f" model_inputs: {model_inputs}")
+ lines.append(f" outputs: {outputs}")
+ return "\n".join(lines)
+
+ def to_dict(self) -> dict[str, Any]:
+ """Convert to a JSON-serializable dictionary."""
+ return {
+ "label": self.label,
+ "default_repo": self.default_repo,
+ "default_dtype": self.default_dtype,
+ "node_params": self.node_params,
+ }
+
+ @classmethod
+ def from_dict(cls, data: dict[str, Any]) -> "MellonPipelineConfig":
+ """
+ Create from a dictionary (loaded from JSON).
+
+ Note: The mellon_params are already in Mellon format when loading from JSON.
+ """
+ instance = cls.__new__(cls)
+ instance.node_specs = None
+ instance._node_params = data.get("node_params", {})
+ instance.label = data.get("label", "")
+ instance.default_repo = data.get("default_repo", "")
+ instance.default_dtype = data.get("default_dtype", "")
+ return instance
+
+ def to_json_string(self) -> str:
+ """Serialize to JSON string."""
+ return json.dumps(self.to_dict(), indent=2, sort_keys=False) + "\n"
+
+ def to_json_file(self, json_file_path: str | os.PathLike):
+ """Save to a JSON file."""
+ with open(json_file_path, "w", encoding="utf-8") as writer:
+ writer.write(self.to_json_string())
+
+ @classmethod
+ def from_json_file(cls, json_file_path: str | os.PathLike) -> "MellonPipelineConfig":
+ """Load from a JSON file."""
+ with open(json_file_path, "r", encoding="utf-8") as reader:
+ data = json.load(reader)
+ return cls.from_dict(data)
+
+ def save(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
+ """Save the mellon pipeline config to a directory."""
+ if os.path.isfile(save_directory):
+ raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+ os.makedirs(save_directory, exist_ok=True)
+ output_path = os.path.join(save_directory, self.config_name)
+ self.to_json_file(output_path)
+ logger.info(f"Pipeline config saved to {output_path}")
+
+ if push_to_hub:
+ commit_message = kwargs.pop("commit_message", None)
+ private = kwargs.pop("private", None)
+ create_pr = kwargs.pop("create_pr", False)
+ token = kwargs.pop("token", None)
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+ repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
+
+ upload_file(
+ path_or_fileobj=output_path,
+ path_in_repo=self.config_name,
+ repo_id=repo_id,
+ token=token,
+ commit_message=commit_message or "Upload MellonPipelineConfig",
+ create_pr=create_pr,
+ )
+ logger.info(f"Pipeline config pushed to hub: {repo_id}")
+
+ @classmethod
+ def load(
+ cls,
+ pretrained_model_name_or_path: str | os.PathLike,
+ **kwargs,
+ ) -> "MellonPipelineConfig":
+ """Load a pipeline config from a local path or Hugging Face Hub."""
+ cache_dir = kwargs.pop("cache_dir", None)
+ local_dir = kwargs.pop("local_dir", None)
+ local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
+ force_download = kwargs.pop("force_download", False)
+ proxies = kwargs.pop("proxies", None)
+ token = kwargs.pop("token", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ revision = kwargs.pop("revision", None)
+ subfolder = kwargs.pop("subfolder", None)
+
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+
+ if os.path.isfile(pretrained_model_name_or_path):
+ config_file = pretrained_model_name_or_path
+ elif os.path.isdir(pretrained_model_name_or_path):
+ config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
+ if not os.path.isfile(config_file):
+ raise EnvironmentError(f"No file named {cls.config_name} found in {pretrained_model_name_or_path}")
+ else:
+ try:
+ config_file = hf_hub_download(
+ pretrained_model_name_or_path,
+ filename=cls.config_name,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ subfolder=subfolder,
+ local_dir=local_dir,
+ local_dir_use_symlinks=local_dir_use_symlinks,
+ )
+ except RepositoryNotFoundError:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
+ " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
+ " token having permission to this repo with `token` or log in with `hf auth login`."
+ )
+ except RevisionNotFoundError:
+ raise EnvironmentError(
+ f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
+ " this model name. Check the model page at"
+ f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
+ )
+ except EntryNotFoundError:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
+ )
+ except HfHubHTTPError as err:
+ raise EnvironmentError(
+ "There was a specific connection error when trying to load"
+ f" {pretrained_model_name_or_path}:\n{err}"
+ )
+ except ValueError:
+ raise EnvironmentError(
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
+ f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
+ f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
+ " run the library in offline mode at"
+ " 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
+ )
+ except EnvironmentError:
+ raise EnvironmentError(
+ f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
+ "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
+ f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
+ f"containing a {cls.config_name} file"
+ )
+
+ try:
+ return cls.from_json_file(config_file)
+ except (json.JSONDecodeError, UnicodeDecodeError):
+ raise EnvironmentError(f"The config file at '{config_file}' is not a valid JSON file.")
+
+ @classmethod
+ def from_blocks(
+ cls,
+ blocks,
+ template: dict[str, dict[str, Any]] | None = None,
+ label: str = "",
+ default_repo: str = "",
+ default_dtype: str = "bfloat16",
+ ) -> "MellonPipelineConfig":
+ """
+ Create MellonPipelineConfig by matching template against actual pipeline blocks.
+ """
+ if template is None:
+ template = DEFAULT_NODE_SPECS
+
+ sub_block_map = dict(blocks.sub_blocks)
+
+ def filter_spec_for_block(template_spec: dict[str, Any], block) -> dict[str, Any] | None:
+ """Filter template spec params based on what the block actually supports."""
+ block_input_names = set(block.input_names)
+ block_output_names = set(block.intermediate_output_names)
+ block_component_names = set(block.component_names)
+
+ filtered_inputs = [
+ p
+ for p in template_spec.get("inputs", [])
+ if p.required_block_params is None
+ or all(name in block_input_names for name in p.required_block_params)
+ ]
+ filtered_model_inputs = [
+ p
+ for p in template_spec.get("model_inputs", [])
+ if p.required_block_params is None
+ or all(name in block_component_names for name in p.required_block_params)
+ ]
+ filtered_outputs = [
+ p
+ for p in template_spec.get("outputs", [])
+ if p.required_block_params is None
+ or all(name in block_output_names for name in p.required_block_params)
+ ]
+
+ filtered_input_names = {p.name for p in filtered_inputs}
+ filtered_model_input_names = {p.name for p in filtered_model_inputs}
+
+ filtered_required_inputs = [
+ r for r in template_spec.get("required_inputs", []) if r in filtered_input_names
+ ]
+ filtered_required_model_inputs = [
+ r for r in template_spec.get("required_model_inputs", []) if r in filtered_model_input_names
+ ]
+
+ return {
+ "inputs": filtered_inputs,
+ "model_inputs": filtered_model_inputs,
+ "outputs": filtered_outputs,
+ "required_inputs": filtered_required_inputs,
+ "required_model_inputs": filtered_required_model_inputs,
+ "block_name": template_spec.get("block_name"),
+ }
+
+ # Build node specs
+ node_specs = {}
+ for node_type, template_spec in template.items():
+ if template_spec is None:
+ node_specs[node_type] = None
+ continue
+
+ block_name = template_spec.get("block_name")
+ if block_name is None or block_name not in sub_block_map:
+ node_specs[node_type] = None
+ continue
+
+ node_specs[node_type] = filter_spec_for_block(template_spec, sub_block_map[block_name])
+
+ return cls(
+ node_specs=node_specs,
+ label=label or getattr(blocks, "model_name", ""),
+ default_repo=default_repo,
+ default_dtype=default_dtype,
+ )
+
+ @classmethod
+ def from_custom_block(
+ cls,
+ block,
+ node_label: str = None,
+ input_types: dict[str, Any] | None = None,
+ output_types: dict[str, Any] | None = None,
+ ) -> "MellonPipelineConfig":
+ """
+ Create a MellonPipelineConfig from a custom block.
+
+ Args:
+ block: A block instance with `inputs`, `outputs`, and `expected_components`/`component_names` properties.
+ Each InputParam/OutputParam should have metadata={"mellon": ""} where type is one of: image,
+ video, text, checkbox, number, slider, dropdown, model. If metadata is None, maps to "custom".
+ node_label: The display label for the node. Defaults to block class name with spaces.
+ input_types:
+ Optional dict mapping input param names to mellon types. Overrides the block's metadata if provided.
+ Example: {"prompt": "textbox", "image": "image"}
+ output_types:
+ Optional dict mapping output param names to mellon types. Overrides the block's metadata if provided.
+ Example: {"prompt": "text", "images": "image"}
+
+ Returns:
+ MellonPipelineConfig instance
+ """
+ if node_label is None:
+ class_name = block.__class__.__name__
+ node_label = "".join([" " + c if c.isupper() else c for c in class_name]).strip()
+
+ if input_types is None:
+ input_types = {}
+ if output_types is None:
+ output_types = {}
+
+ inputs = []
+ model_inputs = []
+ outputs = []
+
+ # Process block inputs
+ for input_param in block.inputs:
+ if input_param.name is None:
+ continue
+ if input_param.name in input_types:
+ input_param = copy.copy(input_param)
+ input_param.metadata = {"mellon": input_types[input_param.name]}
+ print(f" processing input: {input_param.name}, metadata: {input_param.metadata}")
+ inputs.append(input_param_to_mellon_param(input_param))
+
+ # Process block outputs
+ for output_param in block.outputs:
+ if output_param.name is None:
+ continue
+ if output_param.name in output_types:
+ output_param = copy.copy(output_param)
+ output_param.metadata = {"mellon": output_types[output_param.name]}
+ outputs.append(output_param_to_mellon_param(output_param))
+
+ # Process expected components (all map to model inputs)
+ component_names = block.component_names
+ for component_name in component_names:
+ model_inputs.append(MellonParam.Input.model(component_name))
+
+ # Always add doc output
+ outputs.append(MellonParam.doc())
+
+ node_spec = {
+ "inputs": inputs,
+ "model_inputs": model_inputs,
+ "outputs": outputs,
+ "required_inputs": [],
+ "required_model_inputs": [],
+ "block_name": "custom",
+ }
+
+ return cls(
+ node_specs={"custom": node_spec},
+ label=node_label,
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ff94fd8a93959a115d8f87976101536e61c337d7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/__init__.py
@@ -0,0 +1,59 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_minimax_h3"] = ["MiniMaxH3Blocks"]
+ _import_structure["modular_pipeline"] = ["MiniMaxH3ModularPipeline"]
+ _import_structure["references"] = [
+ "MiniMaxH3AudioReference",
+ "MiniMaxH3ImageReference",
+ "MiniMaxH3Reference",
+ "MiniMaxH3VideoReference",
+ ]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_minimax_h3 import MiniMaxH3Blocks
+ from .modular_pipeline import MiniMaxH3ModularPipeline
+ from .references import (
+ MiniMaxH3AudioReference,
+ MiniMaxH3ImageReference,
+ MiniMaxH3Reference,
+ MiniMaxH3VideoReference,
+ )
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef9ba4e982e6e594ab6c5795d7723d5df78d0971
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py
@@ -0,0 +1,1039 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+
+from ...schedulers import MiniMaxH3Scheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import (
+ MINIMAX_H3_AUDIO_CHANNELS,
+ MINIMAX_H3_AUDIO_TAG,
+ MINIMAX_H3_VIDEO_TAG,
+ MiniMaxH3ModularPipeline,
+ align_num_frames,
+ audio_latent_num_frames,
+ resolve_canvas_size,
+ video_latent_num_frames,
+)
+from .references import MiniMaxH3Reference
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Rotary-time constants. One latent frame spans `5/3 * frames_per_latent` rotary units, where the pattern
+# `(1, 4, 4, 4, 4)` mirrors the VAE's 17-pixel-frames-to-5-latent-frames grouping; the spatial axes are normalized
+# by the square root of the latent area and scaled by 32.
+_ROPE_FRAME_RESCALE = 5.0 / 3.0
+_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4)
+_ROPE_SPATIAL_SCALE = 32
+
+
+def patchify_video_latents(latents: torch.Tensor, patch_size: tuple[int, int, int]) -> torch.Tensor:
+ r"""
+ Pack video latents into transformer rows.
+
+ Args:
+ latents (`torch.Tensor` of shape `(batch_size, channels, num_frames, height, width)`):
+ The latents to pack.
+ patch_size (`tuple[int, int, int]`): The `(t, h, w)` patch.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size * num_patches, channels * prod(patch_size))`: The packed rows, ordered
+ frame-major then row-major.
+ """
+ patch_t, patch_h, patch_w = patch_size
+ batch_size, channels, num_frames, height, width = latents.shape
+ if num_frames % patch_t or height % patch_h or width % patch_w:
+ raise ValueError(f"Latents of shape {tuple(latents.shape)} are not divisible by the patch {patch_size}.")
+
+ latents = latents.reshape(
+ batch_size,
+ channels,
+ num_frames // patch_t,
+ patch_t,
+ height // patch_h,
+ patch_h,
+ width // patch_w,
+ patch_w,
+ )
+ latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7)
+ return latents.reshape(-1, channels * patch_t * patch_h * patch_w).contiguous()
+
+
+def _spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor:
+ r"""
+ One aspect-normalized spatial rotary axis: `dim // patch` coordinates centred on the unit interval, scaled up by
+ 32. The right endpoint is excluded, so a square canvas spans `[0, 32)`.
+ """
+ ratio = dim / sqrt_area
+ left = (1.0 - ratio) / 2.0
+ # Built with numpy: `np.linspace(..., endpoint=False)` is `start + arange(num) * (stop - start) / num`, which is
+ # not what `torch.linspace` computes, and the float64 grid has to be reproduced exactly.
+ grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE
+ return torch.from_numpy(grid).to(torch.float64)
+
+
+def _temporal_position_grid(num_latent_frames: int, origin: float) -> torch.Tensor:
+ r"""The rotary time of every latent frame, starting at `origin`. Spacing is non-uniform: `5/3 * (1, 4, 4, 4, 4)`."""
+ spans = torch.tensor(
+ [
+ _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[index % len(_ROPE_FRAMES_PER_LATENT)]
+ for index in range(num_latent_frames)
+ ],
+ dtype=torch.float64,
+ )
+ return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
+
+
+def _temporal_position_span_pairwise(num_latent_frames: int) -> float:
+ r"""
+ The rotary time spanned by `num_latent_frames` latent frames.
+
+ Summed by numpy (pairwise summation) rather than sequentially: the reference computes the keyframe anchor this
+ way and the two summation orders differ in the last ulp from 16 latent frames onwards.
+ """
+ spans = np.ones(num_latent_frames, dtype=np.float64) * _ROPE_FRAME_RESCALE
+ for index in range(len(_ROPE_FRAMES_PER_LATENT)):
+ spans[index :: len(_ROPE_FRAMES_PER_LATENT)] *= _ROPE_FRAMES_PER_LATENT[index]
+ return float(spans.sum())
+
+
+def _temporal_position_span_sequential(num_latent_frames: int) -> float:
+ r"""
+ The rotary time a video reference advances the clock by.
+
+ Summed sequentially in float64, which is *not* how [`_temporal_position_span_pairwise`] sums
+ the same series: that one reproduces a numpy pairwise sum, and the two orders differ in the last ulp from 16
+ latent frames onwards. The reference implementation keeps both, one per call site, so the port has to as well.
+ """
+ return sum(
+ _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[index % len(_ROPE_FRAMES_PER_LATENT)]
+ for index in range(num_latent_frames)
+ )
+
+
+def _frame_position_grid(
+ latent_height: int, latent_width: int, patch_h: int, patch_w: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+ r"""The `(h, w)` rotary coordinates of one latent frame, and the width axis they were built from."""
+ sqrt_area = np.sqrt(latent_height * latent_width)
+ height_grid = _spatial_position_grid(latent_height, patch_h, sqrt_area)
+ width_grid = _spatial_position_grid(latent_width, patch_w, sqrt_area)
+ grids = torch.meshgrid(height_grid, width_grid, indexing="ij")
+ return torch.stack([grid.reshape(-1) for grid in grids], dim=-1), width_grid
+
+
+def _fill_audio_positions(
+ position_ids: torch.Tensor,
+ rows: slice,
+ num_audio_latents: int,
+ rotary_time: float,
+ width_grid: torch.Tensor,
+) -> None:
+ r"""
+ Place one channel-major audio block.
+
+ Audio rows carry no height coordinate and are pinned to the two extremes of the width grid of *their own* block —
+ the target grid for a standalone audio reference, the video's grid for a soundtrack.
+ """
+ time = rotary_time + torch.arange(num_audio_latents, dtype=torch.float64)
+ position_ids[rows, 0] = time.repeat(MINIMAX_H3_AUDIO_CHANNELS)
+ position_ids[rows, 2] = torch.cat(
+ [
+ torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64),
+ torch.full((num_audio_latents,), float(width_grid[-1]), dtype=torch.float64),
+ ]
+ )
+
+
+class MiniMaxH3PrepareLayoutStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resolves the geometry of a `t2va` / `fl2va` request — the canvas, the `17 * n + 5` frame count the video "
+ "VAE can decode and the latent shapes every later block keys off — and builds the packed layout from it: "
+ "`[text | keyframe conditions | target audio | target video]` plus its fp64 rotary grid. MiniMax-H3 runs "
+ "full self-attention over this one sequence, so the layout is what every later block addresses rows "
+ "through."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="text_token_tags",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The per-row modality tag of every row of `prompt_embeds`.",
+ ),
+ InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+ InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ default=124,
+ description=(
+ "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video "
+ "VAE can decode; the resulting duration must stay between 5 and 15 seconds."
+ ),
+ ),
+ InputParam(
+ name="keyframe_anchors",
+ type_hint=tuple,
+ default=(),
+ description="Which end of the video every keyframe is anchored to, in packed order.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."),
+ OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."),
+ OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."),
+ OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."),
+ OutputParam(
+ "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel."
+ ),
+ OutputParam(
+ "position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The `(t, h, w)` rotary coordinate of every row, in float64.",
+ ),
+ OutputParam(
+ "token_tags",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The modality tag of every row.",
+ ),
+ OutputParam(
+ "video_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the video rows, conditioning rows first.",
+ ),
+ OutputParam(
+ "audio_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the audio rows, reference rows first.",
+ ),
+ OutputParam(
+ "text_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the text rows.",
+ ),
+ OutputParam(
+ "num_condition_video_rows",
+ type_hint=int,
+ description="How many leading video rows are conditioning rows rather than generated rows.",
+ ),
+ OutputParam(
+ "num_condition_audio_rows",
+ type_hint=int,
+ description="How many leading audio rows are reference rows rather than generated rows.",
+ ),
+ ]
+
+ @staticmethod
+ def build_packed_sequence(
+ text_token_tags: torch.Tensor,
+ num_latent_frames: int,
+ latent_height: int,
+ latent_width: int,
+ num_audio_latents: int,
+ patch_size: tuple[int, int, int],
+ keyframe_anchors: tuple[str, ...] = (),
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+ r"""
+ Build the `[text | keyframe conditions | target audio | target video]` layout used by the `t2va` and `fl2va`
+ tasks.
+
+ Args:
+ text_token_tags (`torch.Tensor` of shape `(num_text_tokens,)`):
+ The modality tag of every text row. Text is tagged `1`, except for the rows of a keyframe's vision block,
+ which MiniMax-H3 tags `0` (video).
+ num_latent_frames (`int`): Number of target latent frames.
+ latent_height (`int`): Target latent height.
+ latent_width (`int`): Target latent width.
+ num_audio_latents (`int`): Number of target audio latents per channel.
+ patch_size (`tuple[int, int, int]`): The transformer's `(t, h, w)` patch.
+ keyframe_anchors (`tuple[str, ...]`):
+ One entry per keyframe conditioning block, in packed order: `"first"` anchors the block at the first
+ latent frame, `"last"` at the last one.
+
+ Returns:
+ `tuple`: `position_ids`, `token_tags`, `video_indices`, `audio_indices`, `text_indices`, and the number of
+ leading video and audio rows that are conditioning rather than generated.
+ """
+ _, patch_h, patch_w = patch_size
+ rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w)
+ num_text_tokens = text_token_tags.shape[0]
+ num_condition_rows = len(keyframe_anchors) * rows_per_frame
+ num_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
+ num_video_rows = num_latent_frames * rows_per_frame
+ sequence_length = num_text_tokens + num_condition_rows + num_audio_rows + num_video_rows
+
+ condition_start = num_text_tokens
+ audio_start = condition_start + num_condition_rows
+ video_start = audio_start + num_audio_rows
+
+ # 1. The (t, h, w) grid. Text rows sit on the time axis at their row index, and the media rows continue the time
+ # axis from there, so text length shifts the whole media clock.
+ position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64)
+ position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64)
+
+ sqrt_area = np.sqrt(latent_height * latent_width)
+ height_grid = _spatial_position_grid(latent_height, patch_h, sqrt_area)
+ width_grid = _spatial_position_grid(latent_width, patch_w, sqrt_area)
+ frame_grid = torch.stack(
+ [grid.reshape(-1) for grid in torch.meshgrid(height_grid, width_grid, indexing="ij")], -1
+ )
+
+ for index, anchor in enumerate(keyframe_anchors):
+ if anchor == "first":
+ anchor_time = float(num_text_tokens)
+ elif anchor == "last":
+ anchor_time = (
+ float(num_text_tokens) + _temporal_position_span_pairwise(num_latent_frames) - _ROPE_FRAME_RESCALE
+ )
+ else:
+ raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}.")
+ rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame)
+ position_ids[rows, 0] = anchor_time
+ position_ids[rows, 1:] = frame_grid
+
+ # Audio rows are channel-major and share the video's rotary clock: one unit per latent at 40 latents/s equals
+ # 24 fps * 5/3. They carry no height coordinate and are pinned to the two extremes of the width grid.
+ audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64)
+ position_ids[audio_start:video_start, 0] = audio_time.repeat(MINIMAX_H3_AUDIO_CHANNELS)
+ position_ids[audio_start:video_start, 2] = torch.cat(
+ [
+ torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64),
+ torch.full((num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64),
+ ]
+ )
+
+ video_position_ids = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64)
+ video_position_ids[:, :, 0] = _temporal_position_grid(num_latent_frames, float(num_text_tokens))[:, None]
+ video_position_ids[:, :, 1:] = frame_grid[None]
+ position_ids[video_start:] = video_position_ids.reshape(-1, 3)
+
+ # 2. Row indices and modality tags.
+ video_indices = torch.cat(
+ [torch.arange(condition_start, audio_start), torch.arange(video_start, sequence_length)]
+ )
+ audio_indices = torch.arange(audio_start, video_start)
+ text_indices = torch.arange(num_text_tokens)
+
+ token_tags = torch.empty(sequence_length, dtype=torch.long)
+ token_tags[text_indices] = text_token_tags.to(torch.long)
+ token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG
+ token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG
+
+ return position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, 0
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ # Without a keyframe to take the aspect ratio from, MiniMax-H3 generates on its own 16:9 canvas.
+ if block_state.height is None:
+ block_state.height, block_state.width = resolve_canvas_size(16, 9, components.canvas_multiple)
+ if block_state.height % components.canvas_multiple or block_state.width % components.canvas_multiple:
+ raise ValueError(
+ f"`height` and `width` must be multiples of {components.canvas_multiple}, got "
+ f"{block_state.height}x{block_state.width}."
+ )
+
+ frames_per_chunk = components.vae_frames_per_chunk
+ latents_per_chunk = components.vae_latents_per_chunk
+ aligned_num_frames = align_num_frames(block_state.num_frames, frames_per_chunk, latents_per_chunk)
+ if aligned_num_frames != block_state.num_frames:
+ logger.warning(
+ f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} "
+ f"up to {aligned_num_frames}."
+ )
+ block_state.num_frames = aligned_num_frames
+ # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has
+ # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds.
+ duration = block_state.num_frames / components.fps
+ if not components.min_duration <= duration <= components.max_duration:
+ raise ValueError(
+ f"MiniMax-H3 generates between {components.min_duration} and {components.max_duration} seconds "
+ f"at {components.fps} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE "
+ f"can encode, must be between {int(components.min_duration * components.fps)} and "
+ f"{int(components.max_duration * components.fps)}, got {block_state.num_frames}."
+ )
+
+ ratio = components.vae_spatial_compression_ratio
+ block_state.num_latent_frames = video_latent_num_frames(
+ block_state.num_frames, frames_per_chunk, latents_per_chunk
+ )
+ block_state.latent_height = block_state.height // ratio
+ block_state.latent_width = block_state.width // ratio
+ block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames)
+
+ (
+ position_ids,
+ token_tags,
+ video_indices,
+ audio_indices,
+ text_indices,
+ block_state.num_condition_video_rows,
+ block_state.num_condition_audio_rows,
+ ) = self.build_packed_sequence(
+ block_state.text_token_tags,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ block_state.keyframe_anchors,
+ )
+ block_state.position_ids = position_ids.to(device)
+ block_state.token_tags = token_tags.to(device)
+ block_state.video_indices = video_indices.to(device)
+ block_state.audio_indices = audio_indices.to(device)
+ block_state.text_indices = text_indices.to(device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VAPrepareLayoutStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resolves the latent shapes of a `ref2va` request and builds its packed layout — `[text | reference "
+ "blocks | target audio | target video]` — plus its fp64 rotary grid. The reference order advances the "
+ "shared audio/video rotary clock, so it is part of the layout rather than a detail of the presentation."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="text_token_tags",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The per-row modality tag of every row of `prompt_embeds`.",
+ ),
+ InputParam(
+ name="normalized_references",
+ type_hint=list[MiniMaxH3Reference],
+ required=True,
+ description="The references normalized by the setup step, in packed order.",
+ ),
+ InputParam(
+ name="condition_latents",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description=(
+ "The encoded video conditioning latents, one per image and video reference in packed order. "
+ "Their shape is where every reference block's geometry comes from."
+ ),
+ ),
+ InputParam(
+ name="audio_condition_latents",
+ type_hint=list[torch.Tensor],
+ required=True,
+ description="The encoded audio conditioning rows, one per audio-bearing reference in packed order.",
+ ),
+ InputParam.template("height", required=True, description="Height of the generated video in pixels."),
+ InputParam.template("width", required=True, description="Width of the generated video in pixels."),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ required=True,
+ description="Resolved number of frames, of the form 17 * n + 5.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."),
+ OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."),
+ OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."),
+ OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."),
+ OutputParam(
+ "num_audio_latents", type_hint=int, description="Number of generated audio latents per channel."
+ ),
+ OutputParam(
+ "position_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The `(t, h, w)` rotary coordinate of every row, in float64.",
+ ),
+ OutputParam(
+ "token_tags",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The modality tag of every row.",
+ ),
+ OutputParam(
+ "video_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the video rows, conditioning rows first.",
+ ),
+ OutputParam(
+ "audio_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the audio rows, reference rows first.",
+ ),
+ OutputParam(
+ "text_indices",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="Sequence positions of the text rows.",
+ ),
+ OutputParam(
+ "num_condition_video_rows",
+ type_hint=int,
+ description="How many leading video rows are conditioning rows rather than generated rows.",
+ ),
+ OutputParam(
+ "num_condition_audio_rows",
+ type_hint=int,
+ description="How many leading audio rows are reference rows rather than generated rows.",
+ ),
+ ]
+
+ @staticmethod
+ def build_ref2va_packed_sequence(
+ text_token_tags: torch.Tensor,
+ references: list[MiniMaxH3Reference],
+ condition_latents: list[torch.Tensor],
+ audio_condition_latents: list[torch.Tensor],
+ num_latent_frames: int,
+ latent_height: int,
+ latent_width: int,
+ num_audio_latents: int,
+ patch_size: tuple[int, int, int],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+ r"""
+ Build the `[text | reference blocks | target audio | target video]` layout of the `ref2va` task.
+
+ Args:
+ text_token_tags (`torch.Tensor` of shape `(num_text_tokens,)`):
+ The modality tag of every text row. Text is tagged `1`, except for the rows of a reference's vision block,
+ which MiniMax-H3 tags `0` (video).
+ references (`list[MiniMaxH3Reference]`):
+ The references, in packed order. Only their modality is read here; the geometry comes from the latents.
+ condition_latents (`list[torch.Tensor]`):
+ One `(1, channels, num_latent_frames, latent_height, latent_width)` tensor per image and video reference,
+ in packed order, as [`MiniMaxH3Ref2VAReferenceEncoderStep`] produced them.
+ audio_condition_latents (`list[torch.Tensor]`):
+ One `(num_audio_latents * 2, audio_latent_channels)` tensor per audio-bearing reference, in packed order.
+ num_latent_frames (`int`): Number of target latent frames.
+ latent_height (`int`): Target latent height.
+ latent_width (`int`): Target latent width.
+ num_audio_latents (`int`): Number of target audio latents per channel.
+ patch_size (`tuple[int, int, int]`): The transformer's `(t, h, w)` patch.
+
+ Returns:
+ `tuple`: `position_ids`, `token_tags`, `video_indices`, `audio_indices`, `text_indices`, and the number of
+ leading video and audio rows that are references rather than generated.
+ """
+ _, patch_h, patch_w = patch_size
+ num_text_tokens = text_token_tags.shape[0]
+ num_target_video_rows = num_latent_frames * (latent_height // patch_h) * (latent_width // patch_w)
+ num_target_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
+
+ # The geometry of every reference block is the shape of what the encoder produced for it, so the two can never
+ # disagree. Both lists are in packed order but skip the references they do not apply to, so they are consumed as
+ # iterators alongside the reference list rather than indexed by it.
+ visual_geometry = iter(tuple(latents.shape[2:5]) for latents in condition_latents)
+ audio_row_counts = iter(rows.shape[0] for rows in audio_condition_latents)
+ num_reference_video_rows = sum(
+ frames * (height // patch_h) * (width // patch_w)
+ for frames, height, width in (tuple(latents.shape[2:5]) for latents in condition_latents)
+ )
+ num_reference_audio_rows = sum(rows.shape[0] for rows in audio_condition_latents)
+ sequence_length = (
+ num_text_tokens
+ + num_reference_video_rows
+ + num_reference_audio_rows
+ + num_target_audio_rows
+ + num_target_video_rows
+ )
+
+ position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64)
+ position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64)
+ target_frame_grid, target_width_grid = _frame_position_grid(latent_height, latent_width, patch_h, patch_w)
+
+ # Reference blocks, in request order. `rotary_time` is the shared audio/video clock: it starts where the text
+ # rows end and every block pushes it forward by the time that block occupies.
+ video_indices, audio_indices = [], []
+ cursor = num_text_tokens
+ rotary_time = float(num_text_tokens)
+ for reference in references:
+ if reference.kind == "image":
+ num_latent_frames_, reference_height, reference_width = next(visual_geometry)
+ num_video_rows = num_latent_frames_ * (reference_height // patch_h) * (reference_width // patch_w)
+ rows = slice(cursor, cursor + num_video_rows)
+ cursor = rows.stop
+ video_indices.append(torch.arange(rows.start, rows.stop))
+ frame_grid, _ = _frame_position_grid(reference_height, reference_width, patch_h, patch_w)
+ position_ids[rows, 0] = rotary_time
+ position_ids[rows, 1:] = frame_grid
+ # An image is a single frame and takes a single integer rotary slot, not a latent frame's 5/3 units.
+ rotary_time += 1.0
+ elif reference.kind == "audio":
+ num_audio_rows = next(audio_row_counts)
+ reference_audio_latents = num_audio_rows // MINIMAX_H3_AUDIO_CHANNELS
+ rows = slice(cursor, cursor + num_audio_rows)
+ cursor = rows.stop
+ audio_indices.append(torch.arange(rows.start, rows.stop))
+ _fill_audio_positions(position_ids, rows, reference_audio_latents, rotary_time, target_width_grid)
+ rotary_time += float(reference_audio_latents)
+ elif reference.kind == "video":
+ # A video reference's soundtrack rows are packed immediately before its video rows and share their
+ # origin, so the two are rotary-aligned exactly as the generated audio and video are.
+ num_audio_rows = next(audio_row_counts) if reference.has_audio else 0
+ reference_audio_latents = num_audio_rows // MINIMAX_H3_AUDIO_CHANNELS
+ num_latent_frames_, reference_height, reference_width = next(visual_geometry)
+ num_video_rows = num_latent_frames_ * (reference_height // patch_h) * (reference_width // patch_w)
+ audio_rows = slice(cursor, cursor + num_audio_rows)
+ video_rows = slice(audio_rows.stop, audio_rows.stop + num_video_rows)
+ cursor = video_rows.stop
+ audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop))
+ video_indices.append(torch.arange(video_rows.start, video_rows.stop))
+
+ frame_grid, width_grid = _frame_position_grid(reference_height, reference_width, patch_h, patch_w)
+ _fill_audio_positions(position_ids, audio_rows, reference_audio_latents, rotary_time, width_grid)
+ frame_time = _temporal_position_grid(num_latent_frames_, rotary_time)
+ position_ids[video_rows, 0] = frame_time.repeat_interleave(frame_grid.shape[0])
+ position_ids[video_rows, 1:] = frame_grid.repeat(num_latent_frames_, 1)
+ rotary_time += max(
+ float(reference_audio_latents), _temporal_position_span_sequential(num_latent_frames_)
+ )
+ else:
+ raise ValueError(f"A reference must be an 'image', a 'video' or an 'audio', got {reference.kind!r}.")
+
+ # The generated rows. Target audio and target video share the origin the reference blocks left behind.
+ audio_start = cursor
+ video_start = audio_start + num_target_audio_rows
+ _fill_audio_positions(
+ position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time, target_width_grid
+ )
+ frame_time = _temporal_position_grid(num_latent_frames, rotary_time)
+ position_ids[video_start:, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0])
+ position_ids[video_start:, 1:] = target_frame_grid.repeat(num_latent_frames, 1)
+
+ video_indices = torch.cat(video_indices + [torch.arange(video_start, sequence_length)])
+ audio_indices = torch.cat(audio_indices + [torch.arange(audio_start, video_start)])
+ text_indices = torch.arange(num_text_tokens)
+
+ token_tags = torch.empty(sequence_length, dtype=torch.long)
+ token_tags[text_indices] = text_token_tags.to(torch.long)
+ token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG
+ token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG
+
+ return (
+ position_ids,
+ token_tags,
+ video_indices,
+ audio_indices,
+ text_indices,
+ num_reference_video_rows,
+ num_reference_audio_rows,
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ # The canvas and the frame count are settled by the setup step, but in a hand-assembled chain this block
+ # sees user input directly — and a canvas off the multiple floor-divides into a broken decode.
+ multiple = components.canvas_multiple
+ if block_state.height % multiple or block_state.width % multiple:
+ raise ValueError(
+ f"`height` and `width` must be multiples of {multiple}, got {block_state.height}x{block_state.width}."
+ )
+ # A `ref2va` soundtrack is truncated to the generated duration as the references are prepared, so
+ # `num_frames` has to be final before the setup step runs.
+ ratio = components.vae_spatial_compression_ratio
+ block_state.num_latent_frames = video_latent_num_frames(
+ block_state.num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk
+ )
+ block_state.latent_height = block_state.height // ratio
+ block_state.latent_width = block_state.width // ratio
+ block_state.num_audio_latents = audio_latent_num_frames(block_state.num_frames)
+
+ (
+ position_ids,
+ token_tags,
+ video_indices,
+ audio_indices,
+ text_indices,
+ block_state.num_condition_video_rows,
+ block_state.num_condition_audio_rows,
+ ) = self.build_ref2va_packed_sequence(
+ block_state.text_token_tags,
+ block_state.normalized_references,
+ block_state.condition_latents,
+ block_state.audio_condition_latents,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ block_state.num_audio_latents,
+ components.patch_size,
+ )
+ block_state.position_ids = position_ids.to(device)
+ block_state.token_tags = token_tags.to(device)
+ block_state.video_indices = video_indices.to(device)
+ block_state.audio_indices = audio_indices.to(device)
+ block_state.text_indices = text_indices.to(device)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3PrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Draws every noise stream of the request and packs the video rows. MiniMax-H3 draws one condition at a "
+ "time first — noising the encoded anchors to its conditioning level — then the video noise as a latent "
+ "tensor, then the audio noise directly in row layout, all off the request's generator, in that order."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", MiniMaxH3Scheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames."
+ ),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam(
+ name="num_audio_latents",
+ type_hint=int,
+ required=True,
+ description="Number of audio latents per channel.",
+ ),
+ InputParam.template(
+ "generator",
+ description=(
+ "The generator of the request. The video noise is drawn from it first, then the audio noise."
+ ),
+ ),
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description=(
+ "Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, "
+ "used instead of the draw."
+ ),
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ description="Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many conditioning rows the layout reserved, which the packed conditioning must match.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many reference audio rows the layout reserved, which the packed rows must match.",
+ ),
+ InputParam(
+ name="condition_latents",
+ type_hint=list[torch.Tensor],
+ description=(
+ "The encoded video conditioning latents, one `(1, latent_channels, num_latent_frames, "
+ "latent_height, latent_width)` tensor per condition in packed order, or None for a request that "
+ "has none. Noised and packed here."
+ ),
+ ),
+ InputParam(
+ name="audio_condition_latents",
+ type_hint=list[torch.Tensor],
+ description=(
+ "The audio conditioning rows to prepend, one tensor per audio-bearing reference in packed "
+ "order. Empty for a request that has none, which is every `t2va` and `fl2va` one."
+ ),
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ OutputParam(
+ "audio_latents",
+ type_hint=torch.Tensor,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+ patch_size = components.patch_size
+
+ # A request draws every stream from the one generator it is given, and the order is part of what that
+ # generator reproduces: one draw per condition first, then the video noise as a latent tensor, then the audio
+ # noise directly in row layout. Passing `latents` or `audio_latents` skips its draw and shifts the ones after
+ # it.
+ condition_rows = None
+ if block_state.condition_latents:
+ # One draw per condition, in packed order. Each is packed on its own because `ref2va` references are
+ # encoded at their own resolutions, so their latents do not share a shape.
+ packed = []
+ for condition in block_state.condition_latents:
+ noise = randn_tensor(
+ condition.shape, generator=block_state.generator, device=device, dtype=torch.float32
+ )
+ # The anchors are not fully clean: the released model noises them to `t = 0.999` and holds them there
+ # for every step. Mixing before the patchify is the same arithmetic, since patchify only permutes.
+ noised = components.scheduler.scale_noise(condition.to(device), components.keyframe_noise_aug, noise)
+ packed.append(patchify_video_latents(noised, patch_size))
+ condition_rows = torch.cat(packed)
+ # In a hand-assembled chain the canvas reaching the layout is user input, so it can disagree with the
+ # keyframes that were actually encoded. Left alone the mismatch first surfaces as an `index_copy` shape
+ # error inside the transformer, 50 layers deep.
+ if condition_rows.shape[0] != block_state.num_condition_video_rows:
+ raise ValueError(
+ f"The layout reserved {block_state.num_condition_video_rows} conditioning rows but the encoded "
+ f"conditioning latents pack into {condition_rows.shape[0]}. The canvas the layout was built from "
+ "and the one the conditioning was encoded at do not agree."
+ )
+
+ latents = block_state.latents
+ if latents is None:
+ latents = randn_tensor(
+ (
+ 1,
+ components.vae_latent_channels,
+ block_state.num_latent_frames,
+ block_state.latent_height,
+ block_state.latent_width,
+ ),
+ generator=block_state.generator,
+ device=device,
+ dtype=torch.float32,
+ )
+ video_rows = patchify_video_latents(latents.to(device, torch.float32), patch_size)
+
+ if block_state.audio_latents is None:
+ audio_rows = randn_tensor(
+ (block_state.num_audio_latents * components.audio_channels, components.audio_latent_channels),
+ generator=block_state.generator,
+ device=device,
+ dtype=torch.float32,
+ )
+ else:
+ audio_rows = (
+ block_state.audio_latents.to(device, torch.float32)
+ .permute(0, 2, 1)
+ .reshape(-1, components.audio_latent_channels)
+ )
+
+ if condition_rows is not None:
+ video_rows = torch.cat([condition_rows, video_rows])
+ if block_state.audio_condition_latents:
+ num_reference_audio_rows = sum(rows.shape[0] for rows in block_state.audio_condition_latents)
+ if num_reference_audio_rows != block_state.num_condition_audio_rows:
+ raise ValueError(
+ f"The layout reserved {block_state.num_condition_audio_rows} reference audio rows but the encoded "
+ f"soundtracks pack into {num_reference_audio_rows}. The references the layout was built from and "
+ "the ones the audio conditioning was encoded from do not agree."
+ )
+ audio_rows = torch.cat([rows.to(device) for rows in block_state.audio_condition_latents] + [audio_rows])
+ block_state.latents, block_state.audio_latents = video_rows, audio_rows
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Initializes the two schedules — `shift = 12.0` for video, `shift = 3.0` for audio — and stages the "
+ "row-to-timestep plan of every step. One forward serves every modality and every noise level at once: "
+ "the generated rows step down their own schedule while the conditioning rows stay pinned at their "
+ "noise-augmentation level, and that assignment is static per step."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps", required=True),
+ InputParam(
+ name="video_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the video rows, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the audio rows, reference rows first.",
+ ),
+ InputParam(
+ name="text_indices",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Sequence positions of the text rows.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading video rows are conditioning rows.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading audio rows are reference rows.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps of the video schedule."),
+ OutputParam("audio_timesteps", type_hint=torch.Tensor, description="Timesteps of the audio schedule."),
+ OutputParam(
+ "row_timestep_plan",
+ type_hint=list,
+ description=(
+ "One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the "
+ "index of every row into them."
+ ),
+ ),
+ ]
+
+ @staticmethod
+ def build_row_timesteps(
+ video_indices: torch.Tensor,
+ audio_indices: torch.Tensor,
+ num_condition_video_rows: int,
+ num_condition_audio_rows: int,
+ num_text_tokens: int,
+ video_timestep: float,
+ audio_timestep: float,
+ condition_video_timestep: float,
+ condition_audio_timestep: float,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ r"""
+ Assign a timestep to every row of the packed sequence and reduce it to the transformer's `(timestep,
+ timestep_indices)` pair.
+
+ One forward serves rows at different noise levels: the generated video and audio rows step down their own
+ schedules while the conditioning rows stay pinned at their noise-augmentation level. Text rows never reach an
+ output head and inherit the video timestep.
+
+ Args:
+ video_indices (`torch.Tensor`): Sequence positions of the video rows, conditioning rows first.
+ audio_indices (`torch.Tensor`): Sequence positions of the audio rows, reference rows first.
+ num_condition_video_rows (`int`): How many leading video rows are conditioning rows.
+ num_condition_audio_rows (`int`): How many leading audio rows are reference rows.
+ num_text_tokens (`int`): Number of text rows, which never reach an output head.
+ video_timestep (`float`): Timestep of the generated video rows.
+ audio_timestep (`float`): Timestep of the generated audio rows.
+ condition_video_timestep (`float`): Timestep of the video conditioning rows.
+ condition_audio_timestep (`float`): Timestep of the audio reference rows.
+
+ Returns:
+ `tuple[torch.Tensor, torch.Tensor]`: the distinct timesteps, sorted, and the index of every row into them.
+ """
+ sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens)
+ row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)
+ row_timesteps[video_indices[:num_condition_video_rows]] = condition_video_timestep
+ row_timesteps[audio_indices[num_condition_audio_rows:]] = audio_timestep
+ row_timesteps[audio_indices[:num_condition_audio_rows]] = condition_audio_timestep
+ return torch.unique(row_timesteps, sorted=True, return_inverse=True)
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+ block_state.timesteps = components.scheduler.timesteps
+ block_state.audio_timesteps = components.audio_scheduler.timesteps
+
+ block_state.row_timestep_plan = [
+ tuple(
+ tensor.to(device)
+ for tensor in self.build_row_timesteps(
+ block_state.video_indices,
+ block_state.audio_indices,
+ block_state.num_condition_video_rows,
+ block_state.num_condition_audio_rows,
+ block_state.text_indices.numel(),
+ float(timestep),
+ float(audio_timestep),
+ max(float(timestep), components.keyframe_noise_aug),
+ 1.0,
+ )
+ )
+ for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+ ]
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..905c77a5608231b7e3bfbfa99d1f9cd807d2a148
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py
@@ -0,0 +1,458 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import numpy as np
+import PIL
+import torch
+from PIL import Image
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import (
+ MINIMAX_H3_FPS,
+ MiniMaxH3ModularPipeline,
+ align_num_frames,
+ resolve_canvas_size,
+)
+from .references import (
+ MiniMaxH3AudioReference,
+ MiniMaxH3ImageReference,
+ MiniMaxH3Reference,
+ MiniMaxH3VideoReference,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MiniMaxH3ResizeStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Puts the `fl2va` keyframes onto the target canvas — MiniMax-H3's own 768-short-edge geometry for the "
+ "aspect ratio of the first keyframe unless `height` and `width` say otherwise. The canvas resolved here "
+ "is the one the whole request generates at."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="image",
+ type_hint=PIL.Image.Image,
+ description=(
+ "Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is "
+ "derived from its own aspect ratio."
+ ),
+ ),
+ InputParam(
+ name="last_image",
+ type_hint=PIL.Image.Image,
+ description=(
+ "Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with "
+ "`image` it is the follower of the two and is cover-cropped onto the canvas."
+ ),
+ ),
+ InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+ InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam(
+ "keyframes",
+ type_hint=list,
+ description="The keyframes put onto the target canvas, in packed order.",
+ ),
+ OutputParam(
+ "keyframe_anchors",
+ type_hint=tuple,
+ description=(
+ "Which end of the video every keyframe is anchored to, in packed order. Positional with "
+ "`keyframes`, so both are resolved here."
+ ),
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ keyframes = [keyframe for keyframe in (block_state.image, block_state.last_image) if keyframe is not None]
+ block_state.keyframe_anchors = tuple(
+ anchor
+ for anchor, keyframe in (("first", block_state.image), ("last", block_state.last_image))
+ if keyframe is not None
+ )
+ if block_state.height is None:
+ block_state.height, block_state.width = resolve_canvas_size(*keyframes[0].size, components.canvas_multiple)
+
+ prepared = []
+ for index, keyframe in enumerate(keyframes):
+ if keyframe.size == (block_state.width, block_state.height):
+ prepared.append(keyframe)
+ elif index == 0:
+ # The geometry anchor is stretched onto the canvas. `resize_mode="default"` is exactly PIL's
+ # `resize((width, height), LANCZOS)`, verified pixel-identical across aspect ratios.
+ prepared.append(
+ components.image_processor.resize(keyframe, height=block_state.height, width=block_state.width)
+ )
+ else:
+ # The follower is cover-cropped. `VaeImageProcessor`'s `resize_mode="crop"` is *not* a drop-in here:
+ # it sizes with floor division and centres with `w // 2 - src_w // 2`, where MiniMax-H3 rounds and
+ # centres with `(src_w - w) // 2`. The two agree on some aspect ratios and differ by a pixel on
+ # others (106 of 218 sampled), which would move the conditioning latents off the reference
+ # implementation, so the released model's arithmetic is kept.
+ scale = max(block_state.width / keyframe.size[0], block_state.height / keyframe.size[1])
+ resized_size = (
+ max(block_state.width, round(keyframe.size[0] * scale)),
+ max(block_state.height, round(keyframe.size[1] * scale)),
+ )
+ left = max(0, (resized_size[0] - block_state.width) // 2)
+ top = max(0, (resized_size[1] - block_state.height) // 2)
+ resized = keyframe.resize(resized_size, Image.Resampling.LANCZOS)
+ prepared.append(resized.crop((left, top, left + block_state.width, top + block_state.height)))
+ block_state.keyframes = prepared
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VASetupStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ def __init__(
+ self,
+ max_images: int = 9,
+ max_videos: int = 3,
+ max_audios: int = 3,
+ max_references: int = 12,
+ reference_image_short_edge: int = 2048,
+ ):
+ r"""
+ Resolve a `ref2va` request's plan and normalize its references onto MiniMax-H3's own rates and resolutions.
+
+ Args:
+ max_images (`int`, defaults to 9): Image references a request may carry.
+ max_videos (`int`, defaults to 3): Video references a request may carry.
+ max_audios (`int`, defaults to 3): Audio references a request may carry.
+ max_references (`int`, defaults to 12): References of any modality a request may carry in total.
+ reference_image_short_edge (`int`, defaults to 2048):
+ The short edge an image reference is resized to, upscaling included.
+
+ The limits are what MiniMax-H3 documents for the released checkpoint; they bound nothing but this block's own
+ validation, so a fine-tune that packs more can raise them.
+ """
+ self.max_images = max_images
+ self.max_videos = max_videos
+ self.max_audios = max_audios
+ self.max_references = max_references
+ self.reference_image_short_edge = reference_image_short_edge
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resolves the `ref2va` plan — the canvas (MiniMax-H3's own 16:9 unless asked otherwise; references never "
+ "bind the generated geometry) and the `17 * n + 5` frame count — and normalizes every reference onto "
+ "MiniMax-H3's own rates and resolutions."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="references",
+ type_hint=list[MiniMaxH3Reference],
+ required=True,
+ description=(
+ "The references to condition on, **in the order the model should read them**: the order labels "
+ "them in the prompt presentation and lays them out on the shared rotary clock, so a different "
+ "order is a different request. One dataclass per modality, all holding in-memory media — a "
+ "[`MiniMaxH3ImageReference`] (at most 9), a [`MiniMaxH3VideoReference`] at its own `fps` (at most "
+ "3, whose `audio` soundtrack is conditioned on as well), or a [`MiniMaxH3AudioReference`] at its "
+ "own `sample_rate` (at most 3) — for at most 12 references in total, and audio references cannot "
+ "be the only ones. These blocks never open a media file: decode with each class's `from_file` "
+ "classmethod, which brings the rates along."
+ ),
+ ),
+ InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+ InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+ InputParam(
+ name="num_frames",
+ type_hint=int,
+ required=True,
+ description=(
+ "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video "
+ "VAE can decode; the resulting duration must stay between 5 and 15 seconds. To generate a video "
+ "as long as a reference soundtrack, pass `round(samples / sample_rate * 24)`."
+ ),
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."),
+ OutputParam(
+ "normalized_references",
+ type_hint=list[MiniMaxH3Reference],
+ description=(
+ "The references normalized onto MiniMax-H3's own rates and resolutions, in packed order: the "
+ "same public reference types the request passed in, with an image resized to its own "
+ "2048 pixel short edge, a video resampled onto 24 fps and onto the canvas its own aspect ratio "
+ "resolves to, and a soundtrack put on the audio VAE's sample rate and truncated to the generated "
+ "duration."
+ ),
+ ),
+ ]
+
+ @staticmethod
+ def _normalize_video_condition(frames, fps: float, num_frames: int, canvas_multiple: int) -> np.ndarray:
+ r"""
+ Normalize a video reference's frames: any accepted layout, onto `uint8` at 24 fps, truncated to the generated
+ frame count, on the canvas its own aspect ratio resolves to.
+
+ The two passes reproduce the reference implementation's `ffmpeg` decode, in the same order: the constant frame
+ rate resample first (dropping and duplicating whole frames, as `ffmpeg`'s `fps` filter does), the LANCZOS
+ rescale second. Frames handed over at 24 fps and already at the canvas their own aspect ratio resolves to flow
+ through untouched, which is the parity-exact route: the reference rescaled with `ffmpeg`'s own LANCZOS while
+ decoding, so only frames decoded at the canvas reproduce its pixels bit for bit.
+
+ Args:
+ frames (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`):
+ The reference frames: a list of images, a `(num_frames, height, width, 3)` array or a
+ `(num_frames, 3, height, width)` tensor, `uint8` or floating point over `[0, 1]`.
+ fps (`float`): The frame rate `frames` carries.
+ num_frames (`int`): The generated frame count the reference is truncated to.
+ canvas_multiple (`int`): What both canvas axes round to, i.e. `components.canvas_multiple`.
+
+ Returns:
+ `np.ndarray` of shape `(num_frames, height, width, 3)`: the normalized `uint8` RGB frames.
+ """
+ # Any accepted layout onto `uint8` THWC. A `torch.Tensor` is channels-first, as everywhere else in
+ # diffusers, and a `np.ndarray` channels-last; floating point values are read over `[0, 1]`.
+ if isinstance(frames, list):
+ frames = np.stack([np.asarray(frame.convert("RGB")) for frame in frames])
+ if isinstance(frames, torch.Tensor):
+ frames = frames.movedim(-3, -1).cpu().numpy()
+ frames = np.asarray(frames)
+ if frames.dtype != np.uint8:
+ frames = (frames * 255.0).round().clip(0, 255).astype(np.uint8)
+ if frames.ndim != 4 or frames.shape[3] != 3:
+ raise ValueError(
+ f"A reference video must be `(num_frames, height, width, 3)` RGB frames, got {tuple(frames.shape)}."
+ )
+
+ # Onto MiniMax-H3's 24 fps grid: every frame is held until the slot of the next one, and the last one until
+ # the slot the stream's end rounds to.
+ if fps <= 0:
+ raise ValueError(f"A reference video must have a positive frame rate, got {fps}.")
+ if fps != MINIMAX_H3_FPS:
+ scale = MINIMAX_H3_FPS / fps
+ slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64)
+ frames = np.repeat(frames, np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)), axis=0)
+
+ # Truncated to the generated frame count and put on the canvas of its *own* aspect ratio — the same rule the
+ # target canvas follows, unlike an image reference.
+ frames = frames[:num_frames]
+ height, width = resolve_canvas_size(frames.shape[2], frames.shape[1], canvas_multiple)
+ if frames.shape[1:3] == (height, width):
+ return frames
+ return np.stack(
+ [np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS)) for frame in frames]
+ )
+
+ @staticmethod
+ def _normalize_audio_condition(
+ waveform: torch.Tensor, sample_rate: int, target_sample_rate: int, max_duration: float
+ ) -> torch.Tensor:
+ r"""
+ Normalize a reference soundtrack onto the audio VAE's sample rate, as a stereo waveform.
+
+ The reference implementation extracts a soundtrack at a native rate, truncates it there and resamples it once,
+ in torch, which this mirrors: the truncation is applied at `sample_rate` and the resampling is a single
+ `torchaudio` pass. A mono waveform is upmixed by repeating its channel.
+
+ Args:
+ waveform (`torch.Tensor` of shape `(channels, num_samples)`): The soundtrack, mono or stereo.
+ sample_rate (`int`): The sample rate `waveform` carries its samples at.
+ target_sample_rate (`int`): The audio VAE's sample rate, i.e. what the waveform is resampled to.
+ max_duration (`float`): Truncate the reference to this many seconds.
+
+ Returns:
+ `torch.Tensor` of shape `(2, num_samples)`: the float32 waveform.
+ """
+ waveform = torch.as_tensor(waveform)
+ if waveform.ndim != 2 or waveform.shape[0] not in (1, 2):
+ raise ValueError(
+ "A reference soundtrack must be a `(channels, num_samples)` mono or stereo waveform, got "
+ f"{tuple(waveform.shape)}."
+ )
+ waveform = waveform.to(torch.float32)[:, : int(max_duration * sample_rate)]
+ if waveform.shape[0] != 2:
+ waveform = waveform.expand(2, -1).contiguous()
+ if sample_rate == target_sample_rate:
+ return waveform
+
+ try:
+ import torchaudio
+ except ImportError as error:
+ raise ImportError(
+ f"Resampling a MiniMax-H3 reference soundtrack from {sample_rate} Hz to {target_sample_rate} Hz "
+ "needs `torchaudio`. Pass a waveform already at the audio VAE's sample rate to do without it."
+ ) from error
+ return torchaudio.transforms.Resample(sample_rate, target_sample_rate)(waveform)
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # 1. Validate the request.
+ if (block_state.height is None) != (block_state.width is None):
+ raise ValueError("`height` and `width` have to be passed together, or neither of them.")
+ multiple = components.canvas_multiple
+ if block_state.height is not None and (block_state.height % multiple or block_state.width % multiple):
+ raise ValueError(
+ f"`height` and `width` must be multiples of {multiple}, got {block_state.height}x{block_state.width}."
+ )
+ if not block_state.references:
+ raise ValueError("`ref2va` needs at least one reference; use the `t2va` workflow for text-only requests.")
+ for index, entry in enumerate(block_state.references):
+ if not isinstance(entry, MiniMaxH3Reference):
+ raise ValueError(
+ f"`references[{index}]` must be a [`MiniMaxH3ImageReference`], [`MiniMaxH3VideoReference`] or "
+ f"[`MiniMaxH3AudioReference`], got {type(entry)}. MiniMax-H3 blocks never open media files, so a "
+ "request that holds paths decodes them first, with each class's `from_file` classmethod."
+ )
+ kinds = [entry.kind for entry in block_state.references]
+ for kind, limit in (("image", self.max_images), ("video", self.max_videos), ("audio", self.max_audios)):
+ if kinds.count(kind) > limit:
+ raise ValueError(f"MiniMax-H3 accepts at most {limit} {kind} references, got {kinds.count(kind)}.")
+ if len(kinds) > self.max_references:
+ raise ValueError(
+ f"MiniMax-H3 accepts at most {self.max_references} references in total, got {len(kinds)}."
+ )
+ if set(kinds) == {"audio"}:
+ raise ValueError(
+ "An audio reference has to be paired with at least one image or video reference and cannot be used "
+ "on its own."
+ )
+
+ # 2. Resolve the canvas and the frame count. The duration the request generates is the one of the *aligned*
+ # frame count, so that is what the ceiling holds for: 346 frames would otherwise pass the check and then be
+ # rounded up to 362, i.e. 15.083 seconds.
+ if block_state.height is None:
+ block_state.height, block_state.width = resolve_canvas_size(16, 9, multiple)
+ aligned_num_frames = align_num_frames(
+ block_state.num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk
+ )
+ duration = aligned_num_frames / components.fps
+ if not components.min_duration <= duration <= components.max_duration:
+ raise ValueError(
+ f"MiniMax-H3 generates between {components.min_duration} and {components.max_duration} seconds at "
+ f"{components.fps} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can "
+ f"encode, must be between {int(components.min_duration * components.fps)} and "
+ f"{int(components.max_duration * components.fps)}, got {block_state.num_frames} (rounded up to "
+ f"{aligned_num_frames})."
+ )
+ if aligned_num_frames != block_state.num_frames:
+ logger.warning(
+ f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} "
+ f"up to {aligned_num_frames}."
+ )
+ block_state.num_frames = aligned_num_frames
+
+ # 3. Normalize the media, in packed order.
+ normalized = []
+ for entry in block_state.references:
+ waveform = None
+ if entry.has_audio:
+ sample_rate = entry.sample_rate
+ if sample_rate is None:
+ sample_rate = components.audio_sampling_rate
+ waveform = self._normalize_audio_condition(
+ entry.audio,
+ sample_rate,
+ components.audio_sampling_rate,
+ max_duration=block_state.num_frames / components.fps,
+ )
+
+ if entry.kind == "image":
+ # Images are encoded at high detail — a 2048 pixel short edge, upscaling included, with *no* area
+ # cap — unlike video references and the target itself, which share the one canvas rule. Expected in
+ # RGB, as `load_image` returns.
+ image = entry.image
+ if image.size[0] <= 0 or image.size[1] <= 0:
+ raise ValueError(f"A reference image must have a positive size, got {image.size}.")
+ width, height = image.size
+ if width > 4 * height or height > 4 * width:
+ raise ValueError(f"A reference image must be within 1:4 and 4:1, got {width}x{height}.")
+ scale = self.reference_image_short_edge / min(width, height)
+ target_height = max(multiple, round(height * scale / multiple) * multiple)
+ target_width = max(multiple, round(width * scale / multiple) * multiple)
+ if image.size != (target_width, target_height):
+ image = components.image_processor.resize(image, height=target_height, width=target_width)
+ normalized.append(MiniMaxH3ImageReference(image=image))
+ elif entry.kind == "video":
+ normalized.append(
+ MiniMaxH3VideoReference(
+ frames=self._normalize_video_condition(
+ entry.frames, float(entry.fps), block_state.num_frames, multiple
+ ),
+ fps=float(components.fps),
+ audio=waveform,
+ sample_rate=None if waveform is None else components.audio_sampling_rate,
+ )
+ )
+ else:
+ normalized.append(MiniMaxH3AudioReference(audio=waveform, sample_rate=components.audio_sampling_rate))
+ block_state.normalized_references = normalized
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2c596cfd98dc926cdaf0812e9b3c04c40d4c8bc
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/decoders.py
@@ -0,0 +1,246 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import MiniMaxH3ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MiniMaxH3AfterDenoiseStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Turns the denoised rows of the packed sequence back into latents: drops the conditioning rows the loop "
+ "never wrote, then unpacks the video rows into `(1, latent_channels, num_latent_frames, latent_height, "
+ "latent_width)` and the channel-major audio rows into the `(2, audio_latent_channels, num_audio_latents)` "
+ "the mono audio VAE consumes. The decoders then take latents from any source, and popping them leaves "
+ "latents in hand."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The denoised video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The denoised audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading video rows are conditioning rows and are dropped here.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading audio rows are reference rows and are dropped here.",
+ ),
+ InputParam(
+ name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames."
+ ),
+ InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."),
+ InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."),
+ InputParam(
+ name="num_audio_latents",
+ type_hint=int,
+ required=True,
+ description="Number of audio latents per channel.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The generated video latents, of shape `(1, latent_channels, num_latent_frames, "
+ "latent_height, latent_width)`.",
+ ),
+ OutputParam(
+ "audio_latents",
+ type_hint=torch.Tensor,
+ description="The generated audio latents, one batch item per stereo channel.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ patch_t, patch_h, patch_w = components.patch_size
+ channels = components.vae_latent_channels
+
+ # The inverse of the patchify in the prepare-latents step: rows are frame-major then row-major.
+ rows = block_state.latents[block_state.num_condition_video_rows :]
+ rows = rows.reshape(
+ -1,
+ block_state.num_latent_frames // patch_t,
+ block_state.latent_height // patch_h,
+ block_state.latent_width // patch_w,
+ channels,
+ patch_t,
+ patch_h,
+ patch_w,
+ )
+ rows = rows.permute(0, 4, 1, 5, 2, 6, 3, 7)
+ block_state.latents = rows.reshape(
+ -1, channels, block_state.num_latent_frames, block_state.latent_height, block_state.latent_width
+ ).contiguous()
+
+ # Audio rows are channel-major, and the mono audio VAE takes the two stereo channels as two batch items.
+ audio_rows = block_state.audio_latents[block_state.num_condition_audio_rows :]
+ audio_rows = audio_rows.reshape(components.audio_channels, block_state.num_audio_latents, audio_rows.shape[-1])
+ block_state.audio_latents = audio_rows.permute(0, 2, 1).contiguous()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3VideoDecodeStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denormalizes the generated video latents and decodes them into video. The spatial tiling of the video "
+ "VAE covers the canvas exactly, so the decoded frames need no crop back, but the decode itself runs under "
+ "float16 autocast even though the VAE weights are float32, and the VAE produces ImageNet-normalized RGB "
+ "that is reverted here."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLMiniMaxH3),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ # The video VAE decodes into ImageNet-normalized RGB over a [0, 1] base range, which this block
+ # reverts itself, so the processor must not denormalize a second time.
+ config=FrozenDict({"vae_scale_factor": 16, "do_normalize": False}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The generated video latents.",
+ ),
+ InputParam.template("output_type", description="Output format: 'pil', 'np' or 'pt'."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("videos", description="The generated video.")]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ latents_mean = torch.tensor(components.vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1)
+ latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1)
+ latents = block_state.latents * latents_std + latents_mean
+
+ with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
+ video = components.vae.decode(latents, return_dict=False)[0]
+ pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1)
+ pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1)
+ video = (video.float() * pixel_std + pixel_mean).clamp(0, 1)
+ block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3AudioDecodeStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denormalizes the generated audio latents and decodes them into a stereo waveform. The audio VAE is mono "
+ "and takes the two stereo channels as two batch items."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The generated audio latents, one batch item per stereo channel.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "audio",
+ type_hint=torch.Tensor,
+ description="The generated soundtrack, of shape `(1, 2, num_samples)`.",
+ ),
+ OutputParam(
+ "sampling_rate",
+ type_hint=int,
+ description="Sample rate of the generated soundtrack in Hz.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean, device=device).view(1, -1, 1)
+ audio_latents_std = torch.tensor(components.audio_vae.config.latents_std, device=device).view(1, -1, 1)
+ audio_latents = block_state.audio_latents * audio_latents_std + audio_latents_mean
+
+ audio = components.audio_vae.decode(audio_latents, return_dict=False)[0]
+ block_state.audio = audio.float().permute(1, 0, 2)
+ block_state.sampling_rate = components.audio_sampling_rate
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..26df64b9f00b18e7daee9cf4f001fa7ed83e58c0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/denoise.py
@@ -0,0 +1,287 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from ...models import MiniMaxH3Transformer3DModel
+from ...schedulers import MiniMaxH3Scheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import MiniMaxH3ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MiniMaxH3LoopDenoiser(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ def __init__(self, transformer_name: str = "transformer"):
+ r"""
+ Run the one MiniMax-H3 forward pass of a denoising iteration.
+
+ Args:
+ transformer_name (`str`, defaults to `"transformer"`):
+ The component the forward runs against. One repository holds both checkpoint partitions —
+ `transformer/` for `t2va`/`fl2va` and `transformer_ref/` for `ref2va` — under different component
+ names, so which partition a loop drives is configuration, not a separate block.
+ """
+ self.transformer_name = transformer_name
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Runs the one MiniMax-H3 forward pass of a denoising iteration, which predicts the velocity of every row "
+ "of the packed sequence at once. The checkpoint is guidance-distilled, so there is no unconditional pass "
+ "and no guider."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec(self.transformer_name, MiniMaxH3Transformer3DModel)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam.template("prompt_embeds"),
+ InputParam(
+ name="row_timestep_plan",
+ type_hint=list,
+ required=True,
+ description="One `(timestep, timestep_indices)` pair per step.",
+ ),
+ InputParam(
+ kwargs_type="denoiser_input_fields",
+ description=(
+ "The structural description of the packed sequence the transformer reads by name: `token_tags`, "
+ "`position_ids` and the three row-index tensors."
+ ),
+ ),
+ InputParam.template("attention_kwargs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "noise_pred",
+ type_hint=torch.Tensor,
+ description="Predicted velocity of the video rows of the sequence.",
+ ),
+ OutputParam(
+ "audio_noise_pred",
+ type_hint=torch.Tensor,
+ description="Predicted velocity of the audio rows of the sequence.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ transformer = getattr(components, self.transformer_name)
+ unique_timesteps, timestep_indices = block_state.row_timestep_plan[i]
+ # The layout tags its outputs `denoiser_input_fields`, and their names are the transformer's own parameter
+ # names, so the rows of the packed sequence are described to it without this block enumerating them.
+ layout_kwargs = {
+ name: value
+ for name, value in block_state.denoiser_input_fields.items()
+ if name in inspect.signature(transformer.forward).parameters
+ }
+ block_state.noise_pred, block_state.audio_noise_pred = transformer(
+ hidden_states=block_state.latents[None],
+ audio_hidden_states=block_state.audio_latents[None],
+ encoder_hidden_states=block_state.prompt_embeds,
+ timestep=unique_timesteps,
+ timestep_indices=timestep_indices,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **layout_kwargs,
+ )
+ return components, block_state
+
+
+class MiniMaxH3Ref2VALoopDenoiser(MiniMaxH3LoopDenoiser):
+ model_name = "minimax-h3"
+
+ def __init__(self):
+ super().__init__(transformer_name="transformer_ref")
+
+
+class MiniMaxH3LoopSchedulerStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Steps the generated video and audio rows down their own schedule. The conditioning rows are re-imposed "
+ "by construction: only the generated rows are ever written, so the anchors survive the whole loop."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The video rows of the packed sequence, conditioning rows first.",
+ ),
+ InputParam(
+ name="audio_latents",
+ type_hint=torch.Tensor,
+ required=True,
+ description="The channel-major audio rows of the packed sequence, reference rows first.",
+ ),
+ InputParam(
+ name="noise_pred",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Predicted velocity of the video rows.",
+ ),
+ InputParam(
+ name="audio_noise_pred",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Predicted velocity of the audio rows.",
+ ),
+ InputParam(
+ name="audio_timesteps",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Timesteps of the audio schedule.",
+ ),
+ InputParam(
+ name="num_condition_video_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading video rows are conditioning rows.",
+ ),
+ InputParam(
+ name="num_condition_audio_rows",
+ type_hint=int,
+ default=0,
+ description="How many leading audio rows are reference rows.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The video rows of the packed sequence after one step.",
+ ),
+ OutputParam(
+ "audio_latents",
+ type_hint=torch.Tensor,
+ description="The audio rows of the packed sequence after one step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ num_condition_video_rows = block_state.num_condition_video_rows
+ num_condition_audio_rows = block_state.num_condition_audio_rows
+
+ block_state.latents[num_condition_video_rows:] = components.scheduler.step(
+ block_state.noise_pred[0, num_condition_video_rows:].float(),
+ t,
+ block_state.latents[num_condition_video_rows:],
+ return_dict=False,
+ )[0]
+ block_state.audio_latents[num_condition_audio_rows:] = components.audio_scheduler.step(
+ block_state.audio_noise_pred[0, num_condition_audio_rows:].float(),
+ block_state.audio_timesteps[i],
+ block_state.audio_latents[num_condition_audio_rows:],
+ return_dict=False,
+ )[0]
+ return components, block_state
+
+
+class MiniMaxH3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return "Iteratively denoises the packed MiniMax-H3 sequence over the two schedules."
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", MiniMaxH3Scheduler),
+ ComponentSpec("audio_scheduler", MiniMaxH3Scheduler),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("timesteps", required=True, description="Timesteps of the video schedule."),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ with self.progress_bar(total=len(block_state.timesteps)) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ progress_bar.update()
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3DenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+ block_classes = [MiniMaxH3LoopDenoiser, MiniMaxH3LoopSchedulerStep]
+ block_names = ["denoiser", "update"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the `t2va` / `fl2va` MiniMax-H3 denoising loop, one forward pass per step."
+
+
+class MiniMaxH3Ref2VADenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VALoopDenoiser, MiniMaxH3LoopSchedulerStep]
+ block_names = ["denoiser", "update"]
+
+ @property
+ def description(self) -> str:
+ return "Runs the `ref2va` MiniMax-H3 denoising loop, one forward pass per step."
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..ecc556ef1d3d803afa111cf8f4e8cfb68ee02b69
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/encoders.py
@@ -0,0 +1,682 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+from ...models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import MiniMaxH3ModularPipeline
+from .references import MiniMaxH3Reference
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def get_qwen3vl_prompt_embeds(
+ text_encoder,
+ processor,
+ token_ids: list[int],
+ vision_inputs: dict | None = None,
+ text_encoder_layer: int = 50,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+ r"""
+ Encode one tokenized MiniMax-H3 presentation with the Qwen3-VL conditioner.
+
+ Unlike its `get__prompt_embeds` siblings this takes the *tokenized* presentation rather than a prompt
+ string: MiniMax-H3's tokenization is the presentation-building phase that precedes it, where every reference or
+ keyframe prepends a label and a vision block. The `mm_token_type_ids` built here are Qwen-internal (`0` text, `1`
+ image, `2` video — they drive Qwen3-VL's per-modality-run rotary layout and never leave this function) and are not
+ MiniMax-H3's own per-row modality tags, which the presentation phase produces alongside `token_ids`.
+
+ Args:
+ text_encoder (`Qwen3VLForConditionalGeneration`): The conditioner.
+ processor (`Qwen3VLProcessor`): Its processor, which derives the token type ids from the vision pad ids.
+ token_ids (`list[int]`): The tokenized presentation.
+ vision_inputs (`dict`, *optional*):
+ The vision tensors of the presentation's blocks, by the conditioner's own parameter names —
+ `pixel_values` / `image_grid_thw` for images, `pixel_values_videos` / `video_grid_thw` for videos.
+ text_encoder_layer (`int`, *optional*, defaults to 50):
+ Which hidden state conditions the transformer, i.e. `components.text_encoder_layer`.
+ device (`torch.device`, *optional*): The device to run the conditioner on.
+ dtype (`torch.dtype`, *optional*): The dtype of the returned embeddings.
+
+ Returns:
+ `torch.Tensor`: the `(1, num_text_tokens, 5120)` hidden state after decoder layer `text_encoder_layer`.
+ """
+ num_layers = text_encoder.config.text_config.num_hidden_layers
+ if num_layers <= text_encoder_layer:
+ raise ValueError(
+ f"MiniMax-H3 conditions on `hidden_states[{text_encoder_layer}]` of its Qwen3-VL "
+ f"conditioner, which needs more than {text_encoder_layer} decoder layers, but "
+ f"`text_encoder` has {num_layers}. The last hidden state of a stack truncated to exactly "
+ f"{text_encoder_layer} layers is post-norm and is not the conditioning MiniMax-H3 expects."
+ )
+
+ input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
+ mm_token_type_ids = torch.tensor(processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device)
+
+ vision_kwargs = {}
+ for name, value in (vision_inputs or {}).items():
+ vision_kwargs[name] = value.to(device, text_encoder.dtype) if name.startswith("pixel_") else value.to(device)
+
+ # `text_encoder.model` is a submodule, and a CPU-offload hook — accelerate's or the one the
+ # `ComponentsManager` attaches — wraps the *top-level* module's `forward` alone, so calling the submodule
+ # directly would leave the conditioner on the CPU. Fire the hook by hand instead of routing through
+ # `text_encoder(...)`: MiniMax-H3 reads `hidden_states[50]` and never uses the language-model head, whose
+ # vocabulary-wide projection over every token is all the top-level forward would add.
+ # TODO: firing another module's offload hook by hand is not something a block should do. It is here
+ # because MiniMax-H3 reads `hidden_states[50]` off `text_encoder.model` while the hook wraps only the
+ # top-level `forward`. Needs a real answer — an opt-in on the conditioner, or a hook that follows
+ # submodule calls.
+ hook = getattr(text_encoder, "_hf_hook", None)
+ if hook is not None and hasattr(hook, "pre_forward"):
+ hook.pre_forward(text_encoder)
+ outputs = text_encoder.model(
+ input_ids=input_ids,
+ attention_mask=torch.ones_like(input_ids),
+ mm_token_type_ids=mm_token_type_ids,
+ use_cache=False,
+ output_hidden_states=True,
+ **vision_kwargs,
+ )
+ return outputs.hidden_states[text_encoder_layer].to(device=device, dtype=dtype)
+
+
+def encode_vae_condition(
+ vae, pixels: torch.Tensor, pixel_mean: tuple, pixel_std: tuple, encode_seed: int = 42
+) -> torch.Tensor:
+ r"""
+ Encode one visual condition — a keyframe or a `ref2va` reference — into normalized conditioning latents.
+
+ The recipe is the released model's, and every part of it is needed to reproduce its conditioning: the pixels are
+ ImageNet-normalized, the posterior is *sampled* under a fresh generator seeded independently of the request, and
+ the sampled latent is rounded to float16 (~11 bits of every conditioning latent) before being normalized. A
+ single frame is encoded by the (tiled) spatial encoder alone; a frame stack goes through the temporal chunking,
+ which is what turns `17 * n + 5` frames into `5 * n + 2` latent frames.
+
+ Args:
+ vae (`AutoencoderKLMiniMaxH3`): The video VAE.
+ pixels (`torch.Tensor` of shape `(1, 3, num_frames, height, width)`):
+ The condition's `uint8` pixels, on the VAE's device. `num_frames` is 1 for a keyframe or an image
+ reference.
+ pixel_mean (`tuple[float, float, float]`), pixel_std (`tuple[float, float, float]`):
+ The video VAE's pixel convention, i.e. `components.pixel_mean` / `components.pixel_std`.
+ encode_seed (`int`, *optional*, defaults to 42):
+ Seed the posterior is sampled under, i.e. `components.keyframe_encode_seed`.
+
+ Returns:
+ `torch.Tensor`: one `(1, latent_channels, num_latent_frames, latent_height, latent_width)` float32 CPU
+ tensor.
+ """
+ latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1)
+ latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1)
+ pixel_mean = torch.tensor(pixel_mean, device=pixels.device).view(1, -1, 1, 1, 1)
+ pixel_std = torch.tensor(pixel_std, device=pixels.device).view(1, -1, 1, 1, 1)
+
+ pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std
+ posterior = vae.encode(pixels, return_dict=False)[0]
+ latents = posterior.sample(generator=torch.Generator().manual_seed(encode_seed))
+ latents = latents.to(torch.float16).float().cpu()
+ return (latents - latents_mean) / latents_std
+
+
+class MiniMaxH3TextEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes MiniMax-H3's presentation of a `t2va` / `fl2va` request: the prompt verbatim, preceded by a "
+ '`": "` label and a vision block per keyframe, with no chat template and no special tokens. '
+ "The checkpoint is guidance-distilled, so there is no negative prompt and no unconditional branch."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ComponentSpec("processor", Qwen3VLProcessor),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", description="The prompt to guide generation, a single string."),
+ InputParam(
+ name="keyframes",
+ type_hint=list,
+ description="The keyframes put onto the target canvas, in packed order (empty or None for `t2va`).",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template(
+ "prompt_embeds",
+ description=(
+ "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the "
+ "50th decoder layer of the Qwen3-VL conditioner."
+ ),
+ ),
+ OutputParam(
+ "text_token_tags",
+ type_hint=torch.Tensor,
+ description=(
+ "The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video."
+ ),
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if not isinstance(block_state.prompt, str):
+ raise ValueError(
+ "MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got "
+ f"{type(block_state.prompt)}."
+ )
+ tokenizer, processor = components.tokenizer, components.processor
+ text_tag, video_tag = components.text_tag, components.video_tag
+
+ # 1. Vision features of the keyframes, batched. The presentation is text-only without them.
+ vision_inputs, image_grid_thw = {}, None
+ if block_state.keyframes:
+ vision = processor.image_processor(images=block_state.keyframes, return_tensors="pt")
+ image_grid_thw = vision["image_grid_thw"]
+ vision_inputs = {"pixel_values": vision["pixel_values"], "image_grid_thw": image_grid_thw}
+
+ # 2. The presentation, tokenized: a `": "` label and a vision block per keyframe, then the prompt
+ # verbatim. The rows of a vision block are tagged as *video* rather than text, which is what the
+ # transformer's AdaLN modulation keys off.
+ token_ids, token_tags = [], []
+ if block_state.keyframes:
+ merge_size = processor.image_processor.merge_size**2
+ for index in range(len(block_state.keyframes)):
+ num_image_tokens = int(image_grid_thw[index].prod()) // merge_size
+ label_ids = tokenizer(f": ", add_special_tokens=False)["input_ids"]
+ vision_ids = (
+ [tokenizer.convert_tokens_to_ids("<|vision_start|>")]
+ + [tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens
+ + [tokenizer.convert_tokens_to_ids("<|vision_end|>")]
+ )
+ token_ids += label_ids + vision_ids
+ token_tags += [text_tag] * len(label_ids) + [video_tag] * len(vision_ids)
+ prompt_ids = tokenizer(block_state.prompt, add_special_tokens=False)["input_ids"]
+ token_ids += prompt_ids
+ token_tags += [text_tag] * len(prompt_ids)
+
+ # 3. One conditioner call. The block emits the conditioner's dtype, as every other model does — it has no
+ # denoiser of its own, since it is meant to run on its own.
+ block_state.prompt_embeds = get_qwen3vl_prompt_embeds(
+ components.text_encoder,
+ processor,
+ token_ids,
+ vision_inputs,
+ text_encoder_layer=components.text_encoder_layer,
+ device=components._execution_device,
+ dtype=components.text_encoder.dtype,
+ )
+ block_state.text_token_tags = torch.tensor(token_tags, dtype=torch.long)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3KeyframeVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes the `fl2va` keyframes into conditioning latents. They become the anchors of the whole denoising "
+ "loop, which only ever writes the generated rows, so they are never updated again — the prepare-latents "
+ "step noises them to MiniMax-H3's conditioning level and packs them."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKLMiniMaxH3)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="keyframes",
+ type_hint=list,
+ required=True,
+ description="The keyframes put onto the target canvas, in packed order.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "condition_latents",
+ type_hint=list[torch.Tensor],
+ description=(
+ "The normalized video conditioning latents, one `(1, latent_channels, 1, latent_height, "
+ "latent_width)` tensor per keyframe, in packed order. One entry per condition is what the "
+ "prepare-latents step draws its noise against, so the list is the unit the request's generator "
+ "is consumed in."
+ ),
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ # A keyframe is a single frame, so the VAE runs its spatial encoder alone with none of the 17-frame
+ # temporal chunking.
+ block_state.condition_latents = [
+ encode_vae_condition(
+ components.vae,
+ torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None],
+ components.pixel_mean,
+ components.pixel_std,
+ components.keyframe_encode_seed,
+ )
+ for image in block_state.keyframes
+ ]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VATextEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ def __init__(self, video_sample_fps: float = 2.0):
+ r"""
+ Encode the presentation of a `ref2va` request.
+
+ Args:
+ video_sample_fps (`float`, defaults to 2.0):
+ The rate the conditioner reads a reference video at: every `24 / video_sample_fps`-th of the
+ normalized 24 fps frames.
+ """
+ self.video_sample_fps = video_sample_fps
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes MiniMax-H3's presentation of a `ref2va` request: a label per reference, numbered per modality "
+ '(`": "` plus a vision block, `": "` alone, `": "` plus one timestamped '
+ "vision block per merged frame pair), then the prompt verbatim. The checkpoint is guidance-distilled, so "
+ "there is no negative prompt and no unconditional branch."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration),
+ ComponentSpec("tokenizer", Qwen2TokenizerFast),
+ ComponentSpec("processor", Qwen3VLProcessor),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt", description="The prompt to guide generation, a single string."),
+ InputParam(
+ name="normalized_references",
+ type_hint=list[MiniMaxH3Reference],
+ required=True,
+ description="The references normalized by the setup step, in packed order.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template(
+ "prompt_embeds",
+ description=(
+ "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the "
+ "50th decoder layer of the Qwen3-VL conditioner."
+ ),
+ ),
+ OutputParam(
+ "text_token_tags",
+ type_hint=torch.Tensor,
+ description=(
+ "The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video."
+ ),
+ ),
+ ]
+
+ @staticmethod
+ def _sample_video_condition_frames(
+ frames: np.ndarray, fps: float, sample_fps: float, temporal_patch: int
+ ) -> tuple[list[np.ndarray], list[float]]:
+ r"""
+ Sample the frames the conditioner sees from a normalized reference video, and label their vision blocks.
+
+ The conditioner reads a reference at `sample_fps`: every `fps / sample_fps`-th frame, deduplicated. Qwen3-VL
+ then merges the sampled frames in groups of `temporal_patch` — repeating the last one when the count does not
+ divide — and a merged group is labelled with the mean of its timestamps, which `"<{timestamp:.1f} seconds>"`
+ renders with Python's round-half-to-even, so the first block of a 2 fps pair is `"<0.2 seconds>"` rather than
+ `"<0.3 seconds>"`.
+
+ Args:
+ frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The normalized reference video.
+ fps (`float`): The rate `frames` is at, i.e. `components.fps`.
+ sample_fps (`float`): The rate the conditioner reads at, i.e. the block's `video_sample_fps`.
+ temporal_patch (`int`): Qwen3-VL's temporal patch, read off the processor.
+
+ Returns:
+ `tuple[list[np.ndarray], list[float]]`: the sampled frames and one timestamp per vision block.
+ """
+ stride = fps / sample_fps
+ indices, cursor = [], 0.0
+ while round(cursor) < frames.shape[0]:
+ if not indices or round(cursor) > indices[-1]:
+ indices.append(round(cursor))
+ cursor += stride
+
+ timestamps = [index / sample_fps for index in range(len(indices))]
+ timestamps += [timestamps[-1]] * (-len(timestamps) % temporal_patch)
+ block_timestamps = [
+ (timestamps[index] + timestamps[index + temporal_patch - 1]) / 2
+ for index in range(0, len(timestamps), temporal_patch)
+ ]
+ return [frames[index] for index in indices], block_timestamps
+
+ def _gather_vision_features(
+ self, processor, references: list[MiniMaxH3Reference], fps: float
+ ) -> tuple[dict, list[int], list[int], list[list[float]]]:
+ r"""
+ Run the references' pixels through the conditioner's processors, batched per modality.
+
+ The vision tensors are batched per modality while the presentation is tokenized in request order; the two
+ agree because the filtering here preserves relative order within each modality and Qwen3-VL fills the n-th
+ pad *run* of a modality with the n-th entry of that modality's batch. Audio contributes nothing — a waveform
+ never reaches the conditioner.
+
+ Returns:
+ `tuple`: the vision tensors by the conditioner's parameter names, the vision token count per image
+ reference, the per-block token count per video reference, and the block timestamps per video reference.
+ """
+ merge_size = processor.image_processor.merge_size**2
+ vision_inputs = {}
+
+ image_token_counts = []
+ images = [reference.image for reference in references if reference.kind == "image"]
+ if images:
+ image_features = processor.image_processor(images=images, return_tensors="pt")
+ vision_inputs["pixel_values"] = image_features["pixel_values"]
+ vision_inputs["image_grid_thw"] = image_features["image_grid_thw"]
+ image_token_counts = [int(grid.prod()) // merge_size for grid in image_features["image_grid_thw"]]
+
+ video_block_token_counts, video_block_timestamps = [], []
+ videos = [reference for reference in references if reference.kind == "video"]
+ if videos:
+ temporal_patch = processor.video_processor.temporal_patch_size
+ sampled = [
+ self._sample_video_condition_frames(reference.frames, fps, self.video_sample_fps, temporal_patch)
+ for reference in videos
+ ]
+ video_block_timestamps = [timestamps for _, timestamps in sampled]
+ video_features = processor.video_processor(
+ videos=[np.stack(frames) for frames, _ in sampled], do_sample_frames=False, return_tensors="pt"
+ )
+ vision_inputs["pixel_values_videos"] = video_features["pixel_values_videos"]
+ vision_inputs["video_grid_thw"] = video_features["video_grid_thw"]
+ video_block_token_counts = [
+ int(grid[1]) * int(grid[2]) // merge_size for grid in video_features["video_grid_thw"]
+ ]
+ for timestamps, grid in zip(video_block_timestamps, video_features["video_grid_thw"]):
+ if int(grid[0]) != len(timestamps):
+ raise ValueError(
+ f"The processor merged a reference video into {int(grid[0])} vision blocks, but MiniMax-H3 "
+ f"labels {len(timestamps)} of them."
+ )
+
+ return vision_inputs, image_token_counts, video_block_token_counts, video_block_timestamps
+
+ @staticmethod
+ def _build_presentation(
+ tokenizer,
+ prompt: str,
+ references: list[MiniMaxH3Reference],
+ image_token_counts: list[int],
+ video_block_token_counts: list[int],
+ video_block_timestamps: list[list[float]],
+ text_tag: int = 1,
+ video_tag: int = 0,
+ ) -> tuple[list[int], list[int]]:
+ r"""
+ Tokenize MiniMax-H3's presentation of a `ref2va` request.
+
+ Every reference prepends a label, in packed order and numbered per modality: `": "` plus a vision
+ block for an image, `": "` alone for audio — a waveform never reaches the conditioner — and
+ `": "` plus one timestamped vision block per merged frame pair for a video. A video that carries
+ sound is labelled `": "` *before* `": "`, mirroring the order its rows are packed in. The
+ prompt follows verbatim, with no chat template and no special tokens.
+
+ Args:
+ tokenizer (`Qwen2TokenizerFast`): Tokenizer of the conditioner.
+ prompt (`str`): The prompt, appended verbatim.
+ references (`list[MiniMaxH3Reference]`): The normalized references, in packed order.
+ image_token_counts (`list[int]`): Number of vision tokens of every image reference's block.
+ video_block_token_counts (`list[int]`): Number of vision tokens per block of every video reference.
+ video_block_timestamps (`list[list[float]]`): The timestamp of every vision block, per video reference.
+ text_tag (`int`, *optional*, defaults to 1): MiniMax-H3's modality tag for a text row.
+ video_tag (`int`, *optional*, defaults to 0): MiniMax-H3's modality tag for a vision block's rows.
+
+ Returns:
+ `tuple[list[int], list[int]]`: the token ids and their modality tags.
+ """
+
+ def text(value: str) -> tuple[list[int], list[int]]:
+ token_ids = tokenizer(value, add_special_tokens=False)["input_ids"]
+ return token_ids, [text_tag] * len(token_ids)
+
+ def vision(pad_token: str, num_tokens: int) -> tuple[list[int], list[int]]:
+ token_ids = (
+ [tokenizer.convert_tokens_to_ids("<|vision_start|>")]
+ + [tokenizer.convert_tokens_to_ids(pad_token)] * num_tokens
+ + [tokenizer.convert_tokens_to_ids("<|vision_end|>")]
+ )
+ return token_ids, [video_tag] * len(token_ids)
+
+ token_ids, token_tags = [], []
+
+ def emit(segment: tuple[list[int], list[int]]) -> None:
+ token_ids.extend(segment[0])
+ token_tags.extend(segment[1])
+
+ counts = {"image": 0, "video": 0, "audio": 0}
+ for reference in references:
+ if reference.has_audio:
+ counts["audio"] += 1
+ emit(text(f": "))
+ if reference.kind == "image":
+ counts["image"] += 1
+ emit(text(f": "))
+ emit(vision("<|image_pad|>", image_token_counts[counts["image"] - 1]))
+ elif reference.kind == "video":
+ counts["video"] += 1
+ emit(text(f": "))
+ for timestamp in video_block_timestamps[counts["video"] - 1]:
+ # `"{:.1f}"` rounds half to even, so the mean of a 2 fps pair renders as "<0.2 seconds>".
+ emit(text(f"<{timestamp:.1f} seconds>"))
+ emit(vision("<|video_pad|>", video_block_token_counts[counts["video"] - 1]))
+ emit(text(prompt))
+ return token_ids, token_tags
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ if not isinstance(block_state.prompt, str):
+ raise ValueError(
+ "MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got "
+ f"{type(block_state.prompt)}."
+ )
+ references = block_state.normalized_references
+
+ # 1. Vision features, batched per modality — audio contributes nothing to the conditioner.
+ vision_inputs, image_token_counts, video_token_counts, video_timestamps = self._gather_vision_features(
+ components.processor, references, components.fps
+ )
+
+ # 2. The presentation, tokenized in request order.
+ token_ids, token_tags = self._build_presentation(
+ components.tokenizer,
+ block_state.prompt,
+ references,
+ image_token_counts,
+ video_token_counts,
+ video_timestamps,
+ text_tag=components.text_tag,
+ video_tag=components.video_tag,
+ )
+
+ # 3. One conditioner call. The block emits the conditioner's dtype, as every other model does — it has no
+ # denoiser of its own, since it is meant to run on its own.
+ block_state.prompt_embeds = get_qwen3vl_prompt_embeds(
+ components.text_encoder,
+ components.processor,
+ token_ids,
+ vision_inputs,
+ text_encoder_layer=components.text_encoder_layer,
+ device=components._execution_device,
+ dtype=components.text_encoder.dtype,
+ )
+ block_state.text_token_tags = torch.tensor(token_tags, dtype=torch.long)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class MiniMaxH3Ref2VAReferenceEncoderStep(ModularPipelineBlocks):
+ model_name = "minimax-h3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Encodes the `ref2va` references — image and video references through the video VAE, soundtracks through "
+ "the audio VAE. They are the anchors of the whole denoising loop, which only ever writes the generated "
+ "rows; the prepare-latents step noises the visual ones to MiniMax-H3's conditioning level and packs them, "
+ "while soundtracks ride along clean at `t = 1.0`. The latent geometry of every reference is the shape of "
+ "what this emits, which is what the packed layout is built from."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLMiniMaxH3),
+ ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="normalized_references",
+ type_hint=list[MiniMaxH3Reference],
+ required=True,
+ description="The references normalized by the setup step, in packed order.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "condition_latents",
+ type_hint=list[torch.Tensor],
+ description=(
+ "The encoded video conditioning latents of the image and video references, one `(1, "
+ "latent_channels, num_latent_frames, latent_height, latent_width)` tensor each in packed order, "
+ "or None when the references carry none."
+ ),
+ ),
+ OutputParam(
+ "audio_condition_latents",
+ type_hint=list[torch.Tensor],
+ description=(
+ "The clean audio conditioning rows of the reference soundtracks, one `(num_audio_latents * 2, "
+ "audio_latent_channels)` tensor per audio-bearing reference in packed order. One entry per "
+ "reference rather than one concatenated block, because the packed layout is built from the row "
+ "count of each."
+ ),
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean).view(1, 1, -1)
+ audio_latents_std = torch.tensor(components.audio_vae.config.latents_std).view(1, 1, -1)
+
+ condition_latents, audio_condition_latents = [], []
+ for reference in block_state.normalized_references:
+ if reference.kind == "image":
+ pixels = torch.from_numpy(np.array(reference.image)).to(device).permute(2, 0, 1)[None, :, None]
+ condition_latents.append(
+ encode_vae_condition(
+ components.vae,
+ pixels,
+ components.pixel_mean,
+ components.pixel_std,
+ components.keyframe_encode_seed,
+ )
+ )
+ elif reference.kind == "video":
+ # Snap *down* to `17 * n + 5` so the VAE encodes without padding; this only bites when the reference
+ # is shorter than the target, whose own frame count already has that form.
+ frames_per_chunk = components.vae_frames_per_chunk
+ latents_per_chunk = components.vae_latents_per_chunk
+ num_frames = reference.frames.shape[0]
+ num_frames = (
+ max(1, (num_frames - latents_per_chunk) // frames_per_chunk) * frames_per_chunk + latents_per_chunk
+ )
+ pixels = torch.from_numpy(reference.frames[:num_frames].copy()).to(device).permute(3, 0, 1, 2)[None]
+ condition_latents.append(
+ encode_vae_condition(
+ components.vae,
+ pixels,
+ components.pixel_mean,
+ components.pixel_std,
+ components.keyframe_encode_seed,
+ )
+ )
+
+ if reference.has_audio:
+ posterior = components.audio_vae.encode(reference.audio.to(device)[:, None], return_dict=False)[0]
+ # Channel-major rows: the two stereo channels are two batch items of the mono audio VAE. Soundtracks
+ # take the posterior *mean*, and are never sampled.
+ latents = posterior.mode().float().cpu().transpose(1, 2)
+ normalized = (latents - audio_latents_mean) / audio_latents_std
+ audio_condition_latents.append(normalized.reshape(-1, components.audio_latent_channels))
+
+ block_state.condition_latents = condition_latents
+ block_state.audio_condition_latents = audio_condition_latents
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py
new file mode 100644
index 0000000000000000000000000000000000000000..81bad846e31f1753cac2f993c9da4f5549d27830
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_blocks_minimax_h3.py
@@ -0,0 +1,779 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ MiniMaxH3PrepareLatentsStep,
+ MiniMaxH3PrepareLayoutStep,
+ MiniMaxH3Ref2VAPrepareLayoutStep,
+ MiniMaxH3SetTimestepsStep,
+)
+from .before_encoder import MiniMaxH3Ref2VASetupStep, MiniMaxH3ResizeStep
+from .decoders import MiniMaxH3AfterDenoiseStep, MiniMaxH3AudioDecodeStep, MiniMaxH3VideoDecodeStep
+from .denoise import MiniMaxH3DenoiseStep, MiniMaxH3Ref2VADenoiseStep
+from .encoders import (
+ MiniMaxH3KeyframeVaeEncoderStep,
+ MiniMaxH3Ref2VAReferenceEncoderStep,
+ MiniMaxH3Ref2VATextEncoderStep,
+ MiniMaxH3TextEncoderStep,
+)
+
+
+def _generation_outputs() -> list[OutputParam]:
+ r"""Video and audio are generated jointly; muxing them into one file is up to the caller."""
+ return [
+ OutputParam.template("videos", description="The generated video."),
+ OutputParam(
+ "audio",
+ type_hint=torch.Tensor,
+ description="The generated soundtrack, of shape `(1, 2, num_samples)`.",
+ ),
+ OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated soundtrack in Hz."),
+ ]
+
+
+# auto_docstring
+class MiniMaxH3AutoResizeStep(ConditionalPipelineBlocks):
+ """
+ Keyframe canvas block.
+ - MiniMaxH3ResizeStep runs for the `fl2va` workflow, whichever of the two keyframes is given.
+ - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout step falls back to MiniMax-H3's own 16:9 canvas.
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived from its own
+ aspect ratio.
+ last_image (`Image`, *optional*):
+ Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` it is the
+ follower of the two and is cover-cropped onto the canvas.
+ height (`int`, *optional*):
+ Height of the generated video in pixels, a multiple of 32.
+ width (`int`, *optional*):
+ Width of the generated video in pixels, a multiple of 32.
+
+ Outputs:
+ height (`int`):
+ Resolved height of the generated video in pixels.
+ width (`int`):
+ Resolved width of the generated video in pixels.
+ keyframes (`list`):
+ The keyframes put onto the target canvas, in packed order.
+ keyframe_anchors (`tuple`):
+ Which end of the video every keyframe is anchored to, in packed order. Positional with `keyframes`, so both are
+ resolved here.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3ResizeStep]
+ block_names = ["keyframes"]
+ block_trigger_inputs = ["image", "last_image"]
+ default_block_name = None
+
+ def select_block(self, **kwargs) -> str | None:
+ if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
+ return "keyframes"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Keyframe canvas block.\n"
+ + " - MiniMaxH3ResizeStep runs for the `fl2va` workflow, whichever of the two keyframes is given.\n"
+ + " - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped and the layout "
+ "step falls back to MiniMax-H3's own 16:9 canvas."
+ )
+
+
+# auto_docstring
+class MiniMaxH3AutoKeyframeVaeEncoderStep(ConditionalPipelineBlocks):
+ """
+ Keyframe conditioning block.
+ - MiniMaxH3KeyframeVaeEncoderStep runs for the `fl2va` workflow, whichever of the two keyframes is given.
+ - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped.
+
+ Components:
+ vae (`AutoencoderKLMiniMaxH3`)
+
+ Inputs:
+ keyframes (`list`, *optional*):
+ The keyframes put onto the target canvas, in packed order.
+
+ Outputs:
+ condition_latents (`list`):
+ The normalized video conditioning latents, one `(1, latent_channels, 1, latent_height, latent_width)` tensor per
+ keyframe, in packed order. One entry per condition is what the prepare-latents step draws its noise against, so the
+ list is the unit the request's generator is consumed in.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3KeyframeVaeEncoderStep]
+ block_names = ["keyframes"]
+ block_trigger_inputs = ["image", "last_image"]
+ default_block_name = None
+
+ def select_block(self, **kwargs) -> str | None:
+ if kwargs.get("image") is not None or kwargs.get("last_image") is not None:
+ return "keyframes"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Keyframe conditioning block.\n"
+ + " - MiniMaxH3KeyframeVaeEncoderStep runs for the `fl2va` workflow, whichever of the two keyframes is "
+ "given.\n" + " - when neither `image` nor `last_image` is provided (`t2va`), this block is skipped."
+ )
+
+
+# auto_docstring
+class MiniMaxH3AutoBeforeEncodeStep(AutoPipelineBlocks):
+ """
+ Media preparation block.
+ - `MiniMaxH3Ref2VASetupStep` runs when `references` is provided (`ref2va`): it resolves the plan and normalizes every reference onto MiniMax-H3's own rates and resolutions.
+ - otherwise the keyframe canvas block runs, which itself is skipped for a text-only request.
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ references (`list`, *optional*):
+ The references to condition on, **in the order the model should read them**: the order labels them in the prompt
+ presentation and lays them out on the shared rotary clock, so a different order is a different request. One
+ dataclass per modality, all holding in-memory media — a [`MiniMaxH3ImageReference`] (at most 9), a
+ [`MiniMaxH3VideoReference`] at its own `fps` (at most 3, whose `audio` soundtrack is conditioned on as well), or a
+ [`MiniMaxH3AudioReference`] at its own `sample_rate` (at most 3) — for at most 12 references in total, and audio
+ references cannot be the only ones. These blocks never open a media file: decode with each class's `from_file`
+ classmethod, which brings the rates along.
+ height (`int`, *optional*):
+ Height of the generated video in pixels, a multiple of 32.
+ width (`int`, *optional*):
+ Width of the generated video in pixels, a multiple of 32.
+ num_frames (`int`, *optional*):
+ Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode;
+ the resulting duration must stay between 5 and 15 seconds. To generate a video as long as a reference soundtrack,
+ pass `round(samples / sample_rate * 24)`.
+ image (`Image`, *optional*):
+ Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived from its own
+ aspect ratio.
+ last_image (`Image`, *optional*):
+ Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` it is the
+ follower of the two and is cover-cropped onto the canvas.
+
+ Outputs:
+ height (`int`):
+ Resolved height of the generated video in pixels.
+ width (`int`):
+ Resolved width of the generated video in pixels.
+ num_frames (`int`):
+ Resolved number of frames, of the form 17 * n + 5.
+ normalized_references (`list`):
+ The references normalized onto MiniMax-H3's own rates and resolutions, in packed order: the same public reference
+ types the request passed in, with an image resized to its own 2048 pixel short edge, a video resampled onto 24 fps
+ and onto the canvas its own aspect ratio resolves to, and a soundtrack put on the audio VAE's sample rate and
+ truncated to the generated duration.
+ keyframes (`list`):
+ The keyframes put onto the target canvas, in packed order.
+ keyframe_anchors (`tuple`):
+ Which end of the video every keyframe is anchored to, in packed order. Positional with `keyframes`, so both are
+ resolved here.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3AutoResizeStep]
+ block_names = ["ref2va", "keyframes"]
+ block_trigger_inputs = ["references", None]
+
+ @property
+ def description(self):
+ return (
+ "Media preparation block.\n"
+ + " - `MiniMaxH3Ref2VASetupStep` runs when `references` is provided (`ref2va`): it resolves the plan and "
+ "normalizes every reference onto MiniMax-H3's own rates and resolutions.\n"
+ + " - otherwise the keyframe canvas block runs, which itself is skipped for a text-only request."
+ )
+
+
+# auto_docstring
+class MiniMaxH3AutoTextEncoderStep(AutoPipelineBlocks):
+ """
+ Text encoder block. Both branches encode MiniMax-H3's presentation of the request with the Qwen3-VL conditioner.
+ - `MiniMaxH3Ref2VATextEncoderStep` runs when `references` is provided (`ref2va`), labelling every reference in the presentation.
+ - `MiniMaxH3TextEncoderStep` runs otherwise, labelling the keyframes when there are any.
+
+ Components:
+ text_encoder (`Qwen3VLForConditionalGeneration`)
+ tokenizer (`Qwen2Tokenizer`)
+ processor (`Qwen3VLProcessor`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt to guide generation, a single string.
+ normalized_references (`list`, *optional*):
+ The references normalized by the setup step, in packed order.
+ keyframes (`list`, *optional*):
+ The keyframes put onto the target canvas, in packed order (empty or None for `t2va`).
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the 50th decoder layer
+ of the Qwen3-VL conditioner.
+ text_token_tags (`Tensor`):
+ The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VATextEncoderStep, MiniMaxH3TextEncoderStep]
+ block_names = ["ref2va", "t2va"]
+ block_trigger_inputs = ["references", None]
+
+ @property
+ def description(self):
+ return (
+ "Text encoder block. Both branches encode MiniMax-H3's presentation of the request with the Qwen3-VL "
+ "conditioner.\n"
+ + " - `MiniMaxH3Ref2VATextEncoderStep` runs when `references` is provided (`ref2va`), labelling every "
+ "reference in the presentation.\n"
+ + " - `MiniMaxH3TextEncoderStep` runs otherwise, labelling the keyframes when there are any."
+ )
+
+
+# auto_docstring
+class MiniMaxH3AutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ VAE encoder block.
+ - `MiniMaxH3Ref2VAReferenceEncoderStep` runs when `references` is provided (`ref2va`).
+ - otherwise the keyframe conditioning block runs, which itself is skipped for a text-only request.
+
+ Components:
+ vae (`AutoencoderKLMiniMaxH3`)
+ audio_vae (`AutoencoderKLMiniMaxH3Audio`)
+
+ Inputs:
+ normalized_references (`list`, *optional*):
+ The references normalized by the setup step, in packed order.
+ keyframes (`list`, *optional*):
+ The keyframes put onto the target canvas, in packed order.
+
+ Outputs:
+ condition_latents (`list`):
+ The encoded video conditioning latents of the image and video references, one `(1, latent_channels,
+ num_latent_frames, latent_height, latent_width)` tensor each in packed order, or None when the references carry
+ none.
+ audio_condition_latents (`list`):
+ The clean audio conditioning rows of the reference soundtracks, one `(num_audio_latents * 2,
+ audio_latent_channels)` tensor per audio-bearing reference in packed order. One entry per reference rather than one
+ concatenated block, because the packed layout is built from the row count of each.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VAReferenceEncoderStep, MiniMaxH3AutoKeyframeVaeEncoderStep]
+ block_names = ["ref2va", "keyframes"]
+ block_trigger_inputs = ["references", None]
+
+ @property
+ def description(self):
+ return (
+ "VAE encoder block.\n"
+ + " - `MiniMaxH3Ref2VAReferenceEncoderStep` runs when `references` is provided (`ref2va`).\n"
+ + " - otherwise the keyframe conditioning block runs, which itself is skipped for a text-only request."
+ )
+
+
+# auto_docstring
+class MiniMaxH3CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for `t2va` / `fl2va`: builds the packed layout, draws and packs the noise, plans the per-row timesteps and runs the denoising loop against the `transformer` partition.
+
+ Components:
+ scheduler (`MiniMaxH3Scheduler`)
+ audio_scheduler (`MiniMaxH3Scheduler`)
+ transformer (`MiniMaxH3Transformer3DModel`)
+
+ Inputs:
+ text_token_tags (`Tensor`):
+ The per-row modality tag of every row of `prompt_embeds`.
+ height (`int`, *optional*):
+ Height of the generated video in pixels, a multiple of 32.
+ width (`int`, *optional*):
+ Width of the generated video in pixels, a multiple of 32.
+ num_frames (`int`, *optional*, defaults to 124):
+ Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode;
+ the resulting duration must stay between 5 and 15 seconds.
+ keyframe_anchors (`tuple`, *optional*, defaults to ()):
+ Which end of the video every keyframe is anchored to, in packed order.
+ generator (`Generator`, *optional*):
+ The generator of the request. The video noise is drawn from it first, then the audio noise.
+ latents (`Tensor`, *optional*):
+ Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the
+ draw.
+ audio_latents (`Tensor`, *optional*):
+ Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.
+ condition_latents (`list`, *optional*):
+ The encoded video conditioning latents, one `(1, latent_channels, num_latent_frames, latent_height, latent_width)`
+ tensor per condition in packed order, or None for a request that has none. Noised and packed here.
+ audio_condition_latents (`list`, *optional*):
+ The audio conditioning rows to prepend, one tensor per audio-bearing reference in packed order. Empty for a request
+ that has none, which is every `t2va` and `fl2va` one.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ **denoiser_input_fields (`None`, *optional*):
+ The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and
+ the three row-index tensors.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ height (`int`):
+ Resolved height of the generated video in pixels.
+ width (`int`):
+ Resolved width of the generated video in pixels.
+ num_frames (`int`):
+ Resolved number of frames, of the form 17 * n + 5.
+ num_latent_frames (`int`):
+ Number of generated video latent frames.
+ latent_height (`int`):
+ Height of the generated video latents.
+ latent_width (`int`):
+ Width of the generated video latents.
+ num_audio_latents (`int`):
+ Number of generated audio latents per channel.
+ position_ids (`Tensor`):
+ The `(t, h, w)` rotary coordinate of every row, in float64.
+ token_tags (`Tensor`):
+ The modality tag of every row.
+ video_indices (`Tensor`):
+ Sequence positions of the video rows, conditioning rows first.
+ audio_indices (`Tensor`):
+ Sequence positions of the audio rows, reference rows first.
+ text_indices (`Tensor`):
+ Sequence positions of the text rows.
+ num_condition_video_rows (`int`):
+ How many leading video rows are conditioning rows rather than generated rows.
+ num_condition_audio_rows (`int`):
+ How many leading audio rows are reference rows rather than generated rows.
+ latents (`Tensor`):
+ The video rows of the packed sequence, conditioning rows first.
+ audio_latents (`Tensor`):
+ The channel-major audio rows of the packed sequence, reference rows first.
+ timesteps (`Tensor`):
+ Timesteps of the video schedule.
+ audio_timesteps (`Tensor`):
+ Timesteps of the audio schedule.
+ row_timestep_plan (`list`):
+ One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the index of every row
+ into them.
+ noise_pred (`Tensor`):
+ Predicted velocity of the video rows of the sequence.
+ audio_noise_pred (`Tensor`):
+ Predicted velocity of the audio rows of the sequence.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [
+ MiniMaxH3PrepareLayoutStep,
+ MiniMaxH3PrepareLatentsStep,
+ MiniMaxH3SetTimestepsStep,
+ MiniMaxH3DenoiseStep,
+ ]
+ block_names = ["prepare_layout", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core denoising workflow for `t2va` / `fl2va`: builds the packed layout, draws and packs the noise, "
+ "plans the per-row timesteps and runs the denoising loop against the `transformer` partition."
+ )
+
+
+# auto_docstring
+class MiniMaxH3Ref2VACoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for `ref2va`: builds the packed layout with one block per reference, draws and packs the noise, plans the per-row timesteps and runs the denoising loop against the `transformer_ref` partition.
+
+ Components:
+ scheduler (`MiniMaxH3Scheduler`)
+ audio_scheduler (`MiniMaxH3Scheduler`)
+ transformer_ref (`MiniMaxH3Transformer3DModel`)
+
+ Inputs:
+ text_token_tags (`Tensor`):
+ The per-row modality tag of every row of `prompt_embeds`.
+ normalized_references (`list`):
+ The references normalized by the setup step, in packed order.
+ condition_latents (`list`):
+ The encoded video conditioning latents, one per image and video reference in packed order. Their shape is where
+ every reference block's geometry comes from.
+ audio_condition_latents (`list`):
+ The encoded audio conditioning rows, one per audio-bearing reference in packed order.
+ height (`int`):
+ Height of the generated video in pixels.
+ width (`int`):
+ Width of the generated video in pixels.
+ num_frames (`int`):
+ Resolved number of frames, of the form 17 * n + 5.
+ generator (`Generator`, *optional*):
+ The generator of the request. The video noise is drawn from it first, then the audio noise.
+ latents (`Tensor`, *optional*):
+ Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the
+ draw.
+ audio_latents (`Tensor`, *optional*):
+ Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ **denoiser_input_fields (`None`, *optional*):
+ The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and
+ the three row-index tensors.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+
+ Outputs:
+ height (`int`):
+ Resolved height of the generated video in pixels.
+ width (`int`):
+ Resolved width of the generated video in pixels.
+ num_frames (`int`):
+ Resolved number of frames, of the form 17 * n + 5.
+ num_latent_frames (`int`):
+ Number of generated video latent frames.
+ latent_height (`int`):
+ Height of the generated video latents.
+ latent_width (`int`):
+ Width of the generated video latents.
+ num_audio_latents (`int`):
+ Number of generated audio latents per channel.
+ position_ids (`Tensor`):
+ The `(t, h, w)` rotary coordinate of every row, in float64.
+ token_tags (`Tensor`):
+ The modality tag of every row.
+ video_indices (`Tensor`):
+ Sequence positions of the video rows, conditioning rows first.
+ audio_indices (`Tensor`):
+ Sequence positions of the audio rows, reference rows first.
+ text_indices (`Tensor`):
+ Sequence positions of the text rows.
+ num_condition_video_rows (`int`):
+ How many leading video rows are conditioning rows rather than generated rows.
+ num_condition_audio_rows (`int`):
+ How many leading audio rows are reference rows rather than generated rows.
+ latents (`Tensor`):
+ The video rows of the packed sequence, conditioning rows first.
+ audio_latents (`Tensor`):
+ The channel-major audio rows of the packed sequence, reference rows first.
+ timesteps (`Tensor`):
+ Timesteps of the video schedule.
+ audio_timesteps (`Tensor`):
+ Timesteps of the audio schedule.
+ row_timestep_plan (`list`):
+ One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the index of every row
+ into them.
+ noise_pred (`Tensor`):
+ Predicted velocity of the video rows of the sequence.
+ audio_noise_pred (`Tensor`):
+ Predicted velocity of the audio rows of the sequence.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [
+ MiniMaxH3Ref2VAPrepareLayoutStep,
+ MiniMaxH3PrepareLatentsStep,
+ MiniMaxH3SetTimestepsStep,
+ MiniMaxH3Ref2VADenoiseStep,
+ ]
+ block_names = ["prepare_layout", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core denoising workflow for `ref2va`: builds the packed layout with one block per reference, draws and "
+ "packs the noise, plans the per-row timesteps and runs the denoising loop against the `transformer_ref` "
+ "partition."
+ )
+
+
+# auto_docstring
+class MiniMaxH3AutoDenoiseStep(AutoPipelineBlocks):
+ """
+ Denoise block.
+ - `MiniMaxH3Ref2VACoreDenoiseStep` runs when `references` is provided (`ref2va`), against the `transformer_ref` checkpoint partition.
+ - `MiniMaxH3CoreDenoiseStep` runs otherwise, against the `transformer` partition.
+
+ Components:
+ scheduler (`MiniMaxH3Scheduler`)
+ audio_scheduler (`MiniMaxH3Scheduler`)
+ transformer_ref (`MiniMaxH3Transformer3DModel`)
+ transformer (`MiniMaxH3Transformer3DModel`)
+
+ Inputs:
+ text_token_tags (`Tensor`):
+ The per-row modality tag of every row of `prompt_embeds`.
+ normalized_references (`list`, *optional*):
+ The references normalized by the setup step, in packed order.
+ condition_latents (`list`, *optional*):
+ The encoded video conditioning latents, one per image and video reference in packed order. Their shape is where
+ every reference block's geometry comes from.
+ audio_condition_latents (`list`, *optional*):
+ The encoded audio conditioning rows, one per audio-bearing reference in packed order.
+ height (`int`, *optional*):
+ Height of the generated video in pixels.
+ width (`int`, *optional*):
+ Width of the generated video in pixels.
+ num_frames (`int`, *optional*, defaults to 124):
+ Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode;
+ the resulting duration must stay between 5 and 15 seconds.
+ generator (`Generator`, *optional*):
+ The generator of the request. The video noise is drawn from it first, then the audio noise.
+ latents (`Tensor`):
+ Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the
+ draw.
+ audio_latents (`Tensor`):
+ Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ **denoiser_input_fields (`None`, *optional*):
+ The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and
+ the three row-index tensors.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ keyframe_anchors (`tuple`, *optional*, defaults to ()):
+ Which end of the video every keyframe is anchored to, in packed order.
+
+ Outputs:
+ height (`int`):
+ Resolved height of the generated video in pixels.
+ width (`int`):
+ Resolved width of the generated video in pixels.
+ num_frames (`int`):
+ Resolved number of frames, of the form 17 * n + 5.
+ num_latent_frames (`int`):
+ Number of generated video latent frames.
+ latent_height (`int`):
+ Height of the generated video latents.
+ latent_width (`int`):
+ Width of the generated video latents.
+ num_audio_latents (`int`):
+ Number of generated audio latents per channel.
+ position_ids (`Tensor`):
+ The `(t, h, w)` rotary coordinate of every row, in float64.
+ token_tags (`Tensor`):
+ The modality tag of every row.
+ video_indices (`Tensor`):
+ Sequence positions of the video rows, conditioning rows first.
+ audio_indices (`Tensor`):
+ Sequence positions of the audio rows, reference rows first.
+ text_indices (`Tensor`):
+ Sequence positions of the text rows.
+ num_condition_video_rows (`int`):
+ How many leading video rows are conditioning rows rather than generated rows.
+ num_condition_audio_rows (`int`):
+ How many leading audio rows are reference rows rather than generated rows.
+ latents (`Tensor`):
+ The video rows of the packed sequence, conditioning rows first.
+ audio_latents (`Tensor`):
+ The channel-major audio rows of the packed sequence, reference rows first.
+ timesteps (`Tensor`):
+ Timesteps of the video schedule.
+ audio_timesteps (`Tensor`):
+ Timesteps of the audio schedule.
+ row_timestep_plan (`list`):
+ One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the index of every row
+ into them.
+ noise_pred (`Tensor`):
+ Predicted velocity of the video rows of the sequence.
+ audio_noise_pred (`Tensor`):
+ Predicted velocity of the audio rows of the sequence.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VACoreDenoiseStep, MiniMaxH3CoreDenoiseStep]
+ block_names = ["ref2va", "t2va"]
+ block_trigger_inputs = ["references", None]
+
+ @property
+ def description(self):
+ return (
+ "Denoise block.\n"
+ + " - `MiniMaxH3Ref2VACoreDenoiseStep` runs when `references` is provided (`ref2va`), against the "
+ "`transformer_ref` checkpoint partition.\n"
+ + " - `MiniMaxH3CoreDenoiseStep` runs otherwise, against the `transformer` partition."
+ )
+
+
+# auto_docstring
+class MiniMaxH3DecodeStep(SequentialPipelineBlocks):
+ """
+ Decodes the denoised rows of the packed sequence into the generated video and its soundtrack.
+
+ Components:
+ vae (`AutoencoderKLMiniMaxH3`)
+ video_processor (`VideoProcessor`)
+ audio_vae (`AutoencoderKLMiniMaxH3Audio`)
+
+ Inputs:
+ latents (`Tensor`):
+ The generated video latents.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np' or 'pt'.
+ audio_latents (`Tensor`):
+ The generated audio latents, one batch item per stereo channel.
+
+ Outputs:
+ videos (`list`):
+ The generated video.
+ audio (`Tensor`):
+ The generated soundtrack, of shape `(1, 2, num_samples)`.
+ sampling_rate (`int`):
+ Sample rate of the generated soundtrack in Hz.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3VideoDecodeStep, MiniMaxH3AudioDecodeStep]
+ block_names = ["video", "audio"]
+
+ @property
+ def description(self):
+ return "Decodes the denoised rows of the packed sequence into the generated video and its soundtrack."
+
+ @property
+ def outputs(self):
+ return _generation_outputs()
+
+
+# auto_docstring
+class MiniMaxH3Blocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline blocks for joint video + audio generation with MiniMax-H3: the `t2va` (text only), `fl2va` (first and/or last keyframe) and `ref2va` (omni-reference) workflows, selected on the `references` and keyframe inputs. Without a `workflow=`, loading the components pulls **both** 61.7GB transformer partitions — pass one to load only what the task needs.
+
+ Supported workflows:
+ - `t2va`: requires `prompt`
+ - `fl2va`: requires `prompt`, `image`
+ - `fl2va_last_frame`: requires `prompt`, `last_image`
+ - `ref2va`: requires `prompt`, `references`
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+ text_encoder (`Qwen3VLForConditionalGeneration`)
+ tokenizer (`Qwen2Tokenizer`)
+ processor (`Qwen3VLProcessor`)
+ vae (`AutoencoderKLMiniMaxH3`)
+ audio_vae (`AutoencoderKLMiniMaxH3Audio`)
+ scheduler (`MiniMaxH3Scheduler`)
+ audio_scheduler (`MiniMaxH3Scheduler`)
+ transformer_ref (`MiniMaxH3Transformer3DModel`)
+ transformer (`MiniMaxH3Transformer3DModel`)
+ video_processor (`VideoProcessor`)
+
+ Inputs:
+ references (`list`, *optional*):
+ The references to condition on, **in the order the model should read them**: the order labels them in the prompt
+ presentation and lays them out on the shared rotary clock, so a different order is a different request. One
+ dataclass per modality, all holding in-memory media — a [`MiniMaxH3ImageReference`] (at most 9), a
+ [`MiniMaxH3VideoReference`] at its own `fps` (at most 3, whose `audio` soundtrack is conditioned on as well), or a
+ [`MiniMaxH3AudioReference`] at its own `sample_rate` (at most 3) — for at most 12 references in total, and audio
+ references cannot be the only ones. These blocks never open a media file: decode with each class's `from_file`
+ classmethod, which brings the rates along.
+ height (`int`, *optional*):
+ Height of the generated video in pixels, a multiple of 32.
+ width (`int`, *optional*):
+ Width of the generated video in pixels, a multiple of 32.
+ num_frames (`int`, *optional*):
+ Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video VAE can decode;
+ the resulting duration must stay between 5 and 15 seconds. To generate a video as long as a reference soundtrack,
+ pass `round(samples / sample_rate * 24)`.
+ image (`Image`, *optional*):
+ Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is derived from its own
+ aspect ratio.
+ last_image (`Image`, *optional*):
+ Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with `image` it is the
+ follower of the two and is cover-cropped onto the canvas.
+ prompt (`str`):
+ The prompt to guide generation, a single string.
+ generator (`Generator`, *optional*):
+ The generator of the request. The video noise is drawn from it first, then the audio noise.
+ latents (`Tensor`):
+ Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, used instead of the
+ draw.
+ audio_latents (`Tensor`):
+ Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ **denoiser_input_fields (`None`, *optional*):
+ The structural description of the packed sequence the transformer reads by name: `token_tags`, `position_ids` and
+ the three row-index tensors.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np' or 'pt'.
+
+ Outputs:
+ videos (`list`):
+ The generated video.
+ audio (`Tensor`):
+ The generated soundtrack, of shape `(1, 2, num_samples)`.
+ sampling_rate (`int`):
+ Sample rate of the generated soundtrack in Hz.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [
+ MiniMaxH3AutoBeforeEncodeStep,
+ MiniMaxH3AutoTextEncoderStep,
+ MiniMaxH3AutoVaeEncoderStep,
+ MiniMaxH3AutoDenoiseStep,
+ MiniMaxH3AfterDenoiseStep,
+ MiniMaxH3DecodeStep,
+ ]
+ block_names = [
+ "before_encode",
+ "text_encoder",
+ "vae_encoder",
+ "denoise",
+ "after_denoise",
+ "decode",
+ ]
+ # One repository holds both checkpoint partitions (`transformer/` for `t2va`/`fl2va`, `transformer_ref/` for
+ # `ref2va`), so the tasks are workflows over one shared `modular_model_index.json`, each declaring only the
+ # components its own half loads. The keys below are the task names a `workflow=` argument to
+ # `ModularPipeline.from_pretrained` takes, so that the component subset of a task can be resolved from the index
+ # without instantiating the blocks first.
+ _workflow_map = {
+ "t2va": {"prompt": True},
+ "fl2va": {"prompt": True, "image": True},
+ "fl2va_last_frame": {"prompt": True, "last_image": True},
+ "ref2va": {"prompt": True, "references": True},
+ }
+
+ @property
+ def description(self):
+ return (
+ "Auto Modular pipeline blocks for joint video + audio generation with MiniMax-H3: the `t2va` (text "
+ "only), `fl2va` (first and/or last keyframe) and `ref2va` (omni-reference) workflows, selected on the "
+ "`references` and keyframe inputs. Without a `workflow=`, loading the components pulls **both** 61.7GB "
+ "transformer partitions — pass one to load only what the task needs."
+ )
+
+ @property
+ def outputs(self):
+ return _generation_outputs()
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..cf6edfa592ed66f9b40cc16e0ab05d7ef18aa54f
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py
@@ -0,0 +1,288 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Per-row modality tags. They index the transformer's AdaLN table, so the values are a checkpoint contract.
+MINIMAX_H3_VIDEO_TAG = 0
+MINIMAX_H3_TEXT_TAG = 1
+MINIMAX_H3_AUDIO_TAG = 2
+
+# MiniMax-H3 generates at a fixed 24 fps and was released for a 768 pixel short edge only, with a soft area cap of
+# 768x1344. The multiple both axes round to follows from
+# the VAE and the transformer, so it is `MiniMaxH3ModularPipeline.canvas_multiple` below and reaches these helpers
+# as an argument.
+MINIMAX_H3_FPS = 24
+MINIMAX_H3_SHORT_EDGE = 768
+MINIMAX_H3_MAX_PIXELS = 768 * 1344
+MINIMAX_H3_MIN_ASPECT_RATIO = 1 / 4
+MINIMAX_H3_MAX_ASPECT_RATIO = 4
+
+# The audio VAE hops 800 samples at 32 kHz, i.e. 40 latents per second. Stereo is carried as two channel-major
+# blocks of audio rows (and as two batch items at the audio VAE boundary, which is mono).
+MINIMAX_H3_AUDIO_LATENTS_PER_SECOND = 40
+MINIMAX_H3_AUDIO_CHANNELS = 2
+
+
+def resolve_canvas_size(aspect_width: float, aspect_height: float, canvas_multiple: int) -> tuple[int, int]:
+ r"""
+ Resolve a display aspect ratio into a MiniMax-H3 canvas.
+
+ The short edge starts at 768, the area is capped at `768 * 1344` and both axes are then rounded to the nearest
+ `canvas_multiple` — so the final area may end up slightly above the pre-rounding budget. Only the ratio of the two
+ arguments matters; pass either the aspect ratio (`16, 9`) or the source dimensions of a keyframe.
+
+ Args:
+ aspect_width (`float`): Width of the target ratio.
+ aspect_height (`float`): Height of the target ratio.
+ canvas_multiple (`int`):
+ What both axes round to, i.e. `components.canvas_multiple` — 32 for the released checkpoint.
+
+ Returns:
+ `tuple[int, int]`: the `(height, width)` of the canvas.
+ """
+ if aspect_width <= 0 or aspect_height <= 0:
+ raise ValueError(f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}.")
+
+ ratio = aspect_width / aspect_height
+ if not MINIMAX_H3_MIN_ASPECT_RATIO <= ratio <= MINIMAX_H3_MAX_ASPECT_RATIO:
+ raise ValueError(
+ f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height} ({ratio:g})."
+ )
+
+ if ratio >= 1.0:
+ width, height = MINIMAX_H3_SHORT_EDGE * ratio, float(MINIMAX_H3_SHORT_EDGE)
+ else:
+ width, height = float(MINIMAX_H3_SHORT_EDGE), MINIMAX_H3_SHORT_EDGE / ratio
+
+ area = width * height
+ if area > MINIMAX_H3_MAX_PIXELS:
+ scale = (MINIMAX_H3_MAX_PIXELS / area) ** 0.5
+ width, height = width * scale, height * scale
+
+ multiple = canvas_multiple
+ return max(multiple, round(height / multiple) * multiple), max(multiple, round(width / multiple) * multiple)
+
+
+def align_num_frames(num_frames: int, frames_per_chunk: int, latents_per_chunk: int) -> int:
+ r"""
+ Snap a frame count up to the next `frames_per_chunk * n + latents_per_chunk` the video VAE can encode.
+
+ Args:
+ num_frames (`int`): The requested number of frames.
+ frames_per_chunk (`int`): Pixel frames the video VAE encodes per chunk, its `clip_length`.
+ latents_per_chunk (`int`): Latent frames a chunk keeps, the VAE's `tokens_chunk_size`.
+
+ Returns:
+ `int`: The aligned number of frames.
+ """
+ if num_frames < 1:
+ raise ValueError(f"`num_frames` must be positive, got {num_frames}.")
+ while num_frames % frames_per_chunk != latents_per_chunk:
+ num_frames += 1
+ return num_frames
+
+
+def video_latent_num_frames(num_frames: int, frames_per_chunk: int, latents_per_chunk: int) -> int:
+ r"""
+ The number of latent frames the video VAE produces for a `17 * n + 5` frame count.
+
+ Args:
+ num_frames (`int`): An aligned number of frames.
+
+ Returns:
+ `int`: The number of latent frames, `5 * n + 2`.
+ """
+ if num_frames % frames_per_chunk != latents_per_chunk:
+ raise ValueError(
+ f"`num_frames` must be of the form {frames_per_chunk} * n + {latents_per_chunk}, got {num_frames}."
+ )
+ return (num_frames - latents_per_chunk) // frames_per_chunk * latents_per_chunk + 2
+
+
+def audio_latent_num_frames(num_frames: int) -> int:
+ r"""
+ The number of audio latents that covers a video of `num_frames` frames at 24 fps.
+
+ Args:
+ num_frames (`int`): The number of video frames.
+
+ Returns:
+ `int`: The number of audio latents, rounded at the 40 Hz latent grid.
+ """
+ return int(round(num_frames / MINIMAX_H3_FPS * MINIMAX_H3_AUDIO_LATENTS_PER_SECOND))
+
+
+class MiniMaxH3ModularPipeline(ModularPipeline):
+ """
+ A ModularPipeline for joint video + audio generation with MiniMax-H3: the `t2va` (text only) and `fl2va` (first
+ and/or last keyframe) workflows against the `transformer/` checkpoint partition, and the `ref2va` (omni-reference)
+ workflow against `transformer_ref/`. One repository holds both partitions, and selecting a workflow loads only its
+ own:
+
+ ```py
+ pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3", workflow="ref2va")
+ ```
+
+ MiniMax-H3 denoises **one packed sequence** that holds the text conditioning, the keyframe conditioning latents,
+ the audio latents and the video latents at once, which is why the blocks pass a row layout around rather than
+ per-modality tensors, and why the pipeline carries two schedulers (`shift = 12.0` for video, `shift = 3.0` for
+ audio) that are stepped inside a single transformer call.
+
+ The checkpoint is guidance-distilled: guidance is baked into the weights, so there is no guider, no
+ `negative_prompt` and no `guidance_scale`, and every step runs exactly one forward pass.
+
+ MiniMax-H3 is modular only: this pipeline and its blocks are the whole integration, there is no
+ `DiffusionPipeline` half. This module carries the model facts every block keys off — the config-derived geometry
+ as properties, the canvas and frame-count arithmetic as functions — and the conditioning, encoding and noise
+ contracts live on the blocks themselves.
+
+ ```py
+ import torch
+ from diffusers import ModularPipeline
+
+ pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3")
+ pipe.load_components(dtype=torch.bfloat16)
+ ```
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "MiniMaxH3Blocks"
+
+ @property
+ def vae_spatial_compression_ratio(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.spatial_compression_ratio
+ return 16
+
+ @property
+ def vae_latent_channels(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.config.latent_channels
+ return 24
+
+ @property
+ def vae_frames_per_chunk(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.config.clip_length
+ return 17
+
+ @property
+ def vae_latents_per_chunk(self):
+ if getattr(self, "vae", None) is not None:
+ return self.vae.tokens_chunk_size
+ return 5
+
+ @property
+ def audio_sampling_rate(self):
+ if getattr(self, "audio_vae", None) is not None:
+ return self.audio_vae.config.sampling_rate
+ return 32000
+
+ @property
+ def audio_latent_channels(self):
+ if getattr(self, "audio_vae", None) is not None:
+ return self.audio_vae.config.latent_channels
+ return 32
+
+ @property
+ def patch_size(self):
+ # One repository holds both checkpoint partitions — `transformer/` for `t2va`/`fl2va`, `transformer_ref/`
+ # for `ref2va` — and a workflow loads only its own, so read whichever is present.
+ for name in ("transformer", "transformer_ref"):
+ if getattr(self, name, None) is not None:
+ return tuple(getattr(self, name).config.patch_size)
+ return (1, 2, 2)
+
+ @property
+ def canvas_multiple(self):
+ r"""What the generated height and width have to be a multiple of, 32 for the released checkpoint."""
+ # A canvas has to survive the VAE's spatial compression and still be a whole number of patch rows wide, so
+ # the multiple is the product of the two.
+ return self.vae_spatial_compression_ratio * self.patch_size[2]
+
+ @property
+ def fps(self):
+ r"""MiniMax-H3's own frame rate. Everything it generates and conditions on is resampled onto it."""
+ return MINIMAX_H3_FPS
+
+ @property
+ def min_duration(self):
+ r"""Shortest video MiniMax-H3 generates, in seconds."""
+ return 5.0
+
+ @property
+ def max_duration(self):
+ r"""Longest video MiniMax-H3 generates, in seconds."""
+ return 15.0
+
+ @property
+ def audio_channels(self):
+ r"""Channels of the generated soundtrack: MiniMax-H3 is stereo, packed channel-major."""
+ return MINIMAX_H3_AUDIO_CHANNELS
+
+ @property
+ def text_encoder_layer(self):
+ r"""
+ Which Qwen3-VL hidden state conditions the transformer.
+
+ MiniMax-H3 reads `hidden_states[50]`, not the final one: the last layer is post-norm and is not the
+ conditioning the released weights were trained against.
+ """
+ return 50
+
+ @property
+ def pixel_mean(self):
+ r"""Per-channel mean the video VAE's input is normalized by, ImageNet's."""
+ return (0.485, 0.456, 0.406)
+
+ @property
+ def pixel_std(self):
+ r"""Per-channel standard deviation the video VAE's input is normalized by, ImageNet's."""
+ return (0.229, 0.224, 0.225)
+
+ @property
+ def keyframe_encode_seed(self):
+ r"""
+ Seed the conditioning posterior is sampled under, independently of the request's own generator.
+
+ Fixed at 42 in the reference implementation, so the same keyframe always encodes to the same anchor.
+ """
+ return 42
+
+ @property
+ def keyframe_noise_aug(self):
+ r"""
+ The `t` a visual conditioning anchor is held at: 0.999, just short of clean.
+
+ The released model was trained with its anchors very slightly noised, so conditioning on exactly `t = 1.0`
+ is off-distribution.
+ """
+ return 0.999
+
+ @property
+ def text_tag(self):
+ r"""The modality tag of a text row of the packed sequence."""
+ return MINIMAX_H3_TEXT_TAG
+
+ @property
+ def video_tag(self):
+ r"""The modality tag of a video row of the packed sequence, which a vision block's rows also carry."""
+ return MINIMAX_H3_VIDEO_TAG
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/references.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/references.py
new file mode 100644
index 0000000000000000000000000000000000000000..c883b17d0bee9ba5b3e857db3e20fab956739bc7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/minimax_h3/references.py
@@ -0,0 +1,317 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""
+The references of a MiniMax-H3 `ref2va` request: one public dataclass per modality.
+
+`ref2va` conditions on an ordered list of references — images, videos (with their soundtrack) and audio clips —
+packed one block per reference ahead of the generated rows. The order is semantic twice over: it fixes the
+`""` / `""` / `""` labels of the prompt presentation, and it advances the shared
+audio/video rotary clock, so a different order is a different request.
+
+Every reference holds in-memory media plus the rate that media carries. No block of this model opens a media file;
+each class decodes a path or a URL through its `from_file` classmethod, which brings the rates along. That is the
+reason to prefer `from_file` over a bare `load_video`: MiniMax-H3 resamples a reference onto its own 24 fps, so a
+video whose real rate is lost on the way in is conditioned on at the wrong speed, silently.
+"""
+
+import contextlib
+import os
+import tempfile
+from dataclasses import dataclass
+from urllib.parse import unquote, urlparse
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+
+from ...utils import is_av_available, load_image
+from ...utils.constants import DIFFUSERS_REQUEST_TIMEOUT
+from .modular_pipeline import MINIMAX_H3_FPS
+
+
+@dataclass
+class MiniMaxH3Reference:
+ r"""
+ Base class of the three references a [`MiniMaxH3Ref2VABlocks`] request conditions on:
+ [`MiniMaxH3ImageReference`], [`MiniMaxH3VideoReference`] and [`MiniMaxH3AudioReference`].
+
+ References are passed to the blocks as a list, **in the order the model should read them**: the order labels them
+ in the prompt presentation and lays them out on the shared rotary clock, so a different order is a different
+ request.
+
+ Every reference holds in-memory media, and the rate that media carries where there is one — MiniMax-H3 resamples a
+ reference onto its own 24 fps and onto the audio VAE's sample rate, so a rate lost on the way in is a request
+ conditioned at the wrong speed. Each class decodes a file through its `from_file` classmethod, along with the
+ rates.
+
+ ```py
+ >>> import numpy as np
+ >>> from diffusers.modular_pipelines.minimax_h3 import (
+ ... MiniMaxH3AudioReference,
+ ... MiniMaxH3ImageReference,
+ ... MiniMaxH3VideoReference,
+ ... )
+
+ >>> references = [
+ ... MiniMaxH3ImageReference.from_file("subject.png"),
+ ... MiniMaxH3VideoReference.from_file("motion_ref.mp4"), # frames, their `fps`, and the soundtrack
+ ... MiniMaxH3AudioReference.from_file("voice.wav"), # waveform and its `sample_rate`
+ ... ]
+
+ >>> # Media a request produced itself declares the rate it was produced at.
+ >>> frames = np.zeros((30, 480, 854, 3), dtype="uint8")
+ >>> reference = MiniMaxH3VideoReference(frames=frames, fps=30.0)
+ ```
+ """
+
+
+@dataclass
+class MiniMaxH3ImageReference(MiniMaxH3Reference):
+ r"""
+ A subject, style or scene reference: at most 9 per request.
+
+ Attributes:
+ image (`PIL.Image.Image`):
+ The reference image. It never binds the generated geometry — it is encoded at a 2048 pixel short edge of
+ its own aspect ratio, whatever canvas the request generates at.
+ """
+
+ image: Image.Image
+
+ kind = "image"
+ has_audio = False
+
+ @classmethod
+ def from_file(cls, media) -> "MiniMaxH3ImageReference":
+ r"""
+ Load an image file into a [`MiniMaxH3ImageReference`], through [`~utils.load_image`].
+
+ Args:
+ media (`str` or `os.PathLike`): Path or URL of the image.
+ """
+ return cls(image=load_image(str(media)))
+
+
+@dataclass
+class MiniMaxH3VideoReference(MiniMaxH3Reference):
+ r"""
+ A motion and camera reference: at most 3 per request, conditioned on together with its own soundtrack.
+
+ Attributes:
+ frames (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`):
+ The reference frames: a list of images, a `(num_frames, height, width, 3)` array or a `(num_frames, 3,
+ height, width)` tensor, `uint8` or floating point over `[0, 1]`.
+ fps (`float`, *optional*, defaults to 24.0):
+ The frame rate `frames` carries, which is what places the reference's vision blocks on the conditioner's 2
+ fps grid. MiniMax-H3's own clock is 24 fps, so any other rate is resampled onto it by dropping and
+ duplicating whole frames — which makes this the field to get right when the frames came from a file.
+ audio (`torch.Tensor` of shape `(channels, num_samples)`, *optional*):
+ This video's soundtrack, mono or stereo, conditioned on as the reference's own rather than as a reference
+ of its own. Left out, the reference conditions on motion alone.
+ sample_rate (`int`, *optional*):
+ The rate `audio` carries its samples at. Left out, it is the audio VAE's own, which leaves the samples
+ untouched; any other rate is resampled onto it.
+ """
+
+ frames: list[Image.Image] | np.ndarray | torch.Tensor
+ fps: float | None = None
+ audio: torch.Tensor | None = None
+ sample_rate: int | None = None
+
+ kind = "video"
+
+ def __post_init__(self):
+ if self.fps is None:
+ self.fps = float(MINIMAX_H3_FPS)
+
+ @property
+ def has_audio(self) -> bool:
+ r"""Whether this reference contributes audio rows, i.e. whether it carries a soundtrack."""
+ return self.audio is not None
+
+ @classmethod
+ def from_file(cls, media) -> "MiniMaxH3VideoReference":
+ r"""
+ Decode a video file into a [`MiniMaxH3VideoReference`], at the resolution, the frame rate and the soundtrack
+ it carries.
+
+ The rates land on the reference, which is the point of decoding this way rather than with
+ [`~utils.load_video`]: MiniMax-H3 resamples a reference onto its own 24 fps, so a frame rate lost on the way
+ in is a request conditioned at the wrong speed, with nothing to raise about it. A container whose metadata is
+ wrong is corrected by overriding `fps` or `sample_rate` on the returned reference.
+
+ Needs [PyAV](https://github.com/PyAV-Org/PyAV).
+
+ Args:
+ media (`str` or `os.PathLike`): Path or URL of the video.
+
+ Returns:
+ [`MiniMaxH3VideoReference`]: the `(num_frames, height, width, 3)` `uint8` frames at the frame rate the
+ container reports, carrying its soundtrack and that soundtrack's own sample rate when it has an audio
+ stream.
+ """
+ frames, fps, audio, sample_rate = _decode_video_file(media)
+ return cls(frames=frames, fps=fps, audio=audio, sample_rate=sample_rate)
+
+
+@dataclass
+class MiniMaxH3AudioReference(MiniMaxH3Reference):
+ r"""
+ A voice or music reference: at most 3 per request, and never on its own — an audio reference has to be paired with
+ at least one image or video reference. It never reaches the conditioner and is encoded by the audio VAE alone.
+
+ Attributes:
+ audio (`torch.Tensor` of shape `(channels, num_samples)`):
+ The reference waveform, mono or stereo.
+ sample_rate (`int`, *optional*):
+ The rate `audio` carries its samples at. Left out, it is the audio VAE's own, which leaves the samples
+ untouched; any other rate is resampled onto it.
+ """
+
+ audio: torch.Tensor
+ sample_rate: int | None = None
+
+ kind = "audio"
+ has_audio = True
+
+ @classmethod
+ def from_file(cls, media) -> "MiniMaxH3AudioReference":
+ r"""
+ Decode an audio file into a [`MiniMaxH3AudioReference`], at the sample rate it carries.
+
+ Needs [PyAV](https://github.com/PyAV-Org/PyAV).
+
+ Args:
+ media (`str` or `os.PathLike`): Path or URL of the audio, or of a video whose soundtrack is taken.
+
+ Returns:
+ [`MiniMaxH3AudioReference`]: the `(channels, num_samples)` float32 waveform at the sample rate the
+ container reports.
+ """
+ audio, sample_rate = _decode_audio_file(media)
+ return cls(audio=audio, sample_rate=sample_rate)
+
+
+@contextlib.contextmanager
+def _local_media_file(media):
+ r"""The reference media as a local file: a URL is downloaded to a temporary file, removed on the way out."""
+ path = str(media)
+ if not path.startswith(("http://", "https://")):
+ if not os.path.isfile(path):
+ raise ValueError(
+ f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {path} is not a valid path."
+ )
+ yield path
+ return
+
+ response = requests.get(path, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT)
+ if response.status_code != 200:
+ raise ValueError(f"Failed to download {path}. Status code: {response.status_code}")
+ suffix = os.path.splitext(os.path.basename(unquote(urlparse(path).path)))[1]
+ download = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
+ try:
+ with download as file:
+ for chunk in response.iter_content(chunk_size=8192):
+ file.write(chunk)
+ yield download.name
+ finally:
+ os.remove(download.name)
+
+
+def _import_av():
+ r"""PyAV, the soft dependency a media file is decoded with."""
+ if not is_av_available():
+ raise ImportError(
+ "Decoding a MiniMax-H3 reference from a file needs PyAV. You can install it with `pip install av`, or "
+ "build the reference from decoded media itself: frames and the `fps` they carry for a video, a "
+ "`(channels, num_samples)` waveform and its `sample_rate` for audio."
+ )
+
+ import av
+
+ return av
+
+
+def _decode_reference_soundtrack(av, container, stream) -> tuple[torch.Tensor, int]:
+ r"""
+ An audio stream's samples as a `(channels, num_samples)` float32 waveform, at the rate the container carries them.
+
+ Args:
+ av (`module`): PyAV.
+ container (`av.container.InputContainer`): The open container.
+ stream (`av.audio.stream.AudioStream`): The stream to decode.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: the waveform and its sample rate.
+ """
+ sample_rate = int(stream.codec_context.sample_rate)
+ # Planar float is a format conversion only: the sample rate and the channel layout stay the container's own, and a
+ # mono soundtrack is upmixed later, by the setup step's audio normalization.
+ resampler = av.audio.resampler.AudioResampler(format="fltp", layout=stream.layout, rate=sample_rate)
+ chunks = []
+ for frame in container.decode(stream):
+ chunks += [torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(frame)]
+ # Whatever the resampler is still holding.
+ chunks += [torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(None)]
+ return torch.cat(chunks, dim=-1).to(torch.float32), sample_rate
+
+
+def _decode_video_file(media) -> tuple[np.ndarray, float, torch.Tensor | None, int | None]:
+ r"""
+ A video file's frames as `(num_frames, height, width, 3)` `uint8`, the frame rate the container reports, and its
+ soundtrack with that soundtrack's sample rate (`None, None` without an audio stream). The machinery behind
+ [`MiniMaxH3VideoReference.from_file`].
+ """
+ av = _import_av()
+ with _local_media_file(media) as path, av.open(path) as container:
+ stream = container.streams.video[0]
+ frames, rotation = [], 0.0
+ for frame in container.decode(stream):
+ # The display matrix rotation belongs to the stream, and PyAV surfaces it on every frame of it.
+ rotation = frame.rotation
+ frames.append(frame.to_ndarray(format="rgb24"))
+ frame_rate = float(stream.average_rate or stream.guessed_rate)
+ soundtrack = None
+ if container.streams.audio:
+ # Decoding the frames drained the container, so the soundtrack is read in a second pass over it.
+ container.seek(0)
+ soundtrack = _decode_reference_soundtrack(av, container, container.streams.audio[0])
+
+ if not frames:
+ raise ValueError(f"No video frames to decode in {media}.")
+ frames = np.stack(frames)
+ # `ffmpeg` displays a frame upright by undoing the counterclockwise rotation the display matrix carries, which is
+ # what this reproduces, snapped to the nearest quarter turn. A non-square pixel aspect ratio is left alone: the
+ # reference implementation resolved a reference's canvas from its *display* geometry, so a stream that carries a
+ # sample aspect ratio is conditioned on at the wrong shape, and correcting it is untested guesswork here.
+ turns = round(rotation / 90.0) % 4
+ if turns:
+ frames = np.ascontiguousarray(np.rot90(frames, k=-turns, axes=(1, 2)))
+ waveform, sample_rate = soundtrack if soundtrack is not None else (None, None)
+ return frames, frame_rate, waveform, sample_rate
+
+
+def _decode_audio_file(media) -> tuple[torch.Tensor, int]:
+ r"""
+ An audio file's `(channels, num_samples)` float32 waveform, at the sample rate the container reports. The
+ machinery behind [`MiniMaxH3AudioReference.from_file`].
+ """
+ av = _import_av()
+ with _local_media_file(media) as path, av.open(path) as container:
+ if not container.streams.audio:
+ raise ValueError(f"No audio stream to decode in {media}.")
+ waveform, sample_rate = _decode_reference_soundtrack(av, container, container.streams.audio[0])
+ return waveform, sample_rate
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..f59a7e132b102334d643fcc715941beb8a4941bd
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline.py
@@ -0,0 +1,2936 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import gc
+import importlib
+import inspect
+import os
+import sys
+import traceback
+import warnings
+from collections import OrderedDict
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from huggingface_hub import create_repo
+from huggingface_hub.utils import validate_hf_hub_args
+from tqdm.auto import tqdm
+from typing_extensions import Self
+
+from ..configuration_utils import ConfigMixin, FrozenDict
+from ..pipelines.pipeline_loading_utils import (
+ LOADABLE_CLASSES,
+ _fetch_class_library_tuple,
+ _unwrap_model,
+ simple_get_class_obj,
+)
+from ..utils import PushToHubMixin, is_accelerate_available, logging
+from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ..utils.hub_utils import load_or_create_model_card, populate_model_card
+from ..utils.torch_utils import empty_device_cache, is_compiled_module
+from .components_manager import ComponentsManager
+from .modular_pipeline_utils import (
+ MODULAR_MODEL_CARD_TEMPLATE,
+ ComponentSpec,
+ ConfigSpec,
+ InputParam,
+ InsertableDict,
+ OutputParam,
+ _validate_requirements,
+ combine_inputs,
+ combine_outputs,
+ format_components,
+ format_configs,
+ format_workflow,
+ generate_modular_model_card_content,
+ make_doc_string,
+)
+
+
+if is_accelerate_available():
+ import accelerate
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# map regular pipeline to modular pipeline class name
+
+
+def _create_default_map_fn(pipeline_class_name: str):
+ """Create a mapping function that always returns the same pipeline class."""
+
+ def _map_fn(config_dict=None):
+ return pipeline_class_name
+
+ return _map_fn
+
+
+def _flux2_klein_map_fn(config_dict=None):
+ if config_dict is None:
+ return "Flux2KleinModularPipeline"
+
+ if "is_distilled" in config_dict and config_dict["is_distilled"]:
+ return "Flux2KleinModularPipeline"
+ else:
+ return "Flux2KleinBaseModularPipeline"
+
+
+def _wan_map_fn(config_dict=None):
+ if config_dict is None:
+ return "WanModularPipeline"
+
+ if "boundary_ratio" in config_dict and config_dict["boundary_ratio"] is not None:
+ return "Wan22ModularPipeline"
+ else:
+ return "WanModularPipeline"
+
+
+def _wan_i2v_map_fn(config_dict=None):
+ if config_dict is None:
+ return "WanImage2VideoModularPipeline"
+
+ if "boundary_ratio" in config_dict and config_dict["boundary_ratio"] is not None:
+ return "Wan22Image2VideoModularPipeline"
+ else:
+ return "WanImage2VideoModularPipeline"
+
+
+def _krea2_map_fn(config_dict=None):
+ if config_dict is None:
+ return "Krea2ModularPipeline"
+
+ if config_dict.get("is_distilled", False):
+ return "Krea2TurboModularPipeline"
+ else:
+ return "Krea2ModularPipeline"
+
+
+def _helios_pyramid_map_fn(config_dict=None):
+ if config_dict is None:
+ return "HeliosPyramidModularPipeline"
+
+ if config_dict.get("is_distilled", False):
+ return "HeliosPyramidDistilledModularPipeline"
+ else:
+ return "HeliosPyramidModularPipeline"
+
+
+MODULAR_PIPELINE_MAPPING = OrderedDict(
+ [
+ ("stable-diffusion-xl", _create_default_map_fn("StableDiffusionXLModularPipeline")),
+ ("stable-diffusion-3", _create_default_map_fn("StableDiffusion3ModularPipeline")),
+ ("wan", _wan_map_fn),
+ ("wan-i2v", _wan_i2v_map_fn),
+ ("flux", _create_default_map_fn("FluxModularPipeline")),
+ ("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")),
+ ("flux2", _create_default_map_fn("Flux2ModularPipeline")),
+ ("flux2-klein", _flux2_klein_map_fn),
+ ("ideogram4", _create_default_map_fn("Ideogram4ModularPipeline")),
+ ("krea2", _krea2_map_fn),
+ ("qwenimage", _create_default_map_fn("QwenImageModularPipeline")),
+ ("qwenimage-edit", _create_default_map_fn("QwenImageEditModularPipeline")),
+ ("qwenimage-edit-plus", _create_default_map_fn("QwenImageEditPlusModularPipeline")),
+ ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")),
+ ("anima", _create_default_map_fn("AnimaModularPipeline")),
+ ("z-image", _create_default_map_fn("ZImageModularPipeline")),
+ ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")),
+ ("helios", _create_default_map_fn("HeliosModularPipeline")),
+ ("helios-pyramid", _helios_pyramid_map_fn),
+ ("hunyuan-video-1.5", _create_default_map_fn("HunyuanVideo15ModularPipeline")),
+ ("ltx", _create_default_map_fn("LTXModularPipeline")),
+ ("minimax-h3", _create_default_map_fn("MiniMaxH3ModularPipeline")),
+ ("ernie-image", _create_default_map_fn("ErnieImageModularPipeline")),
+ ]
+)
+
+
+@dataclass
+class PipelineState:
+ """
+ [`PipelineState`] stores the state of a pipeline. It is used to pass data between pipeline blocks.
+ """
+
+ values: dict[str, Any] = field(default_factory=dict)
+ kwargs_mapping: dict[str, list[str]] = field(default_factory=dict)
+
+ def set(self, key: str, value: Any, kwargs_type: str = None):
+ """
+ Add a value to the pipeline state.
+
+ Args:
+ key (str): The key for the value
+ value (Any): The value to store
+ kwargs_type (str): The kwargs_type with which the value is associated
+ """
+ self.values[key] = value
+
+ if kwargs_type is not None:
+ if kwargs_type not in self.kwargs_mapping:
+ self.kwargs_mapping[kwargs_type] = [key]
+ else:
+ self.kwargs_mapping[kwargs_type].append(key)
+
+ def get(self, keys: str | list[str], default: Any = None) -> Any | dict[str, Any]:
+ """
+ Get one or multiple values from the pipeline state.
+
+ Args:
+ keys (str | list[str]): Key or list of keys for the values
+ default (Any): The default value to return if not found
+
+ Returns:
+ Any | dict[str, Any]: Single value if keys is str, dictionary of values if keys is list
+ """
+ if isinstance(keys, str):
+ return self.values.get(keys, default)
+ return {key: self.values.get(key, default) for key in keys}
+
+ def get_by_kwargs(self, kwargs_type: str) -> dict[str, Any]:
+ """
+ Get all values with matching kwargs_type.
+
+ Args:
+ kwargs_type (str): The kwargs_type to filter by
+
+ Returns:
+ dict[str, Any]: Dictionary of values with matching kwargs_type
+ """
+ value_names = self.kwargs_mapping.get(kwargs_type, [])
+ return self.get(value_names)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Convert PipelineState to a dictionary.
+ """
+ return {**self.__dict__}
+
+ def __getattr__(self, name):
+ """
+ Allow attribute access to intermediate values. If an attribute is not found in the object, look for it in the
+ intermediates dict.
+ """
+ # Use object.__getattribute__ to avoid infinite recursion during deepcopy
+ try:
+ values = object.__getattribute__(self, "values")
+ except AttributeError:
+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+ if name in values:
+ return values[name]
+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+ def __repr__(self):
+ def format_value(v):
+ if hasattr(v, "shape") and hasattr(v, "dtype"):
+ return f"Tensor(dtype={v.dtype}, shape={v.shape})"
+ elif isinstance(v, list) and len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
+ return f"[Tensor(dtype={v[0].dtype}, shape={v[0].shape}), ...]"
+ else:
+ return repr(v)
+
+ values_str = "\n".join(f" {k}: {format_value(v)}" for k, v in self.values.items())
+ kwargs_mapping_str = "\n".join(f" {k}: {v}" for k, v in self.kwargs_mapping.items())
+
+ return f"PipelineState(\n values={{\n{values_str}\n }},\n kwargs_mapping={{\n{kwargs_mapping_str}\n }}\n)"
+
+
+@dataclass
+class BlockState:
+ """
+ Container for block state data with attribute access and formatted representation.
+ """
+
+ def __init__(self, **kwargs):
+ for key, value in kwargs.items():
+ setattr(self, key, value)
+
+ def __getitem__(self, key: str):
+ # allows block_state["foo"]
+ return getattr(self, key, None)
+
+ def __setitem__(self, key: str, value: Any):
+ # allows block_state["foo"] = "bar"
+ setattr(self, key, value)
+
+ def as_dict(self):
+ """
+ Convert BlockState to a dictionary.
+
+ Returns:
+ dict[str, Any]: Dictionary containing all attributes of the BlockState
+ """
+ return dict(self.__dict__.items())
+
+ def __repr__(self):
+ def format_value(v):
+ # Handle tensors directly
+ if hasattr(v, "shape") and hasattr(v, "dtype"):
+ return f"Tensor(dtype={v.dtype}, shape={v.shape})"
+
+ # Handle lists of tensors
+ elif isinstance(v, list):
+ if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
+ shapes = [t.shape for t in v]
+ return f"list[{len(v)}] of Tensors with shapes {shapes}"
+ return repr(v)
+
+ # Handle tuples of tensors
+ elif isinstance(v, tuple):
+ if len(v) > 0 and hasattr(v[0], "shape") and hasattr(v[0], "dtype"):
+ shapes = [t.shape for t in v]
+ return f"tuple[{len(v)}] of Tensors with shapes {shapes}"
+ return repr(v)
+
+ # Handle dicts with tensor values
+ elif isinstance(v, dict):
+ formatted_dict = {}
+ for k, val in v.items():
+ if hasattr(val, "shape") and hasattr(val, "dtype"):
+ formatted_dict[k] = f"Tensor(shape={val.shape}, dtype={val.dtype})"
+ elif (
+ isinstance(val, list)
+ and len(val) > 0
+ and hasattr(val[0], "shape")
+ and hasattr(val[0], "dtype")
+ ):
+ shapes = [t.shape for t in val]
+ formatted_dict[k] = f"list[{len(val)}] of Tensors with shapes {shapes}"
+ else:
+ formatted_dict[k] = repr(val)
+ return formatted_dict
+
+ # Default case
+ return repr(v)
+
+ attributes = "\n".join(f" {k}: {format_value(v)}" for k, v in self.__dict__.items())
+ return f"BlockState(\n{attributes}\n)"
+
+
+class ModularPipelineBlocks(ConfigMixin, PushToHubMixin):
+ """
+ Base class for all Pipeline Blocks: ConditionalPipelineBlocks, AutoPipelineBlocks, SequentialPipelineBlocks,
+ LoopSequentialPipelineBlocks
+
+ [`ModularPipelineBlocks`] provides method to load and save the definition of pipeline blocks.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ config_name = "modular_config.json"
+ model_name = None
+ _requirements: dict[str, str] | None = None
+ _workflow_map = None
+
+ @classmethod
+ def _get_signature_keys(cls, obj):
+ parameters = inspect.signature(obj.__init__).parameters
+ required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
+ optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
+ expected_modules = set(required_parameters.keys()) - {"self"}
+
+ return expected_modules, optional_parameters
+
+ def __init__(self):
+ self.sub_blocks = InsertableDict()
+
+ @property
+ def description(self) -> str:
+ """Description of the block. Must be implemented by subclasses."""
+ return ""
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return []
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ """list of input parameters. Must be implemented by subclasses."""
+ return []
+
+ def _get_required_inputs(self):
+ input_names = []
+ for input_param in self.inputs:
+ if input_param.required:
+ input_names.append(input_param.name)
+
+ return input_names
+
+ @property
+ def required_inputs(self) -> list[InputParam]:
+ return self._get_required_inputs()
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ """list of intermediate output parameters. Must be implemented by subclasses."""
+ return []
+
+ def _get_outputs(self):
+ return self.intermediate_outputs
+
+ @property
+ def outputs(self) -> list[OutputParam]:
+ return self._get_outputs()
+
+ # currentlyonly ConditionalPipelineBlocks and SequentialPipelineBlocks support `get_execution_blocks`
+ def get_execution_blocks(self, **kwargs):
+ """
+ Get the block(s) that would execute given the inputs. Must be implemented by subclasses that support
+ conditional block selection.
+
+ Args:
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
+ """
+ raise NotImplementedError(f"`get_execution_blocks` is not implemented for {self.__class__.__name__}")
+
+ # currently only SequentialPipelineBlocks support workflows
+ @property
+ def available_workflows(self):
+ """
+ Returns a list of available workflow names. Must be implemented by subclasses that define `_workflow_map`.
+ """
+ raise NotImplementedError(f"`available_workflows` is not implemented for {self.__class__.__name__}")
+
+ def get_workflow(self, workflow_name: str):
+ """
+ Get the execution blocks for a specific workflow. Must be implemented by subclasses that define
+ `_workflow_map`.
+
+ Args:
+ workflow_name: Name of the workflow to retrieve.
+ """
+ raise NotImplementedError(f"`get_workflow` is not implemented for {self.__class__.__name__}")
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ pretrained_model_name_or_path: str,
+ trust_remote_code: bool = False,
+ **kwargs,
+ ):
+ hub_kwargs_names = [
+ "cache_dir",
+ "force_download",
+ "local_files_only",
+ "local_dir",
+ "proxies",
+ "revision",
+ "subfolder",
+ "token",
+ ]
+ hub_kwargs = {name: kwargs.pop(name) for name in hub_kwargs_names if name in kwargs}
+
+ config = cls.load_config(pretrained_model_name_or_path, **hub_kwargs)
+ has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"]
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_remote_code
+ )
+ if not has_remote_code and trust_remote_code:
+ raise ValueError(
+ "Selected model repository does not happear to have any custom code or does not have a valid `config.json` file."
+ )
+
+ if "requirements" in config and config["requirements"] is not None:
+ _ = _validate_requirements(config["requirements"])
+
+ class_ref = config["auto_map"][cls.__name__]
+ module_file, class_name = class_ref.split(".")
+ module_file = module_file + ".py"
+ block_cls = get_class_from_dynamic_module(
+ pretrained_model_name_or_path,
+ module_file=module_file,
+ class_name=class_name,
+ trust_remote_code=trust_remote_code,
+ **hub_kwargs,
+ )
+ expected_kwargs, optional_kwargs = block_cls._get_signature_keys(block_cls)
+ block_kwargs = {
+ name: kwargs.get(name) for name in kwargs if name in expected_kwargs or name in optional_kwargs
+ }
+
+ return block_cls(**block_kwargs)
+
+ def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
+ # TODO: factor out this logic.
+ cls_name = self.__class__.__name__
+
+ full_mod = type(self).__module__
+ module = full_mod.rsplit(".", 1)[-1].replace("__dynamic__", "")
+ parent_module = self.save_pretrained.__func__.__qualname__.split(".", 1)[0]
+ auto_map = {f"{parent_module}": f"{module}.{cls_name}"}
+ self.register_to_config(auto_map=auto_map)
+
+ # resolve requirements
+ requirements = _validate_requirements(getattr(self, "_requirements", None))
+ if requirements:
+ self.register_to_config(requirements=requirements)
+
+ self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
+ config = dict(self.config)
+ self._internal_dict = FrozenDict(config)
+
+ def init_pipeline(
+ self,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
+ ) -> "ModularPipeline":
+ """
+ create a ModularPipeline, optionally accept pretrained_model_name_or_path to load from hub.
+ """
+ map_fn = MODULAR_PIPELINE_MAPPING.get(self.model_name, _create_default_map_fn("ModularPipeline"))
+ pipeline_class_name = map_fn()
+ diffusers_module = importlib.import_module("diffusers")
+ pipeline_class = getattr(diffusers_module, pipeline_class_name)
+
+ modular_pipeline = pipeline_class(
+ blocks=deepcopy(self),
+ pretrained_model_name_or_path=pretrained_model_name_or_path,
+ components_manager=components_manager,
+ collection=collection,
+ )
+ return modular_pipeline
+
+ def get_block_state(self, state: PipelineState) -> dict:
+ """Get all inputs and intermediates in one dictionary"""
+ data = {}
+ state_inputs = self.inputs
+
+ # Check inputs
+ for input_param in state_inputs:
+ if input_param.name:
+ value = state.get(input_param.name)
+ if input_param.required and value is None:
+ raise ValueError(f"Required input '{input_param.name}' is missing")
+ elif value is not None or (value is None and input_param.name not in data):
+ data[input_param.name] = value
+
+ elif input_param.kwargs_type:
+ # if kwargs_type is provided, get all inputs with matching kwargs_type
+ if input_param.kwargs_type not in data:
+ data[input_param.kwargs_type] = {}
+ inputs_kwargs = state.get_by_kwargs(input_param.kwargs_type)
+ if inputs_kwargs:
+ for k, v in inputs_kwargs.items():
+ if v is not None:
+ data[k] = v
+ data[input_param.kwargs_type][k] = v
+
+ return BlockState(**data)
+
+ def set_block_state(self, state: PipelineState, block_state: BlockState):
+ for output_param in self.intermediate_outputs:
+ if not hasattr(block_state, output_param.name):
+ raise ValueError(f"Intermediate output '{output_param.name}' is missing in block state")
+ param = getattr(block_state, output_param.name)
+ state.set(output_param.name, param, output_param.kwargs_type)
+
+ for input_param in self.inputs:
+ if input_param.name and hasattr(block_state, input_param.name):
+ param = getattr(block_state, input_param.name)
+ # Only add if the value is different from what's in the state
+ current_value = state.get(input_param.name)
+ if current_value is not param: # Using identity comparison to check if object was modified
+ state.set(input_param.name, param, input_param.kwargs_type)
+
+ elif input_param.kwargs_type:
+ # if it is a kwargs type, e.g. "denoiser_input_fields", it is likely to be a list of parameters
+ # we need to first find out which inputs are and loop through them.
+ intermediate_kwargs = state.get_by_kwargs(input_param.kwargs_type)
+ for param_name, current_value in intermediate_kwargs.items():
+ if param_name is None:
+ continue
+
+ if not hasattr(block_state, param_name):
+ continue
+
+ param = getattr(block_state, param_name)
+ if current_value is not param: # Using identity comparison to check if object was modified
+ state.set(param_name, param, input_param.kwargs_type)
+
+ @property
+ def input_names(self) -> list[str]:
+ return [input_param.name for input_param in self.inputs if input_param.name is not None]
+
+ @property
+ def intermediate_output_names(self) -> list[str]:
+ return [output_param.name for output_param in self.intermediate_outputs if output_param.name is not None]
+
+ @property
+ def output_names(self) -> list[str]:
+ return [output_param.name for output_param in self.outputs if output_param.name is not None]
+
+ @property
+ def component_names(self) -> list[str]:
+ return [component.name for component in self.expected_components]
+
+ @property
+ def doc(self):
+ return make_doc_string(
+ self.inputs,
+ self.outputs,
+ self.description,
+ class_name=self.__class__.__name__,
+ expected_components=self.expected_components,
+ expected_configs=self.expected_configs,
+ )
+
+
+class ConditionalPipelineBlocks(ModularPipelineBlocks):
+ """
+ A Pipeline Blocks that conditionally selects a block to run based on the inputs. Subclasses must implement the
+ `select_block` method to define the logic for selecting the block. Currently, we only support selection logic based
+ on the presence or absence of inputs (i.e., whether they are `None` or not)
+
+ This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipeline blocks (such as loading or saving etc.)
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+
+ Attributes:
+ block_classes: List of block classes to be used. Must have the same length as `block_names`.
+ block_names: List of names for each block. Must have the same length as `block_classes`.
+ block_trigger_inputs: List of input names that `select_block()` uses to determine which block to run.
+ For `ConditionalPipelineBlocks`, this does not need to correspond to `block_names` and `block_classes`. For
+ `AutoPipelineBlocks`, this must have the same length as `block_names` and `block_classes`, where each
+ element specifies the trigger input for the corresponding block.
+ default_block_name: Name of the default block to run when no trigger inputs match.
+ If None, this block can be skipped entirely when no trigger inputs are provided.
+ """
+
+ block_classes = []
+ block_names = []
+ block_trigger_inputs = []
+ default_block_name = None
+
+ def __init__(self):
+ sub_blocks = InsertableDict()
+ for block_name, block in zip(self.block_names, self.block_classes):
+ if inspect.isclass(block):
+ sub_blocks[block_name] = block()
+ else:
+ sub_blocks[block_name] = block
+ self.sub_blocks = sub_blocks
+ if not (len(self.block_classes) == len(self.block_names)):
+ raise ValueError(
+ f"In {self.__class__.__name__}, the number of block_classes and block_names must be the same."
+ )
+ if self.default_block_name is not None and self.default_block_name not in self.block_names:
+ raise ValueError(
+ f"In {self.__class__.__name__}, default_block_name '{self.default_block_name}' must be one of block_names: {self.block_names}"
+ )
+
+ @property
+ def model_name(self):
+ return next(iter(self.sub_blocks.values())).model_name
+
+ @property
+ def description(self):
+ return ""
+
+ @property
+ def expected_components(self):
+ expected_components = []
+ for block in self.sub_blocks.values():
+ for component in block.expected_components:
+ if component not in expected_components:
+ expected_components.append(component)
+ return expected_components
+
+ @property
+ def expected_configs(self):
+ expected_configs = []
+ for block in self.sub_blocks.values():
+ for config in block.expected_configs:
+ if config not in expected_configs:
+ expected_configs.append(config)
+ return expected_configs
+
+ @property
+ def required_inputs(self) -> list[str]:
+ # no default block means this conditional block can be skipped entirely
+ if self.default_block_name is None:
+ return []
+
+ first_block = next(iter(self.sub_blocks.values()))
+ required_by_all = set(getattr(first_block, "required_inputs", set()))
+
+ # Intersect with required inputs from all other blocks
+ for block in list(self.sub_blocks.values())[1:]:
+ block_required = set(getattr(block, "required_inputs", set()))
+ required_by_all.intersection_update(block_required)
+
+ return list(required_by_all)
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ named_inputs = [(name, block.inputs) for name, block in self.sub_blocks.items()]
+ combined_inputs = combine_inputs(*named_inputs)
+ # mark Required inputs only if that input is required by all the blocks
+ for input_param in combined_inputs:
+ if input_param.name in self.required_inputs:
+ input_param.required = True
+ else:
+ input_param.required = False
+ return combined_inputs
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()]
+ combined_outputs = combine_outputs(*named_outputs)
+ return combined_outputs
+
+ @property
+ def outputs(self) -> list[str]:
+ named_outputs = [(name, block.outputs) for name, block in self.sub_blocks.items()]
+ combined_outputs = combine_outputs(*named_outputs)
+ return combined_outputs
+
+ @property
+ # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements
+ def _requirements(self) -> dict[str, str]:
+ requirements = {}
+ for block_name, block in self.sub_blocks.items():
+ if getattr(block, "_requirements", None):
+ requirements[block_name] = block._requirements
+ return requirements
+
+ # used for `__repr__`
+ def _get_trigger_inputs(self) -> set:
+ """
+ Returns a set of all unique trigger input values found in this block and nested blocks.
+ """
+
+ def fn_recursive_get_trigger(blocks):
+ trigger_values = set()
+
+ if blocks is not None:
+ for name, block in blocks.items():
+ # Check if current block has block_trigger_inputs
+ if hasattr(block, "block_trigger_inputs") and block.block_trigger_inputs is not None:
+ trigger_values.update(t for t in block.block_trigger_inputs if t is not None)
+
+ # If block has sub_blocks, recursively check them
+ if block.sub_blocks:
+ nested_triggers = fn_recursive_get_trigger(block.sub_blocks)
+ trigger_values.update(nested_triggers)
+
+ return trigger_values
+
+ # Start with this block's block_trigger_inputs
+ all_triggers = {t for t in self.block_trigger_inputs if t is not None}
+ # Add nested triggers
+ all_triggers.update(fn_recursive_get_trigger(self.sub_blocks))
+
+ return all_triggers
+
+ def select_block(self, **kwargs) -> str | None:
+ """
+ Select the block to run based on the trigger inputs. Subclasses must implement this method to define the logic
+ for selecting the block.
+
+ Note: When trigger inputs include intermediate outputs from earlier blocks, the selection logic should only
+ depend on the presence or absence of the input (i.e., whether it is None or not), not on its actual value. This
+ is because `get_execution_blocks()` resolves conditions statically by propagating intermediate output names
+ without their runtime values.
+
+ Args:
+ **kwargs: Trigger input names and their values from the state.
+
+ Returns:
+ str | None: The name of the block to run, or None to use default/skip.
+ """
+ raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement the `select_block` method.")
+
+ @torch.no_grad()
+ def __call__(self, pipeline, state: PipelineState) -> PipelineState:
+ trigger_kwargs = {name: state.get(name) for name in self.block_trigger_inputs if name is not None}
+ block_name = self.select_block(**trigger_kwargs)
+
+ if block_name is None:
+ block_name = self.default_block_name
+
+ if block_name is None:
+ logger.info(f"skipping conditional block: {self.__class__.__name__}")
+ return pipeline, state
+
+ block = self.sub_blocks[block_name]
+
+ try:
+ logger.info(f"Running block: {block.__class__.__name__}")
+ return block(pipeline, state)
+ except Exception as e:
+ error_msg = (
+ f"\nError in block: {block.__class__.__name__}\n"
+ f"Error details: {str(e)}\n"
+ f"Traceback:\n{traceback.format_exc()}"
+ )
+ logger.error(error_msg)
+ raise
+
+ def get_execution_blocks(self, **kwargs) -> ModularPipelineBlocks | None:
+ """
+ Get the block(s) that would execute given the inputs.
+
+ Recursively resolves nested ConditionalPipelineBlocks until reaching either:
+ - A leaf block (no sub_blocks or LoopSequentialPipelineBlocks) → returns single `ModularPipelineBlocks`
+ - A `SequentialPipelineBlocks` → delegates to its `get_execution_blocks()` which returns
+ a `SequentialPipelineBlocks` containing the resolved execution blocks
+
+ Args:
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
+
+ Returns:
+ - `ModularPipelineBlocks`: A leaf block or resolved `SequentialPipelineBlocks`
+ - `None`: If this block would be skipped (no trigger matched and no default)
+ """
+ trigger_kwargs = {name: kwargs.get(name) for name in self.block_trigger_inputs if name is not None}
+ block_name = self.select_block(**trigger_kwargs)
+
+ if block_name is None:
+ block_name = self.default_block_name
+
+ if block_name is None:
+ return None
+
+ block = self.sub_blocks[block_name]
+
+ # Recursively resolve until we hit a leaf block
+ if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks):
+ return block.get_execution_blocks(**kwargs)
+
+ return block
+
+ def __repr__(self):
+ class_name = self.__class__.__name__
+ base_class = self.__class__.__bases__[0].__name__
+ header = (
+ f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n"
+ )
+
+ if self._get_trigger_inputs():
+ header += "\n"
+ header += " " + "=" * 100 + "\n"
+ header += " This pipeline contains blocks that are selected at runtime based on inputs.\n"
+ header += f" Trigger Inputs: {sorted(self._get_trigger_inputs())}\n"
+ header += " " + "=" * 100 + "\n\n"
+
+ # Format description with proper indentation
+ desc_lines = self.description.split("\n")
+ desc = []
+ # First line with "Description:" label
+ desc.append(f" Description: {desc_lines[0]}")
+ # Subsequent lines with proper indentation
+ if len(desc_lines) > 1:
+ desc.extend(f" {line}" for line in desc_lines[1:])
+ desc = "\n".join(desc) + "\n"
+
+ # Components section - focus only on expected components
+ expected_components = getattr(self, "expected_components", [])
+ components_str = format_components(expected_components, indent_level=2, add_empty_lines=False)
+
+ # Configs section - use format_configs with add_empty_lines=False
+ expected_configs = getattr(self, "expected_configs", [])
+ configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False)
+
+ # Blocks section
+ blocks_str = " Sub-Blocks:\n"
+ for i, (name, block) in enumerate(self.sub_blocks.items()):
+ if name == self.default_block_name:
+ addtional_str = " [default]"
+ else:
+ addtional_str = ""
+ blocks_str += f" • {name}{addtional_str} ({block.__class__.__name__})\n"
+
+ # Add block description
+ block_desc_lines = block.description.split("\n")
+ indented_desc = block_desc_lines[0]
+ if len(block_desc_lines) > 1:
+ indented_desc += "\n" + "\n".join(" " + line for line in block_desc_lines[1:])
+ blocks_str += f" Description: {indented_desc}\n\n"
+
+ # Build the representation with conditional sections
+ result = f"{header}\n{desc}"
+
+ # Only add components section if it has content
+ if components_str.strip():
+ result += f"\n\n{components_str}"
+
+ # Only add configs section if it has content
+ if configs_str.strip():
+ result += f"\n\n{configs_str}"
+
+ # Always add blocks section
+ result += f"\n\n{blocks_str})"
+
+ return result
+
+ @property
+ def doc(self):
+ return make_doc_string(
+ self.inputs,
+ self.outputs,
+ self.description,
+ class_name=self.__class__.__name__,
+ expected_components=self.expected_components,
+ expected_configs=self.expected_configs,
+ )
+
+
+class AutoPipelineBlocks(ConditionalPipelineBlocks):
+ """
+ A Pipeline Blocks that automatically selects a block to run based on the presence of trigger inputs.
+
+ This is a specialized version of `ConditionalPipelineBlocks` where:
+ - Each block has one corresponding trigger input (1:1 mapping)
+ - Block selection is automatic: the first block whose trigger input is present gets selected
+ - `block_trigger_inputs` must have the same length as `block_names` and `block_classes`
+ - Use `None` in `block_trigger_inputs` to specify the default block, i.e the block that will run if no trigger
+ inputs are present
+
+ Attributes:
+ block_classes:
+ List of block classes to be used. Must have the same length as `block_names` and
+ `block_trigger_inputs`.
+ block_names:
+ List of names for each block. Must have the same length as `block_classes` and `block_trigger_inputs`.
+ block_trigger_inputs:
+ List of input names where each element specifies the trigger input for the corresponding block. Use
+ `None` to mark the default block.
+
+ Example:
+ ```python
+ class MyAutoBlock(AutoPipelineBlocks):
+ block_classes = [InpaintEncoderBlock, ImageEncoderBlock, TextEncoderBlock]
+ block_names = ["inpaint", "img2img", "text2img"]
+ block_trigger_inputs = ["mask_image", "image", None] # text2img is the default
+ ```
+
+ With this definition:
+ - As long as `mask_image` is provided, "inpaint" block runs (regardless of `image` being provided or not)
+ - If `mask_image` is not provided but `image` is provided, "img2img" block runs
+ - Otherwise, "text2img" block runs (default, trigger is `None`)
+ """
+
+ def __init__(self):
+ super().__init__()
+
+ if self.default_block_name is not None:
+ raise ValueError(
+ f"In {self.__class__.__name__}, do not set `default_block_name` for AutoPipelineBlocks. "
+ f"Use `None` in `block_trigger_inputs` to specify the default block."
+ )
+
+ if not (len(self.block_classes) == len(self.block_names) == len(self.block_trigger_inputs)):
+ raise ValueError(
+ f"In {self.__class__.__name__}, the number of block_classes, block_names, and block_trigger_inputs must be the same."
+ )
+
+ if None in self.block_trigger_inputs:
+ idx = self.block_trigger_inputs.index(None)
+ self.default_block_name = self.block_names[idx]
+
+ def select_block(self, **kwargs) -> str | None:
+ """Select block based on which trigger input is present (not None)."""
+ for trigger_input, block_name in zip(self.block_trigger_inputs, self.block_names):
+ if trigger_input is not None and kwargs.get(trigger_input) is not None:
+ return block_name
+ return None
+
+
+class SequentialPipelineBlocks(ModularPipelineBlocks):
+ """
+ A Pipeline Blocks that combines multiple pipeline block classes into one. When called, it will call each block in
+ sequence.
+
+ This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipeline blocks (such as loading or saving etc.)
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+
+ Attributes:
+ block_classes: list of block classes to be used
+ block_names: list of prefixes for each block
+ """
+
+ block_classes = []
+ block_names = []
+
+ @property
+ def description(self):
+ return ""
+
+ @property
+ def model_name(self):
+ return next((block.model_name for block in self.sub_blocks.values() if block.model_name is not None), None)
+
+ @property
+ def expected_components(self):
+ expected_components = []
+ for block in self.sub_blocks.values():
+ for component in block.expected_components:
+ if component not in expected_components:
+ expected_components.append(component)
+ return expected_components
+
+ @property
+ def expected_configs(self):
+ expected_configs = []
+ for block in self.sub_blocks.values():
+ for config in block.expected_configs:
+ if config not in expected_configs:
+ expected_configs.append(config)
+ return expected_configs
+
+ @property
+ def available_workflows(self):
+ if self._workflow_map is None:
+ raise NotImplementedError(
+ f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+ )
+
+ return list(self._workflow_map.keys())
+
+ def get_workflow(self, workflow_name: str):
+ if self._workflow_map is None:
+ raise NotImplementedError(
+ f"workflows is not supported because _workflow_map is not set for {self.__class__.__name__}"
+ )
+
+ if workflow_name not in self._workflow_map:
+ raise ValueError(
+ f"Workflow {workflow_name!r} not found in {self.__class__.__name__}. "
+ f"Available workflows: {list(self._workflow_map)}"
+ )
+
+ trigger_inputs = self._workflow_map[workflow_name]
+ workflow_blocks = self.get_execution_blocks(**trigger_inputs)
+
+ return workflow_blocks
+
+ @classmethod
+ def from_blocks_dict(
+ cls, blocks_dict: dict[str, Any], description: str | None = None
+ ) -> "SequentialPipelineBlocks":
+ """Creates a SequentialPipelineBlocks instance from a dictionary of blocks.
+
+ Args:
+ blocks_dict: Dictionary mapping block names to block classes or instances
+
+ Returns:
+ A new SequentialPipelineBlocks instance
+ """
+ instance = cls()
+
+ # Create instances if classes are provided
+ sub_blocks = InsertableDict()
+ for name, block in blocks_dict.items():
+ if inspect.isclass(block):
+ sub_blocks[name] = block()
+ else:
+ sub_blocks[name] = block
+
+ instance.block_classes = [block.__class__ for block in sub_blocks.values()]
+ instance.block_names = list(sub_blocks.keys())
+ instance.sub_blocks = sub_blocks
+
+ if description is not None:
+ instance.description = description
+
+ return instance
+
+ def __init__(self):
+ sub_blocks = InsertableDict()
+ for block_name, block in zip(self.block_names, self.block_classes):
+ if inspect.isclass(block):
+ sub_blocks[block_name] = block()
+ else:
+ sub_blocks[block_name] = block
+ self.sub_blocks = sub_blocks
+ if not len(self.block_names) == len(self.block_classes):
+ raise ValueError(
+ f"In {self.__class__.__name__}, the number of block_names and block_classes must be the same."
+ )
+
+ def _get_inputs(self):
+ inputs = []
+ outputs = set()
+
+ # Go through all blocks in order
+ for block in self.sub_blocks.values():
+ # Add inputs that aren't in outputs yet
+ for inp in block.inputs:
+ if inp.name not in outputs and inp.name not in {input.name for input in inputs}:
+ inputs.append(inp)
+
+ # Only add outputs if the block cannot be skipped
+ should_add_outputs = True
+ if isinstance(block, ConditionalPipelineBlocks) and block.default_block_name is None:
+ # ConditionalPipelineBlocks without default can be skipped
+ should_add_outputs = False
+
+ if should_add_outputs:
+ # Add this block's outputs
+ block_intermediate_outputs = [out.name for out in block.intermediate_outputs]
+ outputs.update(block_intermediate_outputs)
+
+ return inputs
+
+ # YiYi TODO: add test for this
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return self._get_inputs()
+
+ @property
+ def required_inputs(self) -> list[str]:
+ # Get the first block from the dictionary
+ first_block = next(iter(self.sub_blocks.values()))
+ required_by_any = set(getattr(first_block, "required_inputs", set()))
+
+ # Union with required inputs from all other blocks
+ for block in list(self.sub_blocks.values())[1:]:
+ block_required = set(getattr(block, "required_inputs", set()))
+ required_by_any.update(block_required)
+
+ return list(required_by_any)
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ named_outputs = []
+ for name, block in self.sub_blocks.items():
+ inp_names = {inp.name for inp in block.inputs}
+ # so we only need to list new variables as intermediate_outputs, but if user wants to list these they modified it's still fine (a.k.a we don't enforce)
+ # filter out them here so they do not end up as intermediate_outputs
+ if name not in inp_names:
+ named_outputs.append((name, block.intermediate_outputs))
+ combined_outputs = combine_outputs(*named_outputs)
+ return combined_outputs
+
+ # YiYi TODO: I think we can remove the outputs property
+ @property
+ def outputs(self) -> list[str]:
+ # return next(reversed(self.sub_blocks.values())).intermediate_outputs
+ return self.intermediate_outputs
+
+ @torch.no_grad()
+ def __call__(self, pipeline, state: PipelineState) -> PipelineState:
+ for block_name, block in self.sub_blocks.items():
+ try:
+ pipeline, state = block(pipeline, state)
+ except Exception as e:
+ error_msg = (
+ f"\nError in block: ({block_name}, {block.__class__.__name__})\n"
+ f"Error details: {str(e)}\n"
+ f"Traceback:\n{traceback.format_exc()}"
+ )
+ logger.error(error_msg)
+ raise
+ return pipeline, state
+
+ # used for `__repr__`
+ def _get_trigger_inputs(self):
+ """
+ Returns a set of all unique trigger input values found in the blocks.
+ """
+
+ def fn_recursive_get_trigger(blocks):
+ trigger_values = set()
+
+ if blocks is not None:
+ for name, block in blocks.items():
+ # Check if current block has block_trigger_inputs (ConditionalPipelineBlocks)
+ if hasattr(block, "block_trigger_inputs") and block.block_trigger_inputs is not None:
+ trigger_values.update(t for t in block.block_trigger_inputs if t is not None)
+
+ # If block has sub_blocks, recursively check them
+ if block.sub_blocks:
+ nested_triggers = fn_recursive_get_trigger(block.sub_blocks)
+ trigger_values.update(nested_triggers)
+
+ return trigger_values
+
+ return fn_recursive_get_trigger(self.sub_blocks)
+
+ def get_execution_blocks(self, **kwargs) -> "SequentialPipelineBlocks":
+ """
+ Get the blocks that would execute given the specified inputs.
+
+ As the traversal walks through sequential blocks, intermediate outputs from resolved blocks are added to the
+ active inputs. This means conditional blocks that depend on intermediates (e.g., "run img2img if image_latents
+ is present") will resolve correctly, as long as the condition is based on presence/absence (None or not None),
+ not on the actual value.
+
+
+ Args:
+ **kwargs: Input names and values. Only trigger inputs affect block selection.
+
+ Returns:
+ SequentialPipelineBlocks containing only the blocks that would execute
+ """
+ # Copy kwargs so we can add outputs as we traverse
+ active_inputs = dict(kwargs)
+
+ def fn_recursive_traverse(block, block_name, active_inputs):
+ result_blocks = OrderedDict()
+
+ # ConditionalPipelineBlocks (includes AutoPipelineBlocks)
+ if isinstance(block, ConditionalPipelineBlocks):
+ block = block.get_execution_blocks(**active_inputs)
+ if block is None:
+ return result_blocks
+
+ # Has sub_blocks (SequentialPipelineBlocks/ConditionalPipelineBlocks)
+ if block.sub_blocks and not isinstance(block, LoopSequentialPipelineBlocks):
+ for sub_block_name, sub_block in block.sub_blocks.items():
+ nested_blocks = fn_recursive_traverse(sub_block, sub_block_name, active_inputs)
+ nested_blocks = {f"{block_name}.{k}": v for k, v in nested_blocks.items()}
+ result_blocks.update(nested_blocks)
+ else:
+ # Leaf block: single ModularPipelineBlocks or LoopSequentialPipelineBlocks
+ result_blocks[block_name] = block
+ # Add outputs to active_inputs so subsequent blocks can use them as triggers
+ if hasattr(block, "intermediate_outputs"):
+ for out in block.intermediate_outputs:
+ active_inputs[out.name] = True
+
+ return result_blocks
+
+ all_blocks = OrderedDict()
+ for block_name, block in self.sub_blocks.items():
+ nested_blocks = fn_recursive_traverse(block, block_name, active_inputs)
+ all_blocks.update(nested_blocks)
+
+ return SequentialPipelineBlocks.from_blocks_dict(all_blocks)
+
+ def __repr__(self):
+ class_name = self.__class__.__name__
+ base_class = self.__class__.__bases__[0].__name__
+ header = (
+ f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n"
+ )
+
+ if self._workflow_map is None and self._get_trigger_inputs():
+ header += "\n"
+ header += " " + "=" * 100 + "\n"
+ header += " This pipeline contains blocks that are selected at runtime based on inputs.\n"
+ header += f" Trigger Inputs: {[inp for inp in self._get_trigger_inputs() if inp is not None]}\n"
+ # Get first trigger input as example
+ example_input = next(t for t in self._get_trigger_inputs() if t is not None)
+ header += f" Use `get_execution_blocks()` to see selected blocks (e.g. `get_execution_blocks({example_input}=...)`).\n"
+ header += " " + "=" * 100 + "\n\n"
+
+ description = self.description
+ if self._workflow_map is not None:
+ workflow_str = format_workflow(self._workflow_map)
+ description = f"{self.description}\n\n{workflow_str}"
+
+ # Format description with proper indentation
+ desc_lines = description.split("\n")
+ desc = []
+ # First line with "Description:" label
+ desc.append(f" Description: {desc_lines[0]}")
+ # Subsequent lines with proper indentation
+ if len(desc_lines) > 1:
+ desc.extend(f" {line}" for line in desc_lines[1:])
+ desc = "\n".join(desc) + "\n"
+
+ # Components section - focus only on expected components
+ expected_components = getattr(self, "expected_components", [])
+ components_str = format_components(expected_components, indent_level=2, add_empty_lines=False)
+
+ # Configs section - use format_configs with add_empty_lines=False
+ expected_configs = getattr(self, "expected_configs", [])
+ configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False)
+
+ # Blocks section - moved to the end with simplified format
+ blocks_str = " Sub-Blocks:\n"
+ for i, (name, block) in enumerate(self.sub_blocks.items()):
+ # show execution order
+ blocks_str += f" [{i}] {name} ({block.__class__.__name__})\n"
+
+ # Add block description
+ desc_lines = block.description.split("\n")
+ indented_desc = desc_lines[0]
+ if len(desc_lines) > 1:
+ indented_desc += "\n" + "\n".join(" " + line for line in desc_lines[1:])
+ blocks_str += f" Description: {indented_desc}\n\n"
+
+ # Build the representation with conditional sections
+ result = f"{header}\n{desc}"
+
+ # Only add components section if it has content
+ if components_str.strip():
+ result += f"\n\n{components_str}"
+
+ # Only add configs section if it has content
+ if configs_str.strip():
+ result += f"\n\n{configs_str}"
+
+ # Always add blocks section
+ result += f"\n\n{blocks_str})"
+
+ return result
+
+ @property
+ def doc(self):
+ description = self.description
+ if self._workflow_map is not None:
+ workflow_str = format_workflow(self._workflow_map)
+ description = f"{self.description}\n\n{workflow_str}"
+
+ return make_doc_string(
+ self.inputs,
+ self.outputs,
+ description=description,
+ class_name=self.__class__.__name__,
+ expected_components=self.expected_components,
+ expected_configs=self.expected_configs,
+ )
+
+ @property
+ def _requirements(self) -> dict[str, str]:
+ requirements = {}
+ for block_name, block in self.sub_blocks.items():
+ if getattr(block, "_requirements", None):
+ requirements[block_name] = block._requirements
+ return requirements
+
+
+class LoopSequentialPipelineBlocks(ModularPipelineBlocks):
+ """
+ A Pipeline blocks that combines multiple pipeline block classes into a For Loop. When called, it will call each
+ block in sequence.
+
+ This class inherits from [`ModularPipelineBlocks`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipeline blocks (such as loading or saving etc.)
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+
+ Attributes:
+ block_classes: list of block classes to be used
+ block_names: list of prefixes for each block
+ """
+
+ model_name = None
+ block_classes = []
+ block_names = []
+
+ @property
+ def description(self) -> str:
+ """Description of the block. Must be implemented by subclasses."""
+ raise NotImplementedError("description method must be implemented in subclasses")
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return []
+
+ @property
+ def loop_expected_configs(self) -> list[ConfigSpec]:
+ return []
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ """list of input parameters. Must be implemented by subclasses."""
+ return []
+
+ @property
+ def loop_required_inputs(self) -> list[str]:
+ input_names = []
+ for input_param in self.loop_inputs:
+ if input_param.required:
+ input_names.append(input_param.name)
+ return input_names
+
+ @property
+ def loop_intermediate_outputs(self) -> list[OutputParam]:
+ """list of intermediate output parameters. Must be implemented by subclasses."""
+ return []
+
+ # modified from SequentialPipelineBlocks to include loop_expected_components
+ @property
+ def expected_components(self):
+ expected_components = []
+ for block in self.sub_blocks.values():
+ for component in block.expected_components:
+ if component not in expected_components:
+ expected_components.append(component)
+ for component in self.loop_expected_components:
+ if component not in expected_components:
+ expected_components.append(component)
+ return expected_components
+
+ # modified from SequentialPipelineBlocks to include loop_expected_configs
+ @property
+ def expected_configs(self):
+ expected_configs = []
+ for block in self.sub_blocks.values():
+ for config in block.expected_configs:
+ if config not in expected_configs:
+ expected_configs.append(config)
+ for config in self.loop_expected_configs:
+ if config not in expected_configs:
+ expected_configs.append(config)
+ return expected_configs
+
+ def _get_inputs(self):
+ inputs = []
+ inputs.extend(self.loop_inputs)
+ outputs = set()
+
+ for name, block in self.sub_blocks.items():
+ # Add inputs that aren't in outputs yet
+ for inp in block.inputs:
+ if inp.name not in outputs and inp not in inputs:
+ inputs.append(inp)
+
+ # Add this block's outputs
+ block_intermediate_outputs = [out.name for out in block.intermediate_outputs]
+ outputs.update(block_intermediate_outputs)
+
+ for input_param in inputs:
+ if input_param.name in self.required_inputs:
+ input_param.required = True
+ else:
+ input_param.required = False
+
+ return inputs
+
+ @property
+ # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks.inputs
+ def inputs(self):
+ return self._get_inputs()
+
+ # modified from SequentialPipelineBlocks, if any additionan input required by the loop is required by the block
+ @property
+ def required_inputs(self) -> list[str]:
+ # Get the first block from the dictionary
+ first_block = next(iter(self.sub_blocks.values()))
+ required_by_any = set(getattr(first_block, "required_inputs", set()))
+
+ required_by_loop = set(getattr(self, "loop_required_inputs", set()))
+ required_by_any.update(required_by_loop)
+
+ # Union with required inputs from all other blocks
+ for block in list(self.sub_blocks.values())[1:]:
+ block_required = set(getattr(block, "required_inputs", set()))
+ required_by_any.update(block_required)
+
+ return list(required_by_any)
+
+ # YiYi TODO: this need to be thought about more
+ # modified from SequentialPipelineBlocks to include loop_intermediate_outputs
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ named_outputs = [(name, block.intermediate_outputs) for name, block in self.sub_blocks.items()]
+ combined_outputs = combine_outputs(*named_outputs)
+ for output in self.loop_intermediate_outputs:
+ if output.name not in {output.name for output in combined_outputs}:
+ combined_outputs.append(output)
+ return combined_outputs
+
+ # YiYi TODO: this need to be thought about more
+ @property
+ def outputs(self) -> list[str]:
+ return next(reversed(self.sub_blocks.values())).intermediate_outputs
+
+ @property
+ # Copied from diffusers.modular_pipelines.modular_pipeline.SequentialPipelineBlocks._requirements
+ def _requirements(self) -> dict[str, str]:
+ requirements = {}
+ for block_name, block in self.sub_blocks.items():
+ if getattr(block, "_requirements", None):
+ requirements[block_name] = block._requirements
+ return requirements
+
+ def __init__(self):
+ sub_blocks = InsertableDict()
+ for block_name, block in zip(self.block_names, self.block_classes):
+ if inspect.isclass(block):
+ sub_blocks[block_name] = block()
+ else:
+ sub_blocks[block_name] = block
+ self.sub_blocks = sub_blocks
+
+ # Validate that sub_blocks are only leaf blocks
+ for block_name, block in self.sub_blocks.items():
+ if block.sub_blocks:
+ raise ValueError(
+ f"In {self.__class__.__name__}, sub_blocks must be leaf blocks (no sub_blocks). "
+ f"Block '{block_name}' ({block.__class__.__name__}) has sub_blocks."
+ )
+
+ @classmethod
+ def from_blocks_dict(cls, blocks_dict: dict[str, Any]) -> "LoopSequentialPipelineBlocks":
+ """
+ Creates a LoopSequentialPipelineBlocks instance from a dictionary of blocks.
+
+ Args:
+ blocks_dict: Dictionary mapping block names to block instances
+
+ Returns:
+ A new LoopSequentialPipelineBlocks instance
+ """
+ instance = cls()
+
+ # Create instances if classes are provided
+ sub_blocks = InsertableDict()
+ for name, block in blocks_dict.items():
+ if inspect.isclass(block):
+ sub_blocks[name] = block()
+ else:
+ sub_blocks[name] = block
+
+ instance.block_classes = [block.__class__ for block in blocks_dict.values()]
+ instance.block_names = list(blocks_dict.keys())
+ instance.sub_blocks = blocks_dict
+ return instance
+
+ def loop_step(self, components, state: PipelineState, **kwargs):
+ for block_name, block in self.sub_blocks.items():
+ try:
+ components, state = block(components, state, **kwargs)
+ except Exception as e:
+ error_msg = (
+ f"\nError in block: ({block_name}, {block.__class__.__name__})\n"
+ f"Error details: {str(e)}\n"
+ f"Traceback:\n{traceback.format_exc()}"
+ )
+ logger.error(error_msg)
+ raise
+ return components, state
+
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ raise NotImplementedError("`__call__` method needs to be implemented by the subclass")
+
+ @property
+ def doc(self):
+ return make_doc_string(
+ self.inputs,
+ self.outputs,
+ self.description,
+ class_name=self.__class__.__name__,
+ expected_components=self.expected_components,
+ expected_configs=self.expected_configs,
+ )
+
+ # modified from SequentialPipelineBlocks,
+ # (does not need trigger_inputs related part so removed them,
+ # do not need to support auto block for loop blocks)
+ def __repr__(self):
+ class_name = self.__class__.__name__
+ base_class = self.__class__.__bases__[0].__name__
+ header = (
+ f"{class_name}(\n Class: {base_class}\n" if base_class and base_class != "object" else f"{class_name}(\n"
+ )
+
+ # Format description with proper indentation
+ desc_lines = self.description.split("\n")
+ desc = []
+ # First line with "Description:" label
+ desc.append(f" Description: {desc_lines[0]}")
+ # Subsequent lines with proper indentation
+ if len(desc_lines) > 1:
+ desc.extend(f" {line}" for line in desc_lines[1:])
+ desc = "\n".join(desc) + "\n"
+
+ # Components section - focus only on expected components
+ expected_components = getattr(self, "expected_components", [])
+ components_str = format_components(expected_components, indent_level=2, add_empty_lines=False)
+
+ # Configs section - use format_configs with add_empty_lines=False
+ expected_configs = getattr(self, "expected_configs", [])
+ configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False)
+
+ # Blocks section - moved to the end with simplified format
+ blocks_str = " Sub-Blocks:\n"
+ for i, (name, block) in enumerate(self.sub_blocks.items()):
+ # For SequentialPipelineBlocks, show execution order
+ blocks_str += f" [{i}] {name} ({block.__class__.__name__})\n"
+
+ # Add block description
+ desc_lines = block.description.split("\n")
+ indented_desc = desc_lines[0]
+ if len(desc_lines) > 1:
+ indented_desc += "\n" + "\n".join(" " + line for line in desc_lines[1:])
+ blocks_str += f" Description: {indented_desc}\n\n"
+
+ # Build the representation with conditional sections
+ result = f"{header}\n{desc}"
+
+ # Only add components section if it has content
+ if components_str.strip():
+ result += f"\n\n{components_str}"
+
+ # Only add configs section if it has content
+ if configs_str.strip():
+ result += f"\n\n{configs_str}"
+
+ # Always add blocks section
+ result += f"\n\n{blocks_str})"
+
+ return result
+
+ @torch.compiler.disable
+ def progress_bar(self, iterable=None, total=None):
+ if not hasattr(self, "_progress_bar_config"):
+ self._progress_bar_config = {}
+ elif not isinstance(self._progress_bar_config, dict):
+ raise ValueError(
+ f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
+ )
+
+ if iterable is not None:
+ return tqdm(iterable, **self._progress_bar_config)
+ elif total is not None:
+ return tqdm(total=total, **self._progress_bar_config)
+ else:
+ raise ValueError("Either `total` or `iterable` has to be defined.")
+
+ def set_progress_bar_config(self, **kwargs):
+ self._progress_bar_config = kwargs
+
+
+# YiYi TODO:
+# 1. look into the serialization of modular_model_index.json, make sure the items are properly ordered like model_index.json (currently a mess)
+# 2. do we need ConfigSpec? the are basically just key/val kwargs
+# 3. imnprove docstring and potentially add validator for methods where we accept kwargs to be passed to from_pretrained/save_pretrained/load_components()
+class ModularPipeline(ConfigMixin, PushToHubMixin):
+ """
+ Base class for all Modular pipelines.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+
+ Args:
+ blocks: ModularPipelineBlocks, the blocks to be used in the pipeline
+ """
+
+ config_name = "modular_model_index.json"
+ hf_device_map = None
+ default_blocks_name = None
+
+ # YiYi TODO: add warning for passing multiple ComponentSpec/ConfigSpec with the same name
+ def __init__(
+ self,
+ blocks: ModularPipelineBlocks | None = None,
+ pretrained_model_name_or_path: str | os.PathLike | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
+ workflow: str | None = None,
+ modular_config_dict: dict[str, Any] | None = None,
+ config_dict: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ """
+ Initialize a ModularPipeline instance.
+
+ This method sets up the pipeline by:
+ - creating default pipeline blocks if not provided
+ - gather component and config specifications based on the pipeline blocks's requirement (e.g.
+ expected_components, expected_configs)
+ - update the loading specs of from_pretrained components based on the modular_model_index.json file from
+ huggingface hub if `pretrained_model_name_or_path` is provided
+ - create defaultfrom_config components and register everything
+
+ Args:
+ blocks: `ModularPipelineBlocks` instance. If None, will attempt to load
+ default blocks based on the pipeline class name.
+ pretrained_model_name_or_path: Path to a pretrained pipeline configuration. Can be None if the pipeline
+ does not require any additional loading config. If provided, will first try to load component specs
+ (only for from_pretrained components) and config values from `modular_model_index.json`, then
+ fallback to `model_index.json` for compatibility with standard non-modular repositories.
+ components_manager:
+ Optional ComponentsManager for managing multiple component cross different pipelines and apply
+ offloading strategies.
+ collection: Optional collection name for organizing components in the ComponentsManager.
+ workflow: Optional workflow name. If provided, the blocks are pruned to that workflow's execution
+ blocks (see `ModularPipelineBlocks.get_workflow`), so the pipeline only expects — and
+ `load_components()` only loads — the components that workflow uses.
+ **kwargs: Additional arguments passed to `load_config()` when loading pretrained configuration.
+
+ Examples:
+ ```python
+ # Initialize with custom blocks
+ pipeline = ModularPipeline(blocks=my_custom_blocks)
+
+ # Initialize from pretrained configuration
+ pipeline = ModularPipeline(blocks=my_blocks, pretrained_model_name_or_path="my-repo/modular-pipeline")
+
+ # Initialize with components manager
+ pipeline = ModularPipeline(
+ blocks=my_blocks, components_manager=ComponentsManager(), collection="my_collection"
+ )
+ ```
+
+ Notes:
+ - If blocks is None, the method will try to find default blocks based on the pipeline class name
+ - Components with default_creation_method="from_config" are created immediately, its specs are not included
+ in config dict and will not be saved in `modular_model_index.json`
+ - Components with default_creation_method="from_pretrained" are set to None and can be loaded later with
+ `load_components()` (with or without specific component names)
+ - The pipeline's config dict is populated with component specs (only for from_pretrained components) and
+ config values, which will be saved as `modular_model_index.json` during `save_pretrained`
+ - The pipeline's config dict is also used to store the pipeline blocks's class name, which will be saved as
+ `_blocks_class_name` in the config dict
+ """
+
+ if modular_config_dict is None and config_dict is None and pretrained_model_name_or_path is not None:
+ cache_dir = kwargs.pop("cache_dir", None)
+ force_download = kwargs.pop("force_download", False)
+ proxies = kwargs.pop("proxies", None)
+ token = kwargs.pop("token", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ revision = kwargs.pop("revision", None)
+
+ load_config_kwargs = {
+ "cache_dir": cache_dir,
+ "force_download": force_download,
+ "proxies": proxies,
+ "token": token,
+ "local_files_only": local_files_only,
+ "revision": revision,
+ }
+
+ modular_config_dict, config_dict = self._load_pipeline_config(
+ pretrained_model_name_or_path, **load_config_kwargs
+ )
+
+ if blocks is None:
+ if modular_config_dict is not None:
+ blocks_class_name = modular_config_dict.get("_blocks_class_name")
+ else:
+ blocks_class_name = self.default_blocks_name
+ if blocks_class_name is not None:
+ diffusers_module = importlib.import_module("diffusers")
+ blocks_class = getattr(diffusers_module, blocks_class_name, None)
+ # If the blocks_class is not found or is a base class (e.g. SequentialPipelineBlocks saved by from_blocks_dict) with empty block_classes
+ # fall back to default_blocks_name
+ if blocks_class is None or not blocks_class.block_classes:
+ blocks_class_name = self.default_blocks_name
+ blocks_class = getattr(diffusers_module, blocks_class_name)
+
+ if blocks_class is not None:
+ blocks = blocks_class()
+ else:
+ logger.warning(f"`blocks` is `None`, no default blocks class found for {self.__class__.__name__}")
+
+ if workflow is not None:
+ if blocks is None:
+ raise ValueError(f"`workflow={workflow!r}` requires pipeline blocks, but none could be resolved.")
+ blocks = blocks.get_workflow(workflow)
+
+ self._blocks = blocks
+ self._components_manager = components_manager
+ self._collection = collection
+ self._component_specs = {spec.name: deepcopy(spec) for spec in self._blocks.expected_components}
+ self._config_specs = {spec.name: deepcopy(spec) for spec in self._blocks.expected_configs}
+
+ # update component_specs and config_specs based on modular_model_index.json
+ if modular_config_dict is not None:
+ for name, value in modular_config_dict.items():
+ # all the components in modular_model_index.json are from_pretrained components
+ if name in self._component_specs and isinstance(value, (tuple, list)) and len(value) == 3:
+ library, class_name, component_spec_dict = value
+ component_spec = self._dict_to_component_spec(name, component_spec_dict)
+ component_spec.default_creation_method = "from_pretrained"
+ self._component_specs[name] = component_spec
+
+ elif name in self._config_specs:
+ self._config_specs[name].default = value
+
+ # if `modular_config_dict` is None (i.e. `modular_model_index.json` is not found), update based on `config_dict` (i.e. `model_index.json`)
+ elif config_dict is not None:
+ for name, value in config_dict.items():
+ if name in self._component_specs and isinstance(value, (tuple, list)) and len(value) == 2:
+ library, class_name = value
+ component_spec_dict = {
+ "repo": pretrained_model_name_or_path,
+ "subfolder": name,
+ "type_hint": (library, class_name),
+ }
+ component_spec = self._dict_to_component_spec(name, component_spec_dict)
+ component_spec.default_creation_method = "from_pretrained"
+ self._component_specs[name] = component_spec
+ elif name in self._config_specs:
+ self._config_specs[name].default = value
+
+ if len(kwargs) > 0:
+ logger.warning(f"Unexpected input '{kwargs.keys()}' provided. This input will be ignored.")
+
+ register_components_dict = {}
+ for name, component_spec in self._component_specs.items():
+ if component_spec.default_creation_method == "from_config":
+ component = component_spec.create()
+ else:
+ component = None
+ register_components_dict[name] = component
+ self.register_components(**register_components_dict)
+
+ default_configs = {}
+ for name, config_spec in self._config_specs.items():
+ default_configs[name] = config_spec.default
+ self.register_to_config(**default_configs)
+ self.register_to_config(
+ _blocks_class_name=self._blocks.__class__.__name__ if self._blocks is not None else None
+ )
+
+ self._pretrained_model_name_or_path = pretrained_model_name_or_path
+
+ @property
+ def default_call_parameters(self) -> dict[str, Any]:
+ """
+ Returns:
+ - Dictionary mapping input names to their default values
+ """
+ params = {}
+ for input_param in self._blocks.inputs:
+ params[input_param.name] = input_param.default
+ return params
+
+ @classmethod
+ def _load_pipeline_config(
+ cls,
+ pretrained_model_name_or_path: str | os.PathLike | None,
+ **load_config_kwargs,
+ ):
+ try:
+ # try to load modular_model_index.json
+ modular_config_dict = cls.load_config(pretrained_model_name_or_path, **load_config_kwargs)
+ return modular_config_dict, None
+
+ except EnvironmentError as e:
+ logger.debug(f" modular_model_index.json not found in the repo: {e}")
+
+ try:
+ logger.debug(" try to load model_index.json")
+ from diffusers import DiffusionPipeline
+
+ config_dict = DiffusionPipeline.load_config(pretrained_model_name_or_path, **load_config_kwargs)
+ return None, config_dict
+
+ except EnvironmentError as e:
+ raise EnvironmentError(
+ f"Failed to load config from '{pretrained_model_name_or_path}'. "
+ f"Could not find or load 'modular_model_index.json' or 'model_index.json'."
+ ) from e
+
+ return None, None
+
+ @classmethod
+ @validate_hf_hub_args
+ def from_pretrained(
+ cls,
+ pretrained_model_name_or_path: str | os.PathLike | None,
+ trust_remote_code: bool | None = None,
+ components_manager: ComponentsManager | None = None,
+ collection: str | None = None,
+ workflow: str | None = None,
+ **kwargs,
+ ):
+ """
+ Load a ModularPipeline from a huggingface hub repo.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`, optional):
+ Path to a pretrained pipeline configuration. It will first try to load config from
+ `modular_model_index.json`, then fallback to `model_index.json` for compatibility with standard
+ non-modular repositories. If the pretrained_model_name_or_path does not contain any pipeline config, it
+ will be set to None during initialization.
+ trust_remote_code (`bool`, optional):
+ Whether to trust remote code when loading the pipeline, need to be set to True if you want to create
+ pipeline blocks based on the custom code in `pretrained_model_name_or_path`
+ components_manager (`ComponentsManager`, optional):
+ ComponentsManager instance for managing multiple component cross different pipelines and apply
+ offloading strategies.
+ collection (`str`, optional):`
+ Collection name for organizing components in the ComponentsManager.
+ workflow (`str`, optional):
+ Name of a workflow declared by the pipeline blocks. If provided, the blocks are pruned to that
+ workflow's execution blocks, so the pipeline only expects — and `load_components()` only loads —
+ the components that workflow uses.
+ """
+ from ..pipelines.pipeline_loading_utils import _get_pipeline_class
+
+ try:
+ blocks = ModularPipelineBlocks.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except EnvironmentError as e:
+ logger.debug(f"EnvironmentError: {e}")
+ blocks = None
+
+ cache_dir = kwargs.pop("cache_dir", None)
+ force_download = kwargs.pop("force_download", False)
+ proxies = kwargs.pop("proxies", None)
+ token = kwargs.pop("token", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ revision = kwargs.pop("revision", None)
+
+ load_config_kwargs = {
+ "cache_dir": cache_dir,
+ "force_download": force_download,
+ "proxies": proxies,
+ "token": token,
+ "local_files_only": local_files_only,
+ "revision": revision,
+ }
+
+ modular_config_dict, config_dict = cls._load_pipeline_config(
+ pretrained_model_name_or_path, **load_config_kwargs
+ )
+
+ if modular_config_dict is not None:
+ pipeline_class = _get_pipeline_class(cls, config=modular_config_dict)
+ elif config_dict is not None:
+ from diffusers.pipelines.auto_pipeline import _get_model
+
+ logger.debug(" try to determine the modular pipeline class from model_index.json")
+ standard_pipeline_class = _get_pipeline_class(cls, config=config_dict)
+ model_name = _get_model(standard_pipeline_class.__name__)
+ map_fn = MODULAR_PIPELINE_MAPPING.get(model_name, _create_default_map_fn("ModularPipeline"))
+ pipeline_class_name = map_fn(config_dict)
+ diffusers_module = importlib.import_module("diffusers")
+ pipeline_class = getattr(diffusers_module, pipeline_class_name)
+ else:
+ # there is no config for modular pipeline, assuming that the pipeline block does not need any from_pretrained components
+ pipeline_class = cls
+ pretrained_model_name_or_path = None
+
+ pipeline = pipeline_class(
+ blocks=blocks,
+ pretrained_model_name_or_path=pretrained_model_name_or_path,
+ components_manager=components_manager,
+ collection=collection,
+ workflow=workflow,
+ modular_config_dict=modular_config_dict,
+ config_dict=config_dict,
+ **kwargs,
+ )
+ return pipeline
+
+ def save_pretrained(
+ self,
+ save_directory: str | os.PathLike,
+ safe_serialization: bool = True,
+ variant: str | None = None,
+ max_shard_size: int | str | None = None,
+ push_to_hub: bool = False,
+ **kwargs,
+ ):
+ """
+ Save the pipeline and all its components to a directory, so that it can be re-loaded using the
+ [`~ModularPipeline.from_pretrained`] class method.
+
+ Args:
+ save_directory (`str` or `os.PathLike`):
+ Directory to save the pipeline to. Will be created if it doesn't exist.
+ safe_serialization (`bool`, *optional*, defaults to `True`):
+ Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
+ variant (`str`, *optional*):
+ If specified, weights are saved in the format `pytorch_model..bin`.
+ max_shard_size (`int` or `str`, defaults to `None`):
+ The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
+ lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5GB"`).
+ If expressed as an integer, the unit is bytes.
+ push_to_hub (`bool`, *optional*, defaults to `False`):
+ Whether to push the pipeline to the Hugging Face model hub after saving it.
+ **kwargs: Additional keyword arguments:
+ - `overwrite_modular_index` (`bool`, *optional*, defaults to `False`):
+ When saving a Modular Pipeline, its components in `modular_model_index.json` may reference repos
+ different from the destination repo. Setting this to `True` updates all component references in
+ `modular_model_index.json` so they point to the repo specified by `repo_id`.
+ - `repo_id` (`str`, *optional*):
+ The repository ID to push the pipeline to. Defaults to the last component of `save_directory`.
+ - `commit_message` (`str`, *optional*):
+ Commit message for the push to hub operation.
+ - `private` (`bool`, *optional*):
+ Whether the repository should be private.
+ - `create_pr` (`bool`, *optional*, defaults to `False`):
+ Whether to create a pull request instead of pushing directly.
+ - `token` (`str`, *optional*):
+ The Hugging Face token to use for authentication.
+ """
+ overwrite_modular_index = kwargs.pop("overwrite_modular_index", False)
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+
+ if push_to_hub:
+ commit_message = kwargs.pop("commit_message", None)
+ private = kwargs.pop("private", None)
+ create_pr = kwargs.pop("create_pr", False)
+ token = kwargs.pop("token", None)
+ update_model_card = kwargs.pop("update_model_card", False)
+ repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
+
+ for component_name, component_spec in self._component_specs.items():
+ if component_spec.default_creation_method != "from_pretrained":
+ continue
+
+ component = getattr(self, component_name, None)
+ if component is None:
+ continue
+
+ model_cls = component.__class__
+ if is_compiled_module(component):
+ component = _unwrap_model(component)
+ model_cls = component.__class__
+
+ save_method_name = None
+ for library_name, library_classes in LOADABLE_CLASSES.items():
+ if library_name in sys.modules:
+ library = importlib.import_module(library_name)
+ else:
+ logger.info(
+ f"{library_name} is not installed. Cannot save {component_name} as {library_classes} from {library_name}"
+ )
+ continue
+
+ for base_class, save_load_methods in library_classes.items():
+ class_candidate = getattr(library, base_class, None)
+ if class_candidate is not None and issubclass(model_cls, class_candidate):
+ save_method_name = save_load_methods[0]
+ break
+ if save_method_name is not None:
+ break
+
+ if save_method_name is None:
+ logger.warning(f"self.{component_name}={component} of type {type(component)} cannot be saved.")
+ continue
+
+ save_method = getattr(component, save_method_name)
+ save_method_signature = inspect.signature(save_method)
+ save_method_accept_safe = "safe_serialization" in save_method_signature.parameters
+ save_method_accept_variant = "variant" in save_method_signature.parameters
+ save_method_accept_max_shard_size = "max_shard_size" in save_method_signature.parameters
+
+ save_kwargs = {}
+ if save_method_accept_safe:
+ save_kwargs["safe_serialization"] = safe_serialization
+ if save_method_accept_variant:
+ save_kwargs["variant"] = variant
+ if save_method_accept_max_shard_size and max_shard_size is not None:
+ save_kwargs["max_shard_size"] = max_shard_size
+
+ component_save_path = os.path.join(save_directory, component_name)
+ save_method(component_save_path, **save_kwargs)
+
+ if component_name not in self.config:
+ continue
+
+ has_no_load_id = not hasattr(component, "_diffusers_load_id") or component._diffusers_load_id == "null"
+ if overwrite_modular_index or has_no_load_id:
+ library, class_name, component_spec_dict = self.config[component_name]
+ component_spec_dict["pretrained_model_name_or_path"] = repo_id if push_to_hub else save_directory
+ component_spec_dict["subfolder"] = component_name
+ self.register_to_config(**{component_name: (library, class_name, component_spec_dict)})
+
+ self.save_config(save_directory=save_directory)
+
+ if push_to_hub:
+ card_content = generate_modular_model_card_content(self.blocks)
+ model_card = load_or_create_model_card(
+ repo_id,
+ token=token,
+ is_pipeline=True,
+ model_description=MODULAR_MODEL_CARD_TEMPLATE.format(**card_content),
+ is_modular=True,
+ update_model_card=update_model_card,
+ )
+ model_card = populate_model_card(model_card, tags=card_content["tags"])
+ model_card.save(os.path.join(save_directory, "README.md"))
+
+ self._upload_folder(
+ save_directory,
+ repo_id,
+ token=token,
+ commit_message=commit_message,
+ create_pr=create_pr,
+ )
+
+ @property
+ def doc(self):
+ """
+ Returns:
+ - The docstring of the pipeline blocks
+ """
+ return self._blocks.doc
+
+ @property
+ def blocks(self) -> ModularPipelineBlocks:
+ """
+ Returns:
+ - A copy of the pipeline blocks
+ """
+ return deepcopy(self._blocks)
+
+ def register_components(self, **kwargs):
+ """
+ Register components with their corresponding specifications.
+
+ This method is responsible for:
+ 1. Sets component objects as attributes on the loader (e.g., self.unet = unet)
+ 2. Updates the config dict, which will be saved as `modular_model_index.json` during `save_pretrained` (only
+ for from_pretrained components)
+ 3. Adds components to the component manager if one is attached (only for from_pretrained components)
+
+ This method is called when:
+ - Components are first initialized in __init__:
+ - from_pretrained components not loaded during __init__ so they are registered as None;
+ - non from_pretrained components are created during __init__ and registered as the object itself
+ - Components are updated with the `update_components()` method: e.g. loader.update_components(unet=unet) or
+ loader.update_components(guider=guider_spec)
+ - (from_pretrained) Components are loaded with the `load_components()` method: e.g.
+ loader.load_components(names=["unet"]) or loader.load_components() to load all default components
+
+ Args:
+ **kwargs: Keyword arguments where keys are component names and values are component objects.
+ E.g., register_components(unet=unet_model, text_encoder=encoder_model)
+
+ Notes:
+ - When registering None for a component, it sets attribute to None but still syncs specs with the config
+ dict, which will be saved as `modular_model_index.json` during `save_pretrained`
+ - component_specs are updated to match the new component outside of this method, e.g. in
+ `update_components()` method
+ """
+ for name, module in kwargs.items():
+ # current component spec
+ component_spec = self._component_specs.get(name)
+ if component_spec is None:
+ logger.warning(f"ModularPipeline.register_components: skipping unknown component '{name}'")
+ continue
+
+ # check if it is the first time registration, i.e. calling from __init__
+ is_registered = hasattr(self, name)
+ is_from_pretrained = component_spec.default_creation_method == "from_pretrained"
+
+ if module is not None:
+ # actual library and class name of the module
+ library, class_name = _fetch_class_library_tuple(module) # e.g. ("diffusers", "UNet2DConditionModel")
+ else:
+ # if module is None, e.g. self.register_components(unet=None) during __init__
+ # we do not update the spec,
+ # but we still need to update the modular_model_index.json config based on component spec
+ library, class_name = None, None
+
+ # extract the loading spec from the updated component spec that'll be used as part of modular_model_index.json config
+ # e.g. {"pretrained_model_name_or_path": "stabilityai/stable-diffusion-2-1",
+ # "type_hint": ("diffusers", "UNet2DConditionModel"),
+ # "subfolder": "unet",
+ # "variant": None,
+ # "revision": None}
+ component_spec_dict = self._component_spec_to_dict(component_spec)
+
+ register_dict = {name: (library, class_name, component_spec_dict)}
+
+ # set the component as attribute
+ # if it is not set yet, just set it and skip the process to check and warn below
+ if not is_registered:
+ if is_from_pretrained:
+ self.register_to_config(**register_dict)
+ setattr(self, name, module)
+ if module is not None and is_from_pretrained and self._components_manager is not None:
+ self._components_manager.add(name, module, self._collection)
+ continue
+
+ current_module = getattr(self, name, None)
+ # skip if the component is already registered with the same object
+ if current_module is module:
+ logger.info(
+ f"ModularPipeline.register_components: {name} is already registered with same object, skipping"
+ )
+ continue
+
+ # warn if unregister
+ if current_module is not None and module is None:
+ logger.info(
+ f"ModularPipeline.register_components: setting '{name}' to None "
+ f"(was {current_module.__class__.__name__})"
+ )
+ # same type, new instance → replace but send debug log
+ elif (
+ current_module is not None
+ and module is not None
+ and isinstance(module, current_module.__class__)
+ and current_module != module
+ ):
+ logger.debug(
+ f"ModularPipeline.register_components: replacing existing '{name}' "
+ f"(same type {type(current_module).__name__}, new instance)"
+ )
+
+ # update modular_model_index.json config
+ if is_from_pretrained:
+ self.register_to_config(**register_dict)
+ # finally set models
+ setattr(self, name, module)
+ # add to component manager if one is attached
+ if module is not None and is_from_pretrained and self._components_manager is not None:
+ self._components_manager.add(name, module, self._collection)
+
+ @property
+ def device(self) -> torch.device:
+ r"""
+ Returns:
+ `torch.device`: The torch device on which the pipeline is located.
+ """
+ modules = self.components.values()
+ modules = [m for m in modules if isinstance(m, torch.nn.Module)]
+
+ for module in modules:
+ return module.device
+
+ return torch.device("cpu")
+
+ @property
+ # Modified from diffusers.pipelines.pipeline_utils.DiffusionPipeline._execution_device
+ def _execution_device(self):
+ r"""
+ Returns the device on which the pipeline's models will be executed. After calling
+ [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from
+ Accelerate's module hooks.
+ """
+ from ..hooks.group_offloading import _get_group_onload_device
+
+ for name, model in self.components.items():
+ if not isinstance(model, torch.nn.Module):
+ continue
+ try:
+ return _get_group_onload_device(model)
+ except ValueError:
+ pass
+
+ for name, model in self.components.items():
+ if not isinstance(model, torch.nn.Module):
+ continue
+
+ if not hasattr(model, "_hf_hook"):
+ return self.device
+ for module in model.modules():
+ if (
+ hasattr(module, "_hf_hook")
+ and hasattr(module._hf_hook, "execution_device")
+ and module._hf_hook.execution_device is not None
+ ):
+ return torch.device(module._hf_hook.execution_device)
+ return self.device
+
+ @property
+ def dtype(self) -> torch.dtype:
+ r"""
+ Returns:
+ `torch.dtype`: The torch dtype on which the pipeline is located.
+ """
+ modules = self.components.values()
+ modules = [m for m in modules if isinstance(m, torch.nn.Module)]
+
+ for module in modules:
+ return module.dtype
+
+ return torch.float32
+
+ @property
+ def null_component_names(self) -> list[str]:
+ """
+ Returns:
+ - list of names for components that needs to be loaded
+ """
+ return [name for name in self._component_specs.keys() if hasattr(self, name) and getattr(self, name) is None]
+
+ @property
+ def component_names(self) -> list[str]:
+ """
+ Returns:
+ - list of names for all components
+ """
+ return list(self.components.keys())
+
+ @property
+ def pretrained_component_names(self) -> list[str]:
+ """
+ Returns:
+ - list of names for from_pretrained components
+ """
+ return [
+ name
+ for name in self._component_specs.keys()
+ if self._component_specs[name].default_creation_method == "from_pretrained"
+ ]
+
+ @property
+ def config_component_names(self) -> list[str]:
+ """
+ Returns:
+ - list of names for from_config components
+ """
+ return [
+ name
+ for name in self._component_specs.keys()
+ if self._component_specs[name].default_creation_method == "from_config"
+ ]
+
+ @property
+ def components(self) -> dict[str, Any]:
+ """
+ Returns:
+ - Dictionary mapping component names to their objects (include both from_pretrained and from_config
+ components)
+ """
+ # return only components we've actually set as attributes on self
+ return {name: getattr(self, name) for name in self._component_specs.keys() if hasattr(self, name)}
+
+ def get_component_spec(self, name: str) -> ComponentSpec:
+ """
+ Returns:
+ - a copy of the ComponentSpec object for the given component name
+ """
+ return deepcopy(self._component_specs[name])
+
+ def update_components(self, **kwargs):
+ """
+ Update components and configuration values and specs after the pipeline has been instantiated.
+
+ This method allows you to:
+ 1. Replace existing components with new ones (e.g., updating `self.unet` or `self.text_encoder`)
+ 2. Update configuration values (e.g., changing `self.requires_safety_checker` flag)
+
+ In addition to updating the components and configuration values as pipeline attributes, the method also
+ updates:
+ - the corresponding specs in `_component_specs` and `_config_specs`
+ - the `config` dict, which will be saved as `modular_model_index.json` during `save_pretrained`
+
+ Args:
+ **kwargs: Component objects or configuration values to update:
+ - Component objects: Models loaded with `AutoModel.from_pretrained()` or `ComponentSpec.load()`
+ are automatically tagged with loading information. ConfigMixin objects without weights (e.g.,
+ schedulers, guiders) can be passed directly.
+ - Configuration values: Simple values to update configuration settings
+ (e.g., `requires_safety_checker=False`)
+
+ Examples:
+ ```python
+ # Update pre-trained model
+ pipeline.update_components(unet=new_unet_model, text_encoder=new_text_encoder)
+
+ # Update configuration values
+ pipeline.update_components(requires_safety_checker=False)
+ ```
+
+ Notes:
+ - Components loaded with `AutoModel.from_pretrained()` or `ComponentSpec.load()` will have
+ loading specs preserved for serialization. Custom or locally loaded components without Hub references will
+ have their `modular_model_index.json` entries updated automatically during `save_pretrained()`.
+ - ConfigMixin objects without weights (e.g., schedulers, guiders) can be passed directly.
+ """
+
+ passed_components = {k: kwargs.pop(k) for k in self._component_specs if k in kwargs}
+ passed_config_values = {k: kwargs.pop(k) for k in self._config_specs if k in kwargs}
+
+ for name, component in passed_components.items():
+ current_component_spec = self._component_specs[name]
+
+ # log if type changed
+ if current_component_spec.type_hint is not None and not isinstance(
+ component, current_component_spec.type_hint
+ ):
+ logger.info(
+ f"ModularPipeline.update_components: adding {name} with new type: {component.__class__.__name__}, previous type: {current_component_spec.type_hint.__name__}"
+ )
+ # update _component_specs based on the new component
+ if component is None:
+ new_component_spec = current_component_spec
+ if hasattr(self, name) and getattr(self, name) is not None:
+ logger.warning(f"ModularPipeline.update_components: setting {name} to None (spec unchanged)")
+ elif (
+ current_component_spec.default_creation_method == "from_pretrained"
+ and getattr(component, "_diffusers_load_id", None) is None
+ ):
+ new_component_spec = ComponentSpec(name=name, type_hint=type(component))
+ else:
+ new_component_spec = ComponentSpec.from_component(name, component)
+
+ if new_component_spec.default_creation_method != current_component_spec.default_creation_method:
+ logger.info(
+ f"ModularPipeline.update_components: changing the default_creation_method of {name} from {current_component_spec.default_creation_method} to {new_component_spec.default_creation_method}."
+ )
+
+ self._component_specs[name] = new_component_spec
+
+ if len(kwargs) > 0:
+ logger.warning(f"Unexpected keyword arguments, will be ignored: {kwargs.keys()}")
+
+ self.register_components(**passed_components)
+
+ config_to_register = {}
+ for name, new_value in passed_config_values.items():
+ self._config_specs[name].default = new_value
+ config_to_register[name] = new_value
+ self.register_to_config(**config_to_register)
+
+ def load_components(self, names: list[str] | str | None = None, workflow: str | None = None, **kwargs):
+ """
+ Load selected components from specs.
+
+ Args:
+ names: list of component names to load. If None, will load all components with
+ default_creation_method == "from_pretrained". If provided as a list or string, will load only the
+ specified components.
+ workflow: name of a workflow declared by the pipeline blocks. If provided, only the components that
+ workflow's execution blocks use are loaded. Cannot be combined with `names`.
+ **kwargs: additional kwargs to be passed to `from_pretrained()`.Can be:
+ - a single value to be applied to all components to be loaded, e.g. dtype=torch.bfloat16
+ - a dict, e.g. dtype={"unet": torch.bfloat16, "default": torch.float32}
+ - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g.
+ `pretrained_model_name_or_path`, `variant`, `revision`, etc.
+ - if potentially override ComponentSpec if passed a different loading field in kwargs, e.g.
+ `pretrained_model_name_or_path`, `variant`, `revision`, etc.
+ """
+
+ if workflow is not None and names is not None:
+ raise ValueError("Pass either `names` or `workflow`, not both.")
+
+ if names is None:
+ names = [
+ name
+ for name in self._component_specs.keys()
+ if self._component_specs[name].default_creation_method == "from_pretrained"
+ and self._component_specs[name].pretrained_model_name_or_path is not None
+ and getattr(self, name, None) is None
+ ]
+ elif isinstance(names, str):
+ names = [names]
+ elif not isinstance(names, list):
+ raise ValueError(f"Invalid type for names: {type(names)}")
+
+ if workflow is not None:
+ workflow_component_names = {spec.name for spec in self.blocks.get_workflow(workflow).expected_components}
+ names = [name for name in names if name in workflow_component_names]
+
+ components_to_load = {name for name in names if name in self._component_specs}
+ unknown_names = {name for name in names if name not in self._component_specs}
+ if len(unknown_names) > 0:
+ logger.warning(f"Unknown components will be ignored: {unknown_names}")
+
+ components_to_register = {}
+ for name in components_to_load:
+ spec = self._component_specs[name]
+ component_load_kwargs = {}
+ for key, value in kwargs.items():
+ if not isinstance(value, dict):
+ # if the value is a single value, apply it to all components
+ component_load_kwargs[key] = value
+ else:
+ if name in value:
+ # if it is a dict, check if the component name is in the dict
+ component_load_kwargs[key] = value[name]
+ elif "default" in value:
+ # check if the default is specified
+ component_load_kwargs[key] = value["default"]
+ # Only pass trust_remote_code to components from the same repo as the pipeline.
+ # When a user passes trust_remote_code=True, they intend to trust code from the
+ # pipeline's repo, not from external repos referenced in modular_model_index.json.
+ trust_remote_code_stripped = False
+ if (
+ "trust_remote_code" in component_load_kwargs
+ and self._pretrained_model_name_or_path is not None
+ and spec.pretrained_model_name_or_path != self._pretrained_model_name_or_path
+ ):
+ component_load_kwargs.pop("trust_remote_code")
+ trust_remote_code_stripped = True
+
+ if not spec.pretrained_model_name_or_path:
+ logger.info(f"Skipping component `{name}`: no pretrained model path specified.")
+ continue
+
+ try:
+ components_to_register[name] = spec.load(**component_load_kwargs)
+ except Exception:
+ tb = traceback.format_exc()
+ if trust_remote_code_stripped and "trust_remote_code" in tb:
+ warning_msg = (
+ f"Failed to load component `{name}` from external repository "
+ f"`{spec.pretrained_model_name_or_path}`.\n\n"
+ f"`trust_remote_code=True` was not forwarded to `{name}` because it comes from "
+ f"a different repository than the pipeline (`{self._pretrained_model_name_or_path}`). "
+ f"For safety, `trust_remote_code` is only forwarded to components from the same "
+ f"repository as the pipeline.\n\n"
+ f"You need to load this component manually with `trust_remote_code=True` and pass it "
+ f"to the pipeline via `pipe.update_components()`. For example, if it is a custom model:\n\n"
+ f' {name} = AutoModel.from_pretrained("{spec.pretrained_model_name_or_path}", trust_remote_code=True)\n'
+ f" pipe.update_components({name}={name})\n"
+ )
+ else:
+ warning_msg = (
+ f"Failed to create component {name}:\n"
+ f"- Component spec: {spec}\n"
+ f"- load() called with kwargs: {component_load_kwargs}\n"
+ "If this component is not required for your workflow you can safely ignore this message.\n\n"
+ "Traceback:\n"
+ f"{tb}"
+ )
+ logger.warning(warning_msg)
+
+ # Register all components at once
+ self.register_components(**components_to_register)
+
+ def unload_components(self, names: list[str] | str):
+ """
+ Unload selected components, freeing their memory.
+
+ The component attribute is set back to `None` and, if a ComponentsManager is attached, the component is
+ removed from it. The component spec is untouched, so the component can be loaded again later with
+ `load_components()`.
+
+ Args:
+ names: component name or list of component names to unload.
+ """
+ if isinstance(names, str):
+ names = [names]
+ elif not isinstance(names, list):
+ raise ValueError(f"Invalid type for names: {type(names)}")
+
+ unknown_names = {name for name in names if name not in self._component_specs}
+ if len(unknown_names) > 0:
+ logger.warning(f"Unknown components will be ignored: {unknown_names}")
+
+ for name in names:
+ if name in unknown_names:
+ continue
+ # not holding the full `components` dict across iterations: a lingering reference would keep the
+ # unloaded component alive through the gc pass below
+ component = self.components.get(name)
+ if component is None:
+ continue
+ component_id = f"{name}_{id(component)}"
+ self.register_components(**{name: None})
+ del component
+ if self._components_manager is not None and component_id in self._components_manager.components:
+ self._components_manager.remove(component_id)
+
+ gc.collect()
+ empty_device_cache()
+
+ # Copied from diffusers.pipelines.pipeline_utils.DiffusionPipeline._maybe_raise_error_if_group_offload_active
+ def _maybe_raise_error_if_group_offload_active(
+ self, raise_error: bool = False, module: torch.nn.Module | None = None
+ ) -> bool:
+ from ..hooks.group_offloading import _is_group_offload_enabled
+
+ components = self.components.values() if module is None else [module]
+ components = [component for component in components if isinstance(component, torch.nn.Module)]
+ for component in components:
+ if _is_group_offload_enabled(component):
+ if raise_error:
+ raise ValueError(
+ "You are trying to apply model/sequential CPU offloading to a pipeline that contains components "
+ "with group offloading enabled. This is not supported. Please disable group offloading for "
+ "components of the pipeline to use other offloading methods."
+ )
+ return True
+ return False
+
+ # Modified from diffusers.pipelines.pipeline_utils.DiffusionPipeline.to
+ def to(self, *args, **kwargs) -> Self:
+ r"""
+ Performs Pipeline dtype and/or device conversion. A torch.dtype and torch.device are inferred from the
+ arguments of `self.to(*args, **kwargs).`
+
+ > [!TIP] > If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is.
+ Otherwise, > the returned pipeline is a copy of self with the desired torch.dtype and torch.device.
+
+
+ Here are the ways to call `to`:
+
+ - `to(dtype, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified
+ [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype)
+ - `to(device, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified
+ [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device)
+ - `to(device=None, dtype=None, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the
+ specified [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) and
+ [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype)
+
+ Arguments:
+ dtype (`torch.dtype`, *optional*):
+ Returns a pipeline with the specified
+ [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype)
+ device (`torch.Device`, *optional*):
+ Returns a pipeline with the specified
+ [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device)
+ silence_dtype_warnings (`str`, *optional*, defaults to `False`):
+ Whether to omit warnings if the target `dtype` is not compatible with the target `device`.
+
+ Returns:
+ [`DiffusionPipeline`]: The pipeline converted to specified `dtype` and/or `dtype`.
+ """
+ from ..pipelines.pipeline_utils import _check_bnb_status
+ from ..utils import is_accelerate_available, is_accelerate_version, is_hpu_available, is_transformers_version
+
+ dtype = kwargs.pop("dtype", None)
+ device = kwargs.pop("device", None)
+ silence_dtype_warnings = kwargs.pop("silence_dtype_warnings", False)
+
+ dtype_arg = None
+ device_arg = None
+ if len(args) == 1:
+ if isinstance(args[0], torch.dtype):
+ dtype_arg = args[0]
+ else:
+ device_arg = torch.device(args[0]) if args[0] is not None else None
+ elif len(args) == 2:
+ if isinstance(args[0], torch.dtype):
+ raise ValueError(
+ "When passing two arguments, make sure the first corresponds to `device` and the second to `dtype`."
+ )
+ device_arg = torch.device(args[0]) if args[0] is not None else None
+ dtype_arg = args[1]
+ elif len(args) > 2:
+ raise ValueError("Please make sure to pass at most two arguments (`device` and `dtype`) `.to(...)`")
+
+ if dtype is not None and dtype_arg is not None:
+ raise ValueError(
+ "You have passed `dtype` both as an argument and as a keyword argument. Please only pass one of the two."
+ )
+
+ dtype = dtype or dtype_arg
+
+ if device is not None and device_arg is not None:
+ raise ValueError(
+ "You have passed `device` both as an argument and as a keyword argument. Please only pass one of the two."
+ )
+
+ device = device or device_arg
+ device_type = torch.device(device).type if device is not None else None
+ pipeline_has_bnb = any(any((_check_bnb_status(module))) for _, module in self.components.items())
+
+ # throw warning if pipeline is in "offloaded"-mode but user tries to manually set to GPU.
+ def module_is_sequentially_offloaded(module):
+ if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"):
+ return False
+
+ _, _, is_loaded_in_8bit_bnb = _check_bnb_status(module)
+
+ if is_loaded_in_8bit_bnb:
+ return False
+
+ return hasattr(module, "_hf_hook") and (
+ isinstance(module._hf_hook, accelerate.hooks.AlignDevicesHook)
+ or hasattr(module._hf_hook, "hooks")
+ and isinstance(module._hf_hook.hooks[0], accelerate.hooks.AlignDevicesHook)
+ )
+
+ def module_is_offloaded(module):
+ if not is_accelerate_available() or is_accelerate_version("<", "0.17.0.dev0"):
+ return False
+
+ return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, accelerate.hooks.CpuOffload)
+
+ # .to("cuda") would raise an error if the pipeline is sequentially offloaded, so we raise our own to make it clearer
+ pipeline_is_sequentially_offloaded = any(
+ module_is_sequentially_offloaded(module) for _, module in self.components.items()
+ )
+
+ is_pipeline_device_mapped = self.hf_device_map is not None and len(self.hf_device_map) > 1
+ if is_pipeline_device_mapped:
+ raise ValueError(
+ "It seems like you have activated a device mapping strategy on the pipeline which doesn't allow explicit device placement using `to()`. You can call `reset_device_map()` to remove the existing device map from the pipeline."
+ )
+
+ if device_type in ["cuda", "xpu"]:
+ if pipeline_is_sequentially_offloaded and not pipeline_has_bnb:
+ raise ValueError(
+ "It seems like you have activated sequential model offloading by calling `enable_sequential_cpu_offload`, but are now attempting to move the pipeline to GPU. This is not compatible with offloading. Please, move your pipeline `.to('cpu')` or consider removing the move altogether if you use sequential offloading."
+ )
+ # PR: https://github.com/huggingface/accelerate/pull/3223/
+ elif pipeline_has_bnb and is_accelerate_version("<", "1.1.0.dev0"):
+ raise ValueError(
+ "You are trying to call `.to('cuda')` on a pipeline that has models quantized with `bitsandbytes`. Your current `accelerate` installation does not support it. Please upgrade the installation."
+ )
+
+ # Display a warning in this case (the operation succeeds but the benefits are lost)
+ pipeline_is_offloaded = any(module_is_offloaded(module) for _, module in self.components.items())
+ if pipeline_is_offloaded and device_type in ["cuda", "xpu"]:
+ logger.warning(
+ f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading."
+ )
+
+ # Enable generic support for Intel Gaudi accelerator using GPU/HPU migration
+ if device_type == "hpu" and kwargs.pop("hpu_migration", True) and is_hpu_available():
+ os.environ["PT_HPU_GPU_MIGRATION"] = "1"
+ logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1")
+
+ import habana_frameworks.torch # noqa: F401
+
+ # HPU hardware check
+ if not (hasattr(torch, "hpu") and torch.hpu.is_available()):
+ raise ValueError("You are trying to call `.to('hpu')` but HPU device is unavailable.")
+
+ os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "1"
+ logger.debug("Environment variable set: PT_HPU_MAX_COMPOUND_OP_SIZE=1")
+
+ modules = self.components.values()
+ modules = [m for m in modules if isinstance(m, torch.nn.Module)]
+
+ is_offloaded = pipeline_is_offloaded or pipeline_is_sequentially_offloaded
+ for module in modules:
+ _, is_loaded_in_4bit_bnb, is_loaded_in_8bit_bnb = _check_bnb_status(module)
+ is_group_offloaded = self._maybe_raise_error_if_group_offload_active(module=module)
+
+ if (is_loaded_in_4bit_bnb or is_loaded_in_8bit_bnb) and dtype is not None:
+ logger.warning(
+ f"The module '{module.__class__.__name__}' has been loaded in `bitsandbytes` {'4bit' if is_loaded_in_4bit_bnb else '8bit'} and conversion to {dtype} is not supported. Module is still in {'4bit' if is_loaded_in_4bit_bnb else '8bit'} precision."
+ )
+
+ if is_loaded_in_8bit_bnb and device is not None:
+ logger.warning(
+ f"The module '{module.__class__.__name__}' has been loaded in `bitsandbytes` 8bit and moving it to {device} via `.to()` is not supported. Module is still on {module.device}."
+ )
+
+ # Note: we also handle this at the ModelMixin level. The reason for doing it here too is that modeling
+ # components can be from outside diffusers too, but still have group offloading enabled.
+ if (
+ self._maybe_raise_error_if_group_offload_active(raise_error=False, module=module)
+ and device is not None
+ ):
+ logger.warning(
+ f"The module '{module.__class__.__name__}' is group offloaded and moving it to {device} via `.to()` is not supported."
+ )
+
+ # This can happen for `transformer` models. CPU placement was added in
+ # https://github.com/huggingface/transformers/pull/33122. So, we guard this accordingly.
+ if is_loaded_in_4bit_bnb and device is not None and is_transformers_version(">", "4.44.0"):
+ module.to(device=device)
+ elif not is_loaded_in_4bit_bnb and not is_loaded_in_8bit_bnb and not is_group_offloaded:
+ module.to(device, dtype)
+
+ if (
+ module.dtype == torch.float16
+ and str(device) in ["cpu"]
+ and not silence_dtype_warnings
+ and not is_offloaded
+ ):
+ logger.warning(
+ "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It"
+ " is not recommended to move them to `cpu` as running them will fail. Please make"
+ " sure to use an accelerator to run the pipeline in inference, due to the lack of"
+ " support for`float16` operations on this device in PyTorch. Please, remove the"
+ " `dtype=torch.float16` argument, or use another device for inference."
+ )
+ return self
+
+ @staticmethod
+ def _component_spec_to_dict(component_spec: ComponentSpec) -> Any:
+ """
+ Convert a ComponentSpec into a JSON‐serializable dict for saving as an entry in `modular_model_index.json`. If
+ the `default_creation_method` is not `from_pretrained`, return None.
+
+ This dict contains:
+ - "type_hint": tuple[str, str]
+ Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel"))
+ - All loading fields defined by `component_spec.loading_fields()`, typically:
+ - "pretrained_model_name_or_path": str | None
+ The model pretrained_model_name_or_pathsitory (e.g., "stabilityai/stable-diffusion-xl").
+ - "subfolder": str | None
+ A subfolder within the pretrained_model_name_or_path where this component lives.
+ - "variant": str | None
+ An optional variant identifier for the model.
+ - "revision": str | None
+ A specific git revision (commit hash, tag, or branch).
+ - ... any other loading fields defined on the spec.
+
+ Args:
+ component_spec (ComponentSpec):
+ The spec object describing one pipeline component.
+
+ Returns:
+ dict[str, Any]: A mapping suitable for JSON serialization.
+
+ Example:
+ >>> from diffusers.pipelines.modular_pipeline_utils import ComponentSpec >>> from diffusers import
+ UNet2DConditionModel >>> spec = ComponentSpec(
+ ... name="unet", ... type_hint=UNet2DConditionModel, ... config=None, ...
+ pretrained_model_name_or_path="path/to/pretrained_model_name_or_path", ... subfolder="subfolder", ...
+ variant=None, ... revision=None, ... default_creation_method="from_pretrained",
+ ... ) >>> ModularPipeline._component_spec_to_dict(spec) {
+ "type_hint": ("diffusers", "UNet2DConditionModel"), "pretrained_model_name_or_path": "path/to/repo",
+ "subfolder": "subfolder", "variant": None, "revision": None, "type_hint": ("diffusers",
+ "UNet2DConditionModel"), "pretrained_model_name_or_path": "path/to/repo", "subfolder": "subfolder",
+ "variant": None, "revision": None,
+ }
+ """
+ if component_spec.default_creation_method != "from_pretrained":
+ return None
+
+ if component_spec.type_hint is not None:
+ lib_name, cls_name = _fetch_class_library_tuple(component_spec.type_hint)
+ else:
+ lib_name = None
+ cls_name = None
+ load_spec_dict = {k: getattr(component_spec, k) for k in component_spec.loading_fields()}
+ return {
+ "type_hint": (lib_name, cls_name),
+ **load_spec_dict,
+ }
+
+ @staticmethod
+ def _dict_to_component_spec(name: str, spec_dict: dict[str, Any]) -> ComponentSpec:
+ """
+ Reconstruct a ComponentSpec from a loading specdict.
+
+ This method converts a dictionary representation back into a ComponentSpec object. The dict should contain:
+ - "type_hint": tuple[str, str]
+ Library name and class name of the component. (e.g. ("diffusers", "UNet2DConditionModel"))
+ - All loading fields defined by `component_spec.loading_fields()`, typically:
+ - "pretrained_model_name_or_path": str | None
+ The model repository (e.g., "stabilityai/stable-diffusion-xl").
+ - "subfolder": str | None
+ A subfolder within the pretrained_model_name_or_path where this component lives.
+ - "variant": str | None
+ An optional variant identifier for the model.
+ - "revision": str | None
+ A specific git revision (commit hash, tag, or branch).
+ - ... any other loading fields defined on the spec.
+
+ Args:
+ name (str):
+ The name of the component.
+ specdict (dict[str, Any]):
+ A dictionary containing the component specification data.
+
+ Returns:
+ ComponentSpec: A reconstructed ComponentSpec object.
+
+ Example:
+ >>> spec_dict = { ... "type_hint": ("diffusers", "UNet2DConditionModel"), ...
+ "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl", ... "subfolder": "unet", ... "variant":
+ None, ... "revision": None, ... } >>> ModularPipeline._dict_to_component_spec("unet", spec_dict)
+ ComponentSpec(
+ name="unet", type_hint=UNet2DConditionModel, config=None,
+ pretrained_model_name_or_path="stabilityai/stable-diffusion-xl", subfolder="unet", variant=None,
+ revision=None, default_creation_method="from_pretrained"
+ >>> spec_dict = { ... "type_hint": ("diffusers", "UNet2DConditionModel"), ...
+ "pretrained_model_name_or_path": "stabilityai/stable-diffusion-xl", ... "subfolder": "unet", ... "variant":
+ None, ... "revision": None, ... } >>> ModularPipeline._dict_to_component_spec("unet", spec_dict)
+ ComponentSpec(
+ name="unet", type_hint=UNet2DConditionModel, config=None,
+ pretrained_model_name_or_path="stabilityai/stable-diffusion-xl", subfolder="unet", variant=None,
+ revision=None, default_creation_method="from_pretrained"
+ )
+ """
+ # make a shallow copy so we can pop() safely
+ spec_dict = spec_dict.copy()
+ # pull out and resolve the stored type_hint
+ lib_name, cls_name = spec_dict.pop("type_hint")
+ if lib_name is not None and cls_name is not None:
+ type_hint = simple_get_class_obj(lib_name, cls_name)
+ else:
+ type_hint = None
+
+ # re‐assemble the ComponentSpec
+ return ComponentSpec(
+ name=name,
+ type_hint=type_hint,
+ **spec_dict,
+ )
+
+ def set_progress_bar_config(self, **kwargs):
+ for sub_block_name, sub_block in self._blocks.sub_blocks.items():
+ if hasattr(sub_block, "set_progress_bar_config"):
+ sub_block.set_progress_bar_config(**kwargs)
+
+ def __call__(self, state: PipelineState = None, output: str | list[str] = None, **kwargs):
+ """
+ Execute the pipeline by running the pipeline blocks with the given inputs.
+
+ Args:
+ state (`PipelineState`, optional):
+ PipelineState instance contains inputs and intermediate values. If None, a new `PipelineState` will be
+ created based on the user inputs and the pipeline blocks's requirement.
+ output (`str` or `list[str]`, optional):
+ Optional specification of what to return:
+ - None: Returns the complete `PipelineState` with all inputs and intermediates (default)
+ - str: Returns a specific intermediate value from the state (e.g. `output="image"`)
+ - list[str]: Returns a dictionary of specific intermediate values (e.g. `output=["image",
+ "latents"]`)
+
+
+ Examples:
+ ```python
+ # Get complete pipeline state
+ state = pipeline(prompt="A beautiful sunset", num_inference_steps=20)
+ print(state.intermediates) # All intermediate outputs
+
+ # Get specific output
+ image = pipeline(prompt="A beautiful sunset", output="image")
+
+ # Get multiple specific outputs
+ results = pipeline(prompt="A beautiful sunset", output=["image", "latents"])
+ image, latents = results["image"], results["latents"]
+
+ # Continue from previous state
+ state = pipeline(prompt="A beautiful sunset")
+ new_state = pipeline(state=state, output="image") # Continue processing
+ ```
+
+ Returns:
+ - If `output` is None: Complete `PipelineState` containing all inputs and intermediates
+ - If `output` is str: The specific intermediate value from the state (e.g. `output="image"`)
+ - If `output` is list[str]: Dictionary mapping output names to their values from the state (e.g.
+ `output=["image", "latents"]`)
+ """
+ if state is None:
+ state = PipelineState()
+ else:
+ state = deepcopy(state)
+
+ # Make a copy of the input kwargs
+ passed_kwargs = kwargs.copy()
+
+ # Add inputs to state, using defaults if not provided in the kwargs or the state
+ # if same input already in the state, will override it if provided in the kwargs
+ for expected_input_param in self._blocks.inputs:
+ name = expected_input_param.name
+ default = expected_input_param.default
+ kwargs_type = expected_input_param.kwargs_type
+ if name in passed_kwargs:
+ state.set(name, passed_kwargs.pop(name), kwargs_type)
+ elif kwargs_type is not None and kwargs_type in passed_kwargs:
+ kwargs_dict = passed_kwargs.pop(kwargs_type)
+ for k, v in kwargs_dict.items():
+ state.set(k, v, kwargs_type)
+ elif name is not None and name not in state.values:
+ state.set(name, default, kwargs_type)
+
+ # Warn about unexpected inputs
+ if len(passed_kwargs) > 0:
+ warnings.warn(f"Unexpected input '{passed_kwargs.keys()}' provided. This input will be ignored.")
+ # Run the pipeline
+ with torch.no_grad():
+ try:
+ _, state = self._blocks(self, state)
+ except Exception:
+ error_msg = f"Error in block: ({self._blocks.__class__.__name__}):\n"
+ logger.error(error_msg)
+ raise
+
+ if output is None:
+ return state
+
+ if isinstance(output, str):
+ return state.get(output)
+
+ elif isinstance(output, (list, tuple)):
+ return state.get(output)
+ else:
+ raise ValueError(f"Output '{output}' is not a valid output type")
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5898565ad2ffe083256b911f9e963b98a744e36
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/modular_pipeline_utils.py
@@ -0,0 +1,1314 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import re
+import warnings
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from types import UnionType
+from typing import Any, Literal, Type, Union, get_args, get_origin
+
+import PIL.Image
+import torch
+from packaging.specifiers import InvalidSpecifier, SpecifierSet
+
+from ..configuration_utils import ConfigMixin, FrozenDict
+from ..loaders.single_file_utils import _is_single_file_path_or_url
+from ..utils import DIFFUSERS_LOAD_ID_FIELDS, _resolve_dtype, is_sdnq_available, is_torch_available, logging
+from ..utils.constants import DIFFUSERS_SDNQ_TRANSFORMERS
+from ..utils.import_utils import _is_package_available
+
+
+if is_torch_available():
+ pass
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Template for modular pipeline model card description with placeholders
+MODULAR_MODEL_CARD_TEMPLATE = """{model_description}
+
+## Example Usage
+
+[TODO]
+
+## Pipeline Architecture
+
+This modular pipeline is composed of the following blocks:
+
+{blocks_description} {trigger_inputs_section}
+
+## Model Components
+
+{components_description} {configs_section}
+
+{io_specification_section}
+"""
+
+
+class InsertableDict(OrderedDict):
+ def insert(self, key, value, index):
+ items = list(self.items())
+
+ # Remove key if it already exists to avoid duplicates
+ items = [(k, v) for k, v in items if k != key]
+
+ # Insert at the specified index
+ items.insert(index, (key, value))
+
+ # Clear and update self
+ self.clear()
+ self.update(items)
+
+ # Return self for method chaining
+ return self
+
+ def __repr__(self):
+ if not self:
+ return "InsertableDict()"
+
+ items = []
+ for i, (key, value) in enumerate(self.items()):
+ if isinstance(value, type):
+ # For classes, show class name and
+ obj_repr = f""
+ else:
+ # For objects (instances) and other types, show class name and module
+ obj_repr = f""
+ items.append(f"{i}: ({repr(key)}, {obj_repr})")
+
+ return "InsertableDict([\n " + ",\n ".join(items) + "\n])"
+
+
+# YiYi TODO:
+# 1. validate the dataclass fields
+# 2. improve the docstring and potentially add a validator for load methods, make sure they are valid inputs to pass to from_pretrained()
+@dataclass
+class ComponentSpec:
+ """Specification for a pipeline component.
+
+ A component can be created in two ways:
+ 1. From scratch using __init__ with a config dict
+ 2. using `from_pretrained`
+
+ Attributes:
+ name: Name of the component
+ type_hint: Type of the component (e.g. UNet2DConditionModel)
+ description: Optional description of the component
+ config: Optional config dict for __init__ creation
+ pretrained_model_name_or_path: Optional pretrained_model_name_or_path path for from_pretrained creation
+ subfolder: Optional subfolder in pretrained_model_name_or_path
+ variant: Optional variant in pretrained_model_name_or_path
+ revision: Optional revision in pretrained_model_name_or_path
+ default_creation_method: Preferred creation method - "from_config" or "from_pretrained"
+ """
+
+ name: str | None = None
+ type_hint: Type | None = None
+ description: str | None = None
+ config: FrozenDict | None = None
+ pretrained_model_name_or_path: str | list[str] | None = field(default=None, metadata={"loading": True})
+ subfolder: str | None = field(default="", metadata={"loading": True})
+ variant: str | None = field(default=None, metadata={"loading": True})
+ revision: str | None = field(default=None, metadata={"loading": True})
+ default_creation_method: Literal["from_config", "from_pretrained"] = "from_pretrained"
+
+ # Deprecated
+ repo: str | list[str] | None = field(default=None, metadata={"loading": False})
+
+ def __post_init__(self):
+ repo_value = self.repo
+ if repo_value is not None and self.pretrained_model_name_or_path is None:
+ object.__setattr__(self, "pretrained_model_name_or_path", repo_value)
+
+ def __hash__(self):
+ """Make ComponentSpec hashable, using load_id as the hash value."""
+ return hash((self.name, self.load_id, self.default_creation_method))
+
+ def __eq__(self, other):
+ """Compare ComponentSpec objects based on name and load_id."""
+ if not isinstance(other, ComponentSpec):
+ return False
+ return (
+ self.name == other.name
+ and self.load_id == other.load_id
+ and self.default_creation_method == other.default_creation_method
+ )
+
+ @classmethod
+ def from_component(cls, name: str, component: Any) -> Any:
+ """Create a ComponentSpec from a Component.
+
+ Currently supports:
+ - Components created with `ComponentSpec.load()` method
+ - Components that are ConfigMixin subclasses but not nn.Modules (e.g. schedulers, guiders)
+
+ Args:
+ name: Name of the component
+ component: Component object to create spec from
+
+ Returns:
+ ComponentSpec object
+
+ Raises:
+ ValueError: If component is not supported (e.g. nn.Module without load_id, non-ConfigMixin)
+ """
+
+ # Check if component was created with ComponentSpec.load()
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id != "null":
+ # component has a usable load_id -> from_pretrained, no warning needed
+ default_creation_method = "from_pretrained"
+ else:
+ # Component doesn't have a usable load_id, check if it's a nn.Module
+ if isinstance(component, torch.nn.Module):
+ raise ValueError(
+ "Cannot create ComponentSpec from a nn.Module that was not created with `ComponentSpec.load()` method."
+ )
+ # ConfigMixin objects without weights (e.g. scheduler & guider) can be recreated with from_config
+ elif isinstance(component, ConfigMixin):
+ # warn if component was not created with `ComponentSpec`
+ if not hasattr(component, "_diffusers_load_id"):
+ logger.warning(
+ "Component was not created using `ComponentSpec`, defaulting to `from_config` creation method"
+ )
+ default_creation_method = "from_config"
+ else:
+ # Not a ConfigMixin and not created with `ComponentSpec.load()` method -> throw error
+ raise ValueError(
+ f"Cannot create ComponentSpec from {name}({component.__class__.__name__}). Currently ComponentSpec.from_component() only supports: "
+ f" - components created with `ComponentSpec.load()` method"
+ f" - components that are a subclass of ConfigMixin but not a nn.Module (e.g. guider, scheduler)."
+ )
+
+ type_hint = component.__class__
+
+ if isinstance(component, ConfigMixin) and default_creation_method == "from_config":
+ config = component.config
+ else:
+ config = None
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id != "null":
+ load_spec = cls.decode_load_id(component._diffusers_load_id)
+ else:
+ load_spec = {}
+
+ return cls(
+ name=name, type_hint=type_hint, config=config, default_creation_method=default_creation_method, **load_spec
+ )
+
+ @classmethod
+ def loading_fields(cls) -> list[str]:
+ """
+ Return the names of all loading‐related fields (i.e. those whose field.metadata["loading"] is True).
+ """
+ return DIFFUSERS_LOAD_ID_FIELDS.copy()
+
+ @property
+ def load_id(self) -> str:
+ """
+ Unique identifier for this spec's pretrained load, composed of
+ pretrained_model_name_or_path|subfolder|variant|revision (no empty segments).
+ """
+ if self.default_creation_method == "from_config":
+ return "null"
+ parts = [getattr(self, k) for k in self.loading_fields()]
+ parts = ["null" if p is None else p for p in parts]
+ return "|".join(parts)
+
+ @classmethod
+ def decode_load_id(cls, load_id: str) -> dict[str, str | None]:
+ """
+ Decode a load_id string back into a dictionary of loading fields and values.
+
+ Args:
+ load_id: The load_id string to decode, format: "pretrained_model_name_or_path|subfolder|variant|revision"
+ where None values are represented as "null"
+
+ Returns:
+ Dict mapping loading field names to their values. e.g. {
+ "pretrained_model_name_or_path": "path/to/repo", "subfolder": "subfolder", "variant": "variant",
+ "revision": "revision"
+ } If a segment value is "null", it's replaced with None. Returns None if load_id is "null" (indicating
+ component not created with `load` method).
+ """
+
+ # Get all loading fields in order
+ loading_fields = cls.loading_fields()
+ result = dict.fromkeys(loading_fields)
+
+ if load_id == "null":
+ return result
+
+ # Split the load_id
+ parts = load_id.split("|")
+
+ # Map parts to loading fields by position
+ for i, part in enumerate(parts):
+ if i < len(loading_fields):
+ # Convert "null" string back to None
+ result[loading_fields[i]] = None if part == "null" else part
+
+ return result
+
+ # YiYi TODO: I think we should only support ConfigMixin for this method (after we make guider and image_processors config mixin)
+ # otherwise we cannot do spec -> spec.create() -> component -> ComponentSpec.from_component(component)
+ # the config info is lost in the process
+ # remove error check in from_component spec and ModularPipeline.update_components() if we remove support for non configmixin in `create()` method
+ def create(self, config: FrozenDict | dict[str, Any] | None = None, **kwargs) -> Any:
+ """Create component using from_config with config."""
+
+ if self.type_hint is None or not isinstance(self.type_hint, type):
+ raise ValueError("`type_hint` is required when using from_config creation method.")
+
+ config = config or self.config or {}
+
+ if issubclass(self.type_hint, ConfigMixin):
+ component = self.type_hint.from_config(config, **kwargs)
+ else:
+ signature_params = inspect.signature(self.type_hint.__init__).parameters
+ init_kwargs = {}
+ for k, v in config.items():
+ if k in signature_params:
+ init_kwargs[k] = v
+ for k, v in kwargs.items():
+ if k in signature_params:
+ init_kwargs[k] = v
+ component = self.type_hint(**init_kwargs)
+
+ component._diffusers_load_id = "null"
+ if hasattr(component, "config"):
+ self.config = component.config
+
+ return component
+
+ # YiYi TODO: add guard for type of model, if it is supported by from_pretrained
+ def load(self, **kwargs) -> Any:
+ """Load component using from_pretrained."""
+ torch_dtype = kwargs.pop("torch_dtype", None)
+ if torch_dtype is not None:
+ kwargs["dtype"] = _resolve_dtype(kwargs.get("dtype"), torch_dtype)
+
+ # select loading fields from kwargs passed from user: e.g. pretrained_model_name_or_path, subfolder, variant, revision, note the list could change
+ passed_loading_kwargs = {key: kwargs.pop(key) for key in self.loading_fields() if key in kwargs}
+ # merge loading field value in the spec with user passed values to create load_kwargs
+ load_kwargs = {key: passed_loading_kwargs.get(key, getattr(self, key)) for key in self.loading_fields()}
+
+ pretrained_model_name_or_path = load_kwargs.pop("pretrained_model_name_or_path", None)
+ if pretrained_model_name_or_path is None:
+ raise ValueError(
+ "`pretrained_model_name_or_path` info is required when using `load` method (you can directly set it in `pretrained_model_name_or_path` field of the ComponentSpec or pass it as an argument)"
+ )
+ is_single_file = _is_single_file_path_or_url(pretrained_model_name_or_path)
+ if is_single_file and self.type_hint is None:
+ raise ValueError(
+ f"`type_hint` is required when loading a single file model but is missing for component: {self.name}"
+ )
+
+ from diffusers import AutoModel
+
+ # `dtype` is not an accepted parameter for tokenizers and processors.
+ # As a result, it gets stored in `init_kwargs`, which are written to the config
+ # during save. This causes JSON serialization to fail when saving the component.
+ if self.type_hint is not None and not issubclass(self.type_hint, (torch.nn.Module, AutoModel)):
+ kwargs.pop("dtype", None)
+
+ if self.type_hint is None:
+ try:
+ component = AutoModel.from_pretrained(pretrained_model_name_or_path, **load_kwargs, **kwargs)
+ except Exception as e:
+ raise ValueError(f"Unable to load {self.name} without `type_hint`: {e}")
+ # update type_hint if AutoModel load successfully
+ self.type_hint = component.__class__
+ else:
+ # determine load method
+ load_method = (
+ getattr(self.type_hint, "from_single_file")
+ if is_single_file
+ else getattr(self.type_hint, "from_pretrained")
+ )
+
+ if not is_single_file and DIFFUSERS_SDNQ_TRANSFORMERS and is_sdnq_available():
+ # Opt-in via DIFFUSERS_SDNQ_TRANSFORMERS: import sdnq once so it registers with transformers.
+ from ..quantizers.sdnq.sdnq_quantizer import _ensure_sdnq_registered
+
+ _ensure_sdnq_registered()
+
+ try:
+ component = load_method(pretrained_model_name_or_path, **load_kwargs, **kwargs)
+ except Exception as e:
+ raise ValueError(f"Unable to load {self.name} using load method: {e}")
+
+ self.pretrained_model_name_or_path = pretrained_model_name_or_path
+ for k, v in load_kwargs.items():
+ setattr(self, k, v)
+ component._diffusers_load_id = self.load_id
+
+ return component
+
+
+@dataclass
+class ConfigSpec:
+ """Specification for a pipeline configuration parameter."""
+
+ name: str
+ default: Any
+ description: str | None = None
+
+
+# ======================================================
+# InputParam and OutputParam templates
+# ======================================================
+
+INPUT_PARAM_TEMPLATES = {
+ "prompt": {
+ "type_hint": str,
+ "required": True,
+ "description": "The prompt or prompts to guide image generation.",
+ },
+ "negative_prompt": {
+ "type_hint": str,
+ "description": "The prompt or prompts not to guide the image generation.",
+ },
+ "max_sequence_length": {
+ "type_hint": int,
+ "default": 512,
+ "description": "Maximum sequence length for prompt encoding.",
+ },
+ "height": {
+ "type_hint": int,
+ "description": "The height in pixels of the generated image.",
+ },
+ "width": {
+ "type_hint": int,
+ "description": "The width in pixels of the generated image.",
+ },
+ "num_inference_steps": {
+ "type_hint": int,
+ "default": 50,
+ "description": "The number of denoising steps.",
+ },
+ "num_images_per_prompt": {
+ "type_hint": int,
+ "default": 1,
+ "description": "The number of images to generate per prompt.",
+ },
+ "generator": {
+ "type_hint": torch.Generator,
+ "description": "Torch generator for deterministic generation.",
+ },
+ "sigmas": {
+ "type_hint": list[float],
+ "description": "Custom sigmas for the denoising process.",
+ },
+ "strength": {
+ "type_hint": float,
+ "default": 0.9,
+ "description": "Strength for img2img/inpainting.",
+ },
+ "image": {
+ "type_hint": PIL.Image.Image | list[PIL.Image.Image],
+ "required": True,
+ "description": "Reference image(s) for denoising. Can be a single image or list of images.",
+ },
+ "latents": {
+ "type_hint": torch.Tensor,
+ "description": "Pre-generated noisy latents for image generation.",
+ },
+ "timesteps": {
+ "type_hint": torch.Tensor,
+ "description": "Timesteps for the denoising process.",
+ },
+ "output_type": {
+ "type_hint": str,
+ "default": "pil",
+ "description": "Output format: 'pil', 'np', 'pt'.",
+ },
+ "attention_kwargs": {
+ "type_hint": dict[str, Any],
+ "description": "Additional kwargs for attention processors.",
+ },
+ "denoiser_input_fields": {
+ "name": None,
+ "kwargs_type": "denoiser_input_fields",
+ "description": "conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.",
+ },
+ # inpainting
+ "mask_image": {
+ "type_hint": PIL.Image.Image,
+ "required": True,
+ "description": "Mask image for inpainting.",
+ },
+ "padding_mask_crop": {
+ "type_hint": int,
+ "description": "Padding for mask cropping in inpainting.",
+ },
+ # controlnet
+ "control_image": {
+ "type_hint": PIL.Image.Image,
+ "required": True,
+ "description": "Control image for ControlNet conditioning.",
+ },
+ "control_guidance_start": {
+ "type_hint": float,
+ "default": 0.0,
+ "description": "When to start applying ControlNet.",
+ },
+ "control_guidance_end": {
+ "type_hint": float,
+ "default": 1.0,
+ "description": "When to stop applying ControlNet.",
+ },
+ "controlnet_conditioning_scale": {
+ "type_hint": float,
+ "default": 1.0,
+ "description": "Scale for ControlNet conditioning.",
+ },
+ "layers": {
+ "type_hint": int,
+ "default": 4,
+ "description": "Number of layers to extract from the image",
+ },
+ # common intermediate inputs
+ "prompt_embeds": {
+ "type_hint": torch.Tensor,
+ "required": True,
+ "description": "text embeddings used to guide the image generation. Can be generated from text_encoder step.",
+ },
+ "prompt_embeds_mask": {
+ "type_hint": torch.Tensor,
+ "required": True,
+ "description": "mask for the text embeddings. Can be generated from text_encoder step.",
+ },
+ "negative_prompt_embeds": {
+ "type_hint": torch.Tensor,
+ "description": "negative text embeddings used to guide the image generation. Can be generated from text_encoder step.",
+ },
+ "negative_prompt_embeds_mask": {
+ "type_hint": torch.Tensor,
+ "description": "mask for the negative text embeddings. Can be generated from text_encoder step.",
+ },
+ "image_latents": {
+ "type_hint": torch.Tensor,
+ "required": True,
+ "description": "image latents used to guide the image generation. Can be generated from vae_encoder step.",
+ },
+ "batch_size": {
+ "type_hint": int,
+ "default": 1,
+ "description": "Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ },
+ "dtype": {
+ "type_hint": torch.dtype,
+ "default": torch.float32,
+ "description": "The dtype of the model inputs, can be generated in input step.",
+ },
+}
+
+OUTPUT_PARAM_TEMPLATES = {
+ "images": {
+ "type_hint": list[PIL.Image.Image],
+ "description": "Generated images.",
+ },
+ "videos": {
+ "type_hint": list[PIL.Image.Image],
+ "description": "The generated videos.",
+ },
+ "latents": {
+ "type_hint": torch.Tensor,
+ "description": "Denoised latents.",
+ },
+ # intermediate outputs
+ "prompt_embeds": {
+ "type_hint": torch.Tensor,
+ "kwargs_type": "denoiser_input_fields",
+ "description": "The prompt embeddings.",
+ },
+ "prompt_embeds_mask": {
+ "type_hint": torch.Tensor,
+ "kwargs_type": "denoiser_input_fields",
+ "description": "The encoder attention mask.",
+ },
+ "negative_prompt_embeds": {
+ "type_hint": torch.Tensor,
+ "kwargs_type": "denoiser_input_fields",
+ "description": "The negative prompt embeddings.",
+ },
+ "negative_prompt_embeds_mask": {
+ "type_hint": torch.Tensor,
+ "kwargs_type": "denoiser_input_fields",
+ "description": "The negative prompt embeddings mask.",
+ },
+ "image_latents": {
+ "type_hint": torch.Tensor,
+ "description": "The latent representation of the input image.",
+ },
+}
+
+
+@dataclass
+class InputParam:
+ """Specification for an input parameter."""
+
+ name: str = None
+ type_hint: Any = None
+ default: Any = None
+ required: bool = False
+ description: str = ""
+ kwargs_type: str = None
+ metadata: dict[str, Any] = None
+
+ def __repr__(self):
+ return f"<{self.name}: {'required' if self.required else 'optional'}, default={self.default}>"
+
+ @classmethod
+ def template(cls, template_name: str, note: str = None, **overrides) -> "InputParam":
+ """Get template for name if exists, otherwise raise ValueError."""
+ if template_name not in INPUT_PARAM_TEMPLATES:
+ raise ValueError(f"InputParam template for {template_name} not found")
+
+ template_kwargs = INPUT_PARAM_TEMPLATES[template_name].copy()
+
+ # Determine the actual param name:
+ # 1. From overrides if provided
+ # 2. From template if present
+ # 3. Fall back to template_name
+ name = overrides.pop("name", template_kwargs.pop("name", template_name))
+
+ if note and "description" in template_kwargs:
+ template_kwargs["description"] = f"{template_kwargs['description']} ({note})"
+
+ template_kwargs.update(overrides)
+ return cls(name=name, **template_kwargs)
+
+
+@dataclass
+class OutputParam:
+ """Specification for an output parameter."""
+
+ name: str
+ type_hint: Any = None
+ description: str = ""
+ kwargs_type: str = None
+ metadata: dict[str, Any] = None
+
+ def __repr__(self):
+ return (
+ f"<{self.name}: {self.type_hint.__name__ if hasattr(self.type_hint, '__name__') else str(self.type_hint)}>"
+ )
+
+ @classmethod
+ def template(cls, template_name: str, note: str = None, **overrides) -> "OutputParam":
+ """Get template for name if exists, otherwise raise ValueError."""
+ if template_name not in OUTPUT_PARAM_TEMPLATES:
+ raise ValueError(f"OutputParam template for {template_name} not found")
+
+ template_kwargs = OUTPUT_PARAM_TEMPLATES[template_name].copy()
+
+ # Determine the actual param name:
+ # 1. From overrides if provided
+ # 2. From template if present
+ # 3. Fall back to template_name
+ name = overrides.pop("name", template_kwargs.pop("name", template_name))
+
+ if note and "description" in template_kwargs:
+ template_kwargs["description"] = f"{template_kwargs['description']} ({note})"
+
+ template_kwargs.update(overrides)
+ return cls(name=name, **template_kwargs)
+
+
+def format_inputs_short(inputs):
+ """
+ Format input parameters into a string representation, with required params first followed by optional ones.
+
+ Args:
+ inputs: list of input parameters with 'required' and 'name' attributes, and 'default' for optional params
+
+ Returns:
+ str: Formatted string of input parameters
+
+ Example:
+ >>> inputs = [ ... InputParam(name="prompt", required=True), ... InputParam(name="image", required=True), ...
+ InputParam(name="guidance_scale", required=False, default=7.5), ... InputParam(name="num_inference_steps",
+ required=False, default=50) ... ] >>> format_inputs_short(inputs) 'prompt, image, guidance_scale=7.5,
+ num_inference_steps=50'
+ """
+ required_inputs = [param for param in inputs if param.required]
+ optional_inputs = [param for param in inputs if not param.required]
+
+ required_str = ", ".join(param.name for param in required_inputs)
+ optional_str = ", ".join(f"{param.name}={param.default}" for param in optional_inputs)
+
+ inputs_str = required_str
+ if optional_str:
+ inputs_str = f"{inputs_str}, {optional_str}" if required_str else optional_str
+
+ return inputs_str
+
+
+def format_params(params, header="Args", indent_level=4, max_line_length=115):
+ """Format a list of InputParam or OutputParam objects into a readable string representation.
+
+ Args:
+ params: list of InputParam or OutputParam objects to format
+ header: Header text to use (e.g. "Args" or "Returns")
+ indent_level: Number of spaces to indent each parameter line (default: 4)
+ max_line_length: Maximum length for each line before wrapping (default: 115)
+
+ Returns:
+ A formatted string representing all parameters
+ """
+ if not params:
+ return ""
+
+ base_indent = " " * indent_level
+ param_indent = " " * (indent_level + 4)
+ desc_indent = " " * (indent_level + 8)
+ formatted_params = []
+
+ def get_type_str(type_hint):
+ if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union:
+ type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)]
+ return " | ".join(type_strs)
+ return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint)
+
+ def wrap_text(text, indent, max_length):
+ """Wrap text while preserving markdown links and maintaining indentation."""
+ words = text.split()
+ lines = []
+ current_line = []
+ current_length = 0
+
+ for word in words:
+ word_length = len(word) + (1 if current_line else 0)
+
+ if current_line and current_length + word_length > max_length:
+ lines.append(" ".join(current_line))
+ current_line = [word]
+ current_length = len(word)
+ else:
+ current_line.append(word)
+ current_length += word_length
+
+ if current_line:
+ lines.append(" ".join(current_line))
+
+ return f"\n{indent}".join(lines)
+
+ # Add the header
+ formatted_params.append(f"{base_indent}{header}:")
+
+ for param in params:
+ # Format parameter name and type
+ type_str = get_type_str(param.type_hint) if param.type_hint != Any else ""
+ # YiYi Notes: remove this line if we remove kwargs_type
+ name = f"**{param.kwargs_type}" if param.name is None and param.kwargs_type is not None else param.name
+ param_str = f"{param_indent}{name} (`{type_str}`"
+
+ # Add optional tag and default value if parameter is an InputParam and optional
+ if hasattr(param, "required"):
+ if not param.required:
+ param_str += ", *optional*"
+ if param.default is not None:
+ param_str += f", defaults to {param.default}"
+ param_str += "):"
+
+ # Add description on a new line with additional indentation and wrapping
+ if param.description:
+ desc = re.sub(r"\[(.*?)\]\((https?://[^\s\)]+)\)", r"[\1](\2)", param.description)
+ wrapped_desc = wrap_text(desc, desc_indent, max_line_length)
+ param_str += f"\n{desc_indent}{wrapped_desc}"
+ else:
+ param_str += f"\n{desc_indent}TODO: Add description."
+
+ formatted_params.append(param_str)
+
+ return "\n".join(formatted_params)
+
+
+def format_input_params(input_params, indent_level=4, max_line_length=115):
+ """Format a list of InputParam objects into a readable string representation.
+
+ Args:
+ input_params: list of InputParam objects to format
+ indent_level: Number of spaces to indent each parameter line (default: 4)
+ max_line_length: Maximum length for each line before wrapping (default: 115)
+
+ Returns:
+ A formatted string representing all input parameters
+ """
+ return format_params(input_params, "Inputs", indent_level, max_line_length)
+
+
+def format_output_params(output_params, indent_level=4, max_line_length=115):
+ """Format a list of OutputParam objects into a readable string representation.
+
+ Args:
+ output_params: list of OutputParam objects to format
+ indent_level: Number of spaces to indent each parameter line (default: 4)
+ max_line_length: Maximum length for each line before wrapping (default: 115)
+
+ Returns:
+ A formatted string representing all output parameters
+ """
+ return format_params(output_params, "Outputs", indent_level, max_line_length)
+
+
+def format_params_markdown(params, header="Inputs"):
+ """Format a list of InputParam or OutputParam objects as a markdown bullet-point list.
+
+ Suitable for model cards rendered on Hugging Face Hub.
+
+ Args:
+ params: list of InputParam or OutputParam objects to format
+ header: Header text (e.g. "Inputs" or "Outputs")
+
+ Returns:
+ A formatted markdown string, or empty string if params is empty.
+ """
+ if not params:
+ return ""
+
+ def get_type_str(type_hint):
+ if isinstance(type_hint, UnionType) or get_origin(type_hint) is Union:
+ type_strs = [t.__name__ if hasattr(t, "__name__") else str(t) for t in get_args(type_hint)]
+ return " | ".join(type_strs)
+ return type_hint.__name__ if hasattr(type_hint, "__name__") else str(type_hint)
+
+ lines = [f"**{header}:**\n"] if header else []
+ for param in params:
+ type_str = get_type_str(param.type_hint) if param.type_hint != Any else ""
+ name = f"**{param.kwargs_type}" if param.name is None and param.kwargs_type is not None else param.name
+ param_str = f"- `{name}` (`{type_str}`"
+
+ if hasattr(param, "required") and not param.required:
+ param_str += ", *optional*"
+ if param.default is not None:
+ param_str += f", defaults to `{param.default}`"
+ param_str += ")"
+
+ desc = param.description if param.description else "No description provided"
+ param_str += f": {desc}"
+ lines.append(param_str)
+
+ return "\n".join(lines)
+
+
+def format_components(components, indent_level=4, max_line_length=115, add_empty_lines=True):
+ """Format a list of ComponentSpec objects into a readable string representation.
+
+ Args:
+ components: list of ComponentSpec objects to format
+ indent_level: Number of spaces to indent each component line (default: 4)
+ max_line_length: Maximum length for each line before wrapping (default: 115)
+ add_empty_lines: Whether to add empty lines between components (default: True)
+
+ Returns:
+ A formatted string representing all components
+ """
+ if not components:
+ return ""
+
+ base_indent = " " * indent_level
+ component_indent = " " * (indent_level + 4)
+ formatted_components = []
+
+ # Add the header
+ formatted_components.append(f"{base_indent}Components:")
+ if add_empty_lines:
+ formatted_components.append("")
+
+ # Add each component with optional empty lines between them
+ for i, component in enumerate(components):
+ # Get type name, handling special cases
+ type_name = (
+ component.type_hint.__name__ if hasattr(component.type_hint, "__name__") else str(component.type_hint)
+ )
+
+ component_desc = f"{component_indent}{component.name} (`{type_name}`)"
+ if component.description:
+ component_desc += f": {component.description}"
+
+ # Get the loading fields dynamically
+ loading_field_values = []
+ for field_name in component.loading_fields():
+ field_value = getattr(component, field_name)
+ if field_value:
+ loading_field_values.append(f"{field_name}={field_value}")
+
+ # Add loading field information if available
+ if loading_field_values:
+ component_desc += f" [{', '.join(loading_field_values)}]"
+
+ formatted_components.append(component_desc)
+
+ # Add an empty line after each component except the last one
+ if add_empty_lines and i < len(components) - 1:
+ formatted_components.append("")
+
+ return "\n".join(formatted_components)
+
+
+def format_configs(configs, indent_level=4, max_line_length=115, add_empty_lines=True):
+ """Format a list of ConfigSpec objects into a readable string representation.
+
+ Args:
+ configs: list of ConfigSpec objects to format
+ indent_level: Number of spaces to indent each config line (default: 4)
+ max_line_length: Maximum length for each line before wrapping (default: 115)
+ add_empty_lines: Whether to add empty lines between configs (default: True)
+
+ Returns:
+ A formatted string representing all configs
+ """
+ if not configs:
+ return ""
+
+ base_indent = " " * indent_level
+ config_indent = " " * (indent_level + 4)
+ formatted_configs = []
+
+ # Add the header
+ formatted_configs.append(f"{base_indent}Configs:")
+ if add_empty_lines:
+ formatted_configs.append("")
+
+ # Add each config with optional empty lines between them
+ for i, config in enumerate(configs):
+ config_desc = f"{config_indent}{config.name} (default: {config.default})"
+ if config.description:
+ config_desc += f": {config.description}"
+ formatted_configs.append(config_desc)
+
+ # Add an empty line after each config except the last one
+ if add_empty_lines and i < len(configs) - 1:
+ formatted_configs.append("")
+
+ return "\n".join(formatted_configs)
+
+
+def format_workflow(workflow_map):
+ """Format a workflow map into a readable string representation.
+
+ Args:
+ workflow_map: Dictionary mapping workflow names to trigger inputs
+
+ Returns:
+ A formatted string representing all workflows
+ """
+ if workflow_map is None:
+ return ""
+
+ lines = ["Supported workflows:"]
+ for workflow_name, trigger_inputs in workflow_map.items():
+ required_inputs = [k for k, v in trigger_inputs.items() if v]
+ if required_inputs:
+ inputs_str = ", ".join(f"`{t}`" for t in required_inputs)
+ lines.append(f" - `{workflow_name}`: requires {inputs_str}")
+ else:
+ lines.append(f" - `{workflow_name}`: default (no additional inputs required)")
+
+ return "\n".join(lines)
+
+
+def make_doc_string(
+ inputs,
+ outputs,
+ description="",
+ class_name=None,
+ expected_components=None,
+ expected_configs=None,
+):
+ """
+ Generates a formatted documentation string describing the pipeline block's parameters and structure.
+
+ Args:
+ inputs: list of input parameters
+ outputs: list of output parameters
+ description (str, *optional*): Description of the block
+ class_name (str, *optional*): Name of the class to include in the documentation
+ expected_components (list[ComponentSpec], *optional*): list of expected components
+ expected_configs (list[ConfigSpec], *optional*): list of expected configurations
+
+ Returns:
+ str: A formatted string containing information about components, configs, call parameters,
+ intermediate inputs/outputs, and final outputs.
+ """
+ output = ""
+
+ # Add class name if provided
+ if class_name:
+ output += f"class {class_name}\n\n"
+
+ # Add description
+ if description:
+ desc_lines = description.strip().split("\n")
+ aligned_desc = "\n".join(" " + line.rstrip() for line in desc_lines)
+ output += aligned_desc + "\n\n"
+
+ # Add components section if provided
+ if expected_components and len(expected_components) > 0:
+ components_str = format_components(expected_components, indent_level=2, add_empty_lines=False)
+ output += components_str + "\n\n"
+
+ # Add configs section if provided
+ if expected_configs and len(expected_configs) > 0:
+ configs_str = format_configs(expected_configs, indent_level=2, add_empty_lines=False)
+ output += configs_str + "\n\n"
+
+ # Add inputs section
+ output += format_input_params(inputs, indent_level=2)
+
+ # Add outputs section
+ output += "\n\n"
+ output += format_output_params(outputs, indent_level=2)
+
+ return output
+
+
+def _validate_requirements(reqs):
+ if reqs is None:
+ normalized_reqs = {}
+ else:
+ if not isinstance(reqs, dict):
+ raise ValueError(
+ "Requirements must be provided as a dictionary mapping package names to version specifiers."
+ )
+ normalized_reqs = _normalize_requirements(reqs)
+
+ if not normalized_reqs:
+ return {}
+
+ final: dict[str, str] = {}
+ for req, specified_ver in normalized_reqs.items():
+ req_available, req_actual_ver = _is_package_available(req)
+ if not req_available:
+ logger.warning(f"{req} was specified in the requirements but wasn't found in the current environment.")
+
+ if specified_ver:
+ try:
+ specifier = SpecifierSet(specified_ver)
+ except InvalidSpecifier as err:
+ raise ValueError(f"Requirement specifier '{specified_ver}' for {req} is invalid.") from err
+
+ if req_actual_ver == "N/A":
+ logger.warning(
+ f"Version of {req} could not be determined to validate requirement '{specified_ver}'. Things might work unexpected."
+ )
+ elif not specifier.contains(req_actual_ver, prereleases=True):
+ logger.warning(
+ f"{req} requirement '{specified_ver}' is not satisfied by the installed version {req_actual_ver}. Things might work unexpected."
+ )
+
+ final[req] = specified_ver
+
+ return final
+
+
+def _normalize_requirements(reqs):
+ if not reqs:
+ return {}
+
+ normalized: "OrderedDict[str, str]" = OrderedDict()
+
+ def _accumulate(mapping: dict[str, Any]):
+ for pkg, spec in mapping.items():
+ if isinstance(spec, dict):
+ # This is recursive because blocks are composable. This way, we can merge requirements
+ # from multiple blocks.
+ _accumulate(spec)
+ continue
+
+ pkg_name = str(pkg).strip()
+ if not pkg_name:
+ raise ValueError("Requirement package name cannot be empty.")
+
+ spec_str = "" if spec is None else str(spec).strip()
+ if spec_str and not spec_str.startswith(("<", ">", "=", "!", "~")):
+ spec_str = f"=={spec_str}"
+
+ existing_spec = normalized.get(pkg_name)
+ if existing_spec is not None:
+ if not existing_spec and spec_str:
+ normalized[pkg_name] = spec_str
+ elif existing_spec and spec_str and existing_spec != spec_str:
+ try:
+ combined_spec = SpecifierSet(",".join(filter(None, [existing_spec, spec_str])))
+ except InvalidSpecifier:
+ logger.warning(
+ f"Conflicting requirements for '{pkg_name}' detected: '{existing_spec}' vs '{spec_str}'. Keeping '{existing_spec}'."
+ )
+ else:
+ normalized[pkg_name] = str(combined_spec)
+ continue
+
+ normalized[pkg_name] = spec_str
+
+ _accumulate(reqs)
+
+ return normalized
+
+
+def combine_inputs(*named_input_lists: list[tuple[str, list[InputParam]]]) -> list[InputParam]:
+ """
+ Combines multiple lists of InputParam objects from different blocks. For duplicate inputs, updates only if current
+ default value is None and new default value is not None. Warns if multiple non-None default values exist for the
+ same input.
+
+ Args:
+ named_input_lists: List of tuples containing (block_name, input_param_list) pairs
+
+ Returns:
+ List[InputParam]: Combined list of unique InputParam objects
+ """
+ combined_dict = {} # name -> InputParam
+ value_sources = {} # name -> block_name
+
+ for block_name, inputs in named_input_lists:
+ for input_param in inputs:
+ if input_param.name is None and input_param.kwargs_type is not None:
+ input_name = "*_" + input_param.kwargs_type
+ else:
+ input_name = input_param.name
+ if input_name in combined_dict:
+ current_param = combined_dict[input_name]
+ if (
+ current_param.default is not None
+ and input_param.default is not None
+ and current_param.default != input_param.default
+ ):
+ warnings.warn(
+ f"Multiple different default values found for input '{input_name}': "
+ f"{current_param.default} (from block '{value_sources[input_name]}') and "
+ f"{input_param.default} (from block '{block_name}'). Using {current_param.default}."
+ )
+ if current_param.default is None and input_param.default is not None:
+ combined_dict[input_name] = input_param
+ value_sources[input_name] = block_name
+ else:
+ combined_dict[input_name] = input_param
+ value_sources[input_name] = block_name
+
+ return list(combined_dict.values())
+
+
+def combine_outputs(*named_output_lists: list[tuple[str, list[OutputParam]]]) -> list[OutputParam]:
+ """
+ Combines multiple lists of OutputParam objects from different blocks. For duplicate outputs, keeps the first
+ occurrence of each output name.
+
+ Args:
+ named_output_lists: List of tuples containing (block_name, output_param_list) pairs
+
+ Returns:
+ List[OutputParam]: Combined list of unique OutputParam objects
+ """
+ combined_dict = {} # name -> OutputParam
+
+ for block_name, outputs in named_output_lists:
+ for output_param in outputs:
+ if (output_param.name not in combined_dict) or (
+ combined_dict[output_param.name].kwargs_type is None and output_param.kwargs_type is not None
+ ):
+ combined_dict[output_param.name] = output_param
+
+ return list(combined_dict.values())
+
+
+def generate_modular_model_card_content(blocks) -> dict[str, Any]:
+ """
+ Generate model card content for a modular pipeline.
+
+ This function creates a comprehensive model card with descriptions of the pipeline's architecture, components,
+ configurations, inputs, and outputs.
+
+ Args:
+ blocks: The pipeline's blocks object containing all pipeline specifications
+
+ Returns:
+ Dict[str, Any]: A dictionary containing formatted content sections:
+ - pipeline_name: Name of the pipeline
+ - model_description: Overall description with pipeline type
+ - blocks_description: Detailed architecture of blocks
+ - components_description: List of required components
+ - configs_section: Configuration parameters section
+ - io_specification_section: Input/Output specification (per-workflow or unified)
+ - trigger_inputs_section: Conditional execution information
+ - tags: List of relevant tags for the model card
+ """
+ blocks_class_name = blocks.__class__.__name__
+ pipeline_name = blocks_class_name.replace("Blocks", " Pipeline")
+ description = getattr(blocks, "description", "A modular diffusion pipeline.")
+
+ # generate blocks architecture description
+ blocks_desc_parts = []
+ sub_blocks = getattr(blocks, "sub_blocks", None) or {}
+ if sub_blocks:
+ for i, (name, block) in enumerate(sub_blocks.items()):
+ block_class = block.__class__.__name__
+ block_desc = block.description.split("\n")[0] if getattr(block, "description", "") else ""
+ blocks_desc_parts.append(f"{i + 1}. **{name}** (`{block_class}`)")
+ if block_desc:
+ blocks_desc_parts.append(f" - {block_desc}")
+
+ blocks_description = "\n".join(blocks_desc_parts) if blocks_desc_parts else "No blocks defined."
+
+ components = getattr(blocks, "expected_components", [])
+ if components:
+ components_str = format_components(components, indent_level=0, add_empty_lines=False)
+ # remove the "Components:" header since template has its own
+ components_description = components_str.replace("Components:\n", "").strip()
+ if components_description:
+ # Convert to enumerated list
+ lines = [line.strip() for line in components_description.split("\n") if line.strip()]
+ enumerated_lines = [f"{i + 1}. {line}" for i, line in enumerate(lines)]
+ components_description = "\n".join(enumerated_lines)
+ else:
+ components_description = "No specific components required."
+ else:
+ components_description = "No specific components required. Components can be loaded dynamically."
+
+ configs = getattr(blocks, "expected_configs", [])
+ configs_section = ""
+ if configs:
+ configs_str = format_configs(configs, indent_level=0, add_empty_lines=False)
+ configs_description = configs_str.replace("Configs:\n", "").strip()
+ if configs_description:
+ configs_section = f"\n\n## Configuration Parameters\n\n{configs_description}"
+
+ # Branch on whether workflows are defined
+ has_workflows = getattr(blocks, "_workflow_map", None) is not None
+
+ if has_workflows:
+ workflow_map = blocks._workflow_map
+ parts = []
+
+ # If blocks overrides outputs (e.g. to return just "images" instead of all intermediates),
+ # use that as the shared output for all workflows
+ blocks_outputs = blocks.outputs
+ blocks_intermediate = getattr(blocks, "intermediate_outputs", None)
+ shared_outputs = (
+ blocks_outputs if blocks_intermediate is not None and blocks_outputs != blocks_intermediate else None
+ )
+
+ parts.append("## Workflow Input Specification\n")
+
+ # Per-workflow details: show trigger inputs with full param descriptions
+ for wf_name, trigger_inputs in workflow_map.items():
+ trigger_input_names = set(trigger_inputs.keys())
+ try:
+ workflow_blocks = blocks.get_workflow(wf_name)
+ except Exception:
+ parts.append(f"\n{wf_name} \n")
+ parts.append("*Could not resolve workflow blocks.*\n")
+ parts.append(" \n")
+ continue
+
+ wf_inputs = workflow_blocks.inputs
+ # Show only trigger inputs with full parameter descriptions
+ trigger_params = [p for p in wf_inputs if p.name in trigger_input_names]
+
+ parts.append(f"\n{wf_name} \n")
+
+ inputs_str = format_params_markdown(trigger_params, header=None)
+ parts.append(inputs_str if inputs_str else "No additional inputs required.")
+ parts.append("")
+
+ parts.append(" \n")
+
+ # Common Inputs & Outputs section (like non-workflow pipelines)
+ all_inputs = blocks.inputs
+ all_outputs = shared_outputs if shared_outputs is not None else blocks.outputs
+
+ inputs_str = format_params_markdown(all_inputs, "Inputs")
+ outputs_str = format_params_markdown(all_outputs, "Outputs")
+ inputs_description = inputs_str if inputs_str else "No specific inputs defined."
+ outputs_description = outputs_str if outputs_str else "Standard pipeline outputs."
+
+ parts.append(f"\n## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}")
+
+ io_specification_section = "\n".join(parts)
+ # Suppress trigger_inputs_section when workflows are shown (it's redundant)
+ trigger_inputs_section = ""
+ else:
+ # Unified I/O section (original behavior)
+ inputs = blocks.inputs
+ outputs = blocks.outputs
+ inputs_str = format_params_markdown(inputs, "Inputs")
+ outputs_str = format_params_markdown(outputs, "Outputs")
+ inputs_description = inputs_str if inputs_str else "No specific inputs defined."
+ outputs_description = outputs_str if outputs_str else "Standard pipeline outputs."
+ io_specification_section = f"## Input/Output Specification\n\n{inputs_description}\n\n{outputs_description}"
+
+ trigger_inputs_section = ""
+ if hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
+ trigger_inputs_list = sorted([t for t in blocks.trigger_inputs if t is not None])
+ if trigger_inputs_list:
+ trigger_inputs_str = ", ".join(f"`{t}`" for t in trigger_inputs_list)
+ trigger_inputs_section = f"""
+### Conditional Execution
+
+This pipeline contains blocks that are selected at runtime based on inputs:
+- **Trigger Inputs**: {trigger_inputs_str}
+"""
+
+ # generate tags based on pipeline characteristics
+ tags = ["modular-diffusers", "diffusers"]
+
+ if hasattr(blocks, "model_name") and blocks.model_name:
+ tags.append(blocks.model_name)
+
+ if has_workflows:
+ # Derive tags from workflow names
+ workflow_names = set(blocks._workflow_map.keys())
+ if any("inpainting" in wf for wf in workflow_names):
+ tags.append("inpainting")
+ if any("image2image" in wf for wf in workflow_names):
+ tags.append("image-to-image")
+ if any("controlnet" in wf for wf in workflow_names):
+ tags.append("controlnet")
+ if any("text2image" in wf for wf in workflow_names):
+ tags.append("text-to-image")
+ elif hasattr(blocks, "trigger_inputs") and blocks.trigger_inputs:
+ triggers = blocks.trigger_inputs
+ if any(t in triggers for t in ["mask", "mask_image"]):
+ tags.append("inpainting")
+ if any(t in triggers for t in ["image", "image_latents"]):
+ tags.append("image-to-image")
+ if any(t in triggers for t in ["control_image", "controlnet_cond"]):
+ tags.append("controlnet")
+ if not any(t in triggers for t in ["image", "mask", "image_latents", "mask_image"]):
+ tags.append("text-to-image")
+ else:
+ tags.append("text-to-image")
+
+ block_count = len(blocks.sub_blocks)
+ model_description = f"""This is a modular diffusion pipeline built with 🧨 Diffusers' modular pipeline framework.
+
+**Pipeline Type**: {blocks_class_name}
+
+**Description**: {description}
+
+This pipeline uses a {block_count}-block architecture that can be customized and extended."""
+
+ return {
+ "pipeline_name": pipeline_name,
+ "model_description": model_description,
+ "blocks_description": blocks_description,
+ "components_description": components_description,
+ "configs_section": configs_section,
+ "io_specification_section": io_specification_section,
+ "trigger_inputs_section": trigger_inputs_section,
+ "tags": tags,
+ }
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a30a7276f424d4715362d5cbfbb718d3d608ce8e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/__init__.py
@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_qwenimage"] = ["QwenImageAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_edit"] = ["QwenImageEditAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_edit_plus"] = ["QwenImageEditPlusAutoBlocks"]
+ _import_structure["modular_blocks_qwenimage_layered"] = ["QwenImageLayeredAutoBlocks"]
+ _import_structure["modular_pipeline"] = [
+ "QwenImageEditModularPipeline",
+ "QwenImageEditPlusModularPipeline",
+ "QwenImageLayeredModularPipeline",
+ "QwenImageModularPipeline",
+ ]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_qwenimage import QwenImageAutoBlocks
+ from .modular_blocks_qwenimage_edit import QwenImageEditAutoBlocks
+ from .modular_blocks_qwenimage_edit_plus import QwenImageEditPlusAutoBlocks
+ from .modular_blocks_qwenimage_layered import QwenImageLayeredAutoBlocks
+ from .modular_pipeline import (
+ QwenImageEditModularPipeline,
+ QwenImageEditPlusModularPipeline,
+ QwenImageLayeredModularPipeline,
+ QwenImageModularPipeline,
+ )
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..455c7717dbaf6843d384ec2ee4f11513c8d7c3f4
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/before_denoise.py
@@ -0,0 +1,1330 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import numpy as np
+import torch
+
+from ...models import QwenImageControlNetModel, QwenImageMultiControlNetModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils.torch_utils import randn_tensor, unwrap_module
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier
+
+
+# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# modified from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps
+def get_timesteps(scheduler, num_inference_steps, strength):
+ # get the original timestep using init_timestep
+ init_timestep = min(num_inference_steps * strength, num_inference_steps)
+
+ t_start = int(max(num_inference_steps - init_timestep, 0))
+ timesteps = scheduler.timesteps[t_start * scheduler.order :]
+ if hasattr(scheduler, "set_begin_index"):
+ scheduler.set_begin_index(t_start * scheduler.order)
+
+ return timesteps, num_inference_steps - t_start
+
+
+# ====================
+# 1. PREPARE LATENTS
+# ====================
+
+
+# auto_docstring
+class QwenImagePrepareLatentsStep(ModularPipelineBlocks):
+ """
+ Prepare initial random noise for the generation process
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ dtype (`dtype`, *optional*, defaults to torch.float32):
+ The dtype of the model inputs, can be generated in input step.
+
+ Outputs:
+ height (`int`):
+ if not set, updated to default value
+ width (`int`):
+ if not set, updated to default value
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Prepare initial random noise for the generation process"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("generator"),
+ InputParam.template("batch_size"),
+ InputParam.template("dtype"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="height", type_hint=int, description="if not set, updated to default value"),
+ OutputParam(name="width", type_hint=int, description="if not set, updated to default value"),
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ height=block_state.height,
+ width=block_state.width,
+ vae_scale_factor=components.vae_scale_factor,
+ )
+
+ device = components._execution_device
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+ # we can update the height and width here since it's used to generate the initial
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ latent_height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+ latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+
+ shape = (batch_size, components.num_channels_latents, 1, latent_height, latent_width)
+ if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+ if block_state.latents is None:
+ block_state.latents = randn_tensor(
+ shape, generator=block_state.generator, device=device, dtype=block_state.dtype
+ )
+ block_state.latents = components.pachifier.pack_latents(block_state.latents)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredPrepareLatentsStep(ModularPipelineBlocks):
+ """
+ Prepare initial random noise (B, layers+1, C, H, W) for the generation process
+
+ Components:
+ pachifier (`QwenImageLayeredPachifier`)
+
+ Inputs:
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ layers (`int`, *optional*, defaults to 4):
+ Number of layers to extract from the image
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ dtype (`dtype`, *optional*, defaults to torch.float32):
+ The dtype of the model inputs, can be generated in input step.
+
+ Outputs:
+ height (`int`):
+ if not set, updated to default value
+ width (`int`):
+ if not set, updated to default value
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Prepare initial random noise (B, layers+1, C, H, W) for the generation process"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("latents"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("layers"),
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("generator"),
+ InputParam.template("batch_size"),
+ InputParam.template("dtype"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="height", type_hint=int, description="if not set, updated to default value"),
+ OutputParam(name="width", type_hint=int, description="if not set, updated to default value"),
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ height=block_state.height,
+ width=block_state.width,
+ vae_scale_factor=components.vae_scale_factor,
+ )
+
+ device = components._execution_device
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+ # we can update the height and width here since it's used to generate the initial
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ latent_height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+ latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+
+ shape = (batch_size, block_state.layers + 1, components.num_channels_latents, latent_height, latent_width)
+ if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+ if block_state.latents is None:
+ block_state.latents = randn_tensor(
+ shape, generator=block_state.generator, device=device, dtype=block_state.dtype
+ )
+ block_state.latents = components.pachifier.pack_latents(block_state.latents)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImagePrepareLatentsWithStrengthStep(ModularPipelineBlocks):
+ """
+ Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps,
+ prepare_latents. Both noise and image latents should alreadybe patchified.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ latents (`Tensor`):
+ The initial random noised, can be generated in prepare latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be
+ generated from vae encoder and updated in input step.)
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+
+ Outputs:
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ latents (`Tensor`):
+ The scaled noisy latents to use for inpainting/image-to-image denoising.
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps, prepare_latents. Both noise and image latents should alreadybe patchified."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial random noised, can be generated in prepare latent step.",
+ ),
+ InputParam.template("image_latents", note="Can be generated from vae encoder and updated in input step."),
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="initial_noise",
+ type_hint=torch.Tensor,
+ description="The initial random noised used for inpainting denoising.",
+ ),
+ OutputParam(
+ name="latents",
+ type_hint=torch.Tensor,
+ description="The scaled noisy latents to use for inpainting/image-to-image denoising.",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(image_latents, latents):
+ if image_latents.shape[0] != latents.shape[0]:
+ raise ValueError(
+ f"`image_latents` must have have same batch size as `latents`, but got {image_latents.shape[0]} and {latents.shape[0]}"
+ )
+
+ if image_latents.ndim != 3:
+ raise ValueError(f"`image_latents` must have 3 dimensions (patchified), but got {image_latents.ndim}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ image_latents=block_state.image_latents,
+ latents=block_state.latents,
+ )
+
+ # prepare latent timestep
+ latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
+
+ # make copy of initial_noise
+ block_state.initial_noise = block_state.latents
+
+ # scale noise
+ block_state.latents = components.scheduler.scale_noise(
+ block_state.image_latents, latent_timestep, block_state.latents
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageCreateMaskLatentsStep(ModularPipelineBlocks):
+ """
+ Step that creates mask latents from preprocessed mask_image by interpolating to latent space.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ processed_mask_image (`Tensor`):
+ The processed mask to use for the inpainting process.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ dtype (`dtype`, *optional*, defaults to torch.float32):
+ The dtype of the model inputs, can be generated in input step.
+
+ Outputs:
+ mask (`Tensor`):
+ The mask to use for the inpainting process.
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that creates mask latents from preprocessed mask_image by interpolating to latent space."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="processed_mask_image",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The processed mask to use for the inpainting process.",
+ ),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("dtype"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+
+ height_latents = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+ width_latents = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+
+ block_state.mask = torch.nn.functional.interpolate(
+ block_state.processed_mask_image,
+ size=(height_latents, width_latents),
+ )
+
+ block_state.mask = block_state.mask.unsqueeze(2)
+ block_state.mask = block_state.mask.repeat(1, components.num_channels_latents, 1, 1, 1)
+ block_state.mask = block_state.mask.to(device=device, dtype=block_state.dtype)
+
+ block_state.mask = components.pachifier.pack_latents(block_state.mask)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# ====================
+# 2. SET TIMESTEPS
+# ====================
+
+
+# auto_docstring
+class QwenImageSetTimestepsStep(ModularPipelineBlocks):
+ """
+ Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`):
+ The initial random noised latents for the denoising process. Can be generated in prepare latents step.
+
+ Outputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial random noised latents for the denoising process. Can be generated in prepare latents step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process"
+ ),
+ ]
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ sigmas = (
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+ if block_state.sigmas is None
+ else block_state.sigmas
+ )
+
+ mu = calculate_shift(
+ image_seq_len=block_state.latents.shape[1],
+ base_seq_len=components.scheduler.config.get("base_image_seq_len", 256),
+ max_seq_len=components.scheduler.config.get("max_image_seq_len", 4096),
+ base_shift=components.scheduler.config.get("base_shift", 0.5),
+ max_shift=components.scheduler.config.get("max_shift", 1.15),
+ )
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ scheduler=components.scheduler,
+ num_inference_steps=block_state.num_inference_steps,
+ device=device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+
+ components.scheduler.set_begin_index(0)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredSetTimestepsStep(ModularPipelineBlocks):
+ """
+ Set timesteps step for QwenImage Layered with custom mu calculation based on image_latents.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process.
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Set timesteps step for QwenImage Layered with custom mu calculation based on image_latents."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ InputParam.template("image_latents"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process."
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ # Layered-specific mu calculation
+ base_seqlen = 256 * 256 / 16 / 16 # = 256
+ mu = (block_state.image_latents.shape[1] / base_seqlen) ** 0.5
+
+ # Default sigmas if not provided
+ sigmas = (
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+ if block_state.sigmas is None
+ else block_state.sigmas
+ )
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+
+ components.scheduler.set_begin_index(0)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageSetTimestepsWithStrengthStep(ModularPipelineBlocks):
+ """
+ Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare
+ latents step.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ latents (`Tensor`):
+ The latents to use for the denoising process. Can be generated in prepare latents step.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+
+ Outputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process.
+ num_inference_steps (`int`):
+ The number of denoising steps to perform at inference time. Updated based on strength.
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare latents step."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_inference_steps"),
+ InputParam.template("sigmas"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to use for the denoising process. Can be generated in prepare latents step.",
+ ),
+ InputParam.template("strength", default=0.9),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="timesteps",
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process.",
+ ),
+ OutputParam(
+ name="num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time. Updated based on strength.",
+ ),
+ ]
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ sigmas = (
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+ if block_state.sigmas is None
+ else block_state.sigmas
+ )
+
+ mu = calculate_shift(
+ image_seq_len=block_state.latents.shape[1],
+ base_seq_len=components.scheduler.config.get("base_image_seq_len", 256),
+ max_seq_len=components.scheduler.config.get("max_image_seq_len", 4096),
+ base_shift=components.scheduler.config.get("base_shift", 0.5),
+ max_shift=components.scheduler.config.get("max_shift", 1.15),
+ )
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ scheduler=components.scheduler,
+ num_inference_steps=block_state.num_inference_steps,
+ device=device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+
+ block_state.timesteps, block_state.num_inference_steps = get_timesteps(
+ scheduler=components.scheduler,
+ num_inference_steps=block_state.num_inference_steps,
+ strength=block_state.strength,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# ====================
+# 3. OTHER INPUTS FOR DENOISER
+# ====================
+
+## RoPE inputs for denoiser
+
+
+# auto_docstring
+class QwenImageRoPEInputsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step
+
+ Inputs:
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+
+ Outputs:
+ img_shapes (`list`):
+ The shapes of the images latents, used for RoPE calculation
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("batch_size"),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("prompt_embeds_mask"),
+ InputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="img_shapes",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[list[tuple[int, int, int]]],
+ description="The shapes of the images latents, used for RoPE calculation",
+ ),
+ ]
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.img_shapes = [
+ [
+ (
+ 1,
+ block_state.height // components.vae_scale_factor // 2,
+ block_state.width // components.vae_scale_factor // 2,
+ )
+ ]
+ ] * block_state.batch_size
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditRoPEInputsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after
+ prepare_latents step
+
+ Inputs:
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ image_height (`int`):
+ The height of the reference image. Can be generated in input step.
+ image_width (`int`):
+ The width of the reference image. Can be generated in input step.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+
+ Outputs:
+ img_shapes (`list`):
+ The shapes of the images latents, used for RoPE calculation
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit. Should be placed after prepare_latents step"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("batch_size"),
+ InputParam(
+ name="image_height",
+ required=True,
+ type_hint=int,
+ description="The height of the reference image. Can be generated in input step.",
+ ),
+ InputParam(
+ name="image_width",
+ required=True,
+ type_hint=int,
+ description="The width of the reference image. Can be generated in input step.",
+ ),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("prompt_embeds_mask"),
+ InputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="img_shapes",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[list[tuple[int, int, int]]],
+ description="The shapes of the images latents, used for RoPE calculation",
+ ),
+ ]
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # for edit, image size can be different from the target size (height/width)
+ block_state.img_shapes = [
+ [
+ (
+ 1,
+ block_state.height // components.vae_scale_factor // 2,
+ block_state.width // components.vae_scale_factor // 2,
+ ),
+ (
+ 1,
+ block_state.image_height // components.vae_scale_factor // 2,
+ block_state.image_width // components.vae_scale_factor // 2,
+ ),
+ ]
+ ] * block_state.batch_size
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditPlusRoPEInputsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit Plus.
+ Unlike Edit, Edit Plus handles lists of image_height/image_width for multiple reference images. Should be placed
+ after prepare_latents step.
+
+ Inputs:
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ image_height (`list`):
+ The heights of the reference images. Can be generated in input step.
+ image_width (`list`):
+ The widths of the reference images. Can be generated in input step.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+
+ Outputs:
+ img_shapes (`list`):
+ The shapes of the image latents, used for RoPE calculation
+ txt_seq_lens (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation
+ negative_txt_seq_lens (`list`):
+ The sequence lengths of the negative prompt embeds, used for RoPE calculation
+ """
+
+ model_name = "qwenimage-edit-plus"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that prepares the RoPE inputs for denoising process. This is used in QwenImage Edit Plus.\n"
+ "Unlike Edit, Edit Plus handles lists of image_height/image_width for multiple reference images.\n"
+ "Should be placed after prepare_latents step."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("batch_size"),
+ InputParam(
+ name="image_height",
+ required=True,
+ type_hint=list[int],
+ description="The heights of the reference images. Can be generated in input step.",
+ ),
+ InputParam(
+ name="image_width",
+ required=True,
+ type_hint=list[int],
+ description="The widths of the reference images. Can be generated in input step.",
+ ),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("prompt_embeds_mask"),
+ InputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="img_shapes",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[list[tuple[int, int, int]]],
+ description="The shapes of the image latents, used for RoPE calculation",
+ ),
+ OutputParam(
+ name="txt_seq_lens",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the prompt embeds, used for RoPE calculation",
+ ),
+ OutputParam(
+ name="negative_txt_seq_lens",
+ kwargs_type="denoiser_input_fields",
+ type_hint=list[int],
+ description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
+ ),
+ ]
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae_scale_factor = components.vae_scale_factor
+
+ # Edit Plus: image_height and image_width are lists
+ block_state.img_shapes = [
+ [
+ (1, block_state.height // vae_scale_factor // 2, block_state.width // vae_scale_factor // 2),
+ *[
+ (1, img_height // vae_scale_factor // 2, img_width // vae_scale_factor // 2)
+ for img_height, img_width in zip(block_state.image_height, block_state.image_width)
+ ],
+ ]
+ ] * block_state.batch_size
+
+ block_state.txt_seq_lens = (
+ block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None
+ )
+ block_state.negative_txt_seq_lens = (
+ block_state.negative_prompt_embeds_mask.sum(dim=1).tolist()
+ if block_state.negative_prompt_embeds_mask is not None
+ else None
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredRoPEInputsStep(ModularPipelineBlocks):
+ """
+ Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step
+
+ Inputs:
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ layers (`int`, *optional*, defaults to 4):
+ Number of layers to extract from the image
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+
+ Outputs:
+ img_shapes (`list`):
+ The shapes of the image latents, used for RoPE calculation
+ txt_seq_lens (`list`):
+ The sequence lengths of the prompt embeds, used for RoPE calculation
+ negative_txt_seq_lens (`list`):
+ The sequence lengths of the negative prompt embeds, used for RoPE calculation
+ additional_t_cond (`Tensor`):
+ The additional t cond, used for RoPE calculation
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that prepares the RoPE inputs for the denoising process. Should be place after prepare_latents step"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("batch_size"),
+ InputParam.template("layers"),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("prompt_embeds_mask"),
+ InputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="img_shapes",
+ type_hint=list[list[tuple[int, int, int]]],
+ kwargs_type="denoiser_input_fields",
+ description="The shapes of the image latents, used for RoPE calculation",
+ ),
+ OutputParam(
+ name="txt_seq_lens",
+ type_hint=list[int],
+ kwargs_type="denoiser_input_fields",
+ description="The sequence lengths of the prompt embeds, used for RoPE calculation",
+ ),
+ OutputParam(
+ name="negative_txt_seq_lens",
+ type_hint=list[int],
+ kwargs_type="denoiser_input_fields",
+ description="The sequence lengths of the negative prompt embeds, used for RoPE calculation",
+ ),
+ OutputParam(
+ name="additional_t_cond",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The additional t cond, used for RoPE calculation",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ # All shapes are the same for Layered
+ shape = (
+ 1,
+ block_state.height // components.vae_scale_factor // 2,
+ block_state.width // components.vae_scale_factor // 2,
+ )
+
+ # layers+1 output shapes + 1 condition shape (all same)
+ block_state.img_shapes = [[shape] * (block_state.layers + 2)] * block_state.batch_size
+
+ # txt_seq_lens
+ block_state.txt_seq_lens = (
+ block_state.prompt_embeds_mask.sum(dim=1).tolist() if block_state.prompt_embeds_mask is not None else None
+ )
+ block_state.negative_txt_seq_lens = (
+ block_state.negative_prompt_embeds_mask.sum(dim=1).tolist()
+ if block_state.negative_prompt_embeds_mask is not None
+ else None
+ )
+
+ block_state.additional_t_cond = torch.tensor([0] * block_state.batch_size).to(device=device, dtype=torch.long)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+## ControlNet inputs for denoiser
+
+
+# auto_docstring
+class QwenImageControlNetBeforeDenoiserStep(ModularPipelineBlocks):
+ """
+ step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step.
+
+ Components:
+ controlnet (`QwenImageControlNetModel`)
+
+ Inputs:
+ control_guidance_start (`float`, *optional*, defaults to 0.0):
+ When to start applying ControlNet.
+ control_guidance_end (`float`, *optional*, defaults to 1.0):
+ When to stop applying ControlNet.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning.
+ control_image_latents (`Tensor`):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+
+ Outputs:
+ controlnet_keep (`list`):
+ The controlnet keep values
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("controlnet", QwenImageControlNetModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "step that prepare inputs for controlnet. Insert before the Denoise Step, after set_timesteps step."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("control_guidance_start"),
+ InputParam.template("control_guidance_end"),
+ InputParam.template("controlnet_conditioning_scale"),
+ InputParam(
+ name="control_image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The control image latents to use for the denoising process. Can be generated in controlnet vae encoder step.",
+ ),
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ controlnet = unwrap_module(components.controlnet)
+
+ # control_guidance_start/control_guidance_end (align format)
+ if not isinstance(block_state.control_guidance_start, list) and isinstance(
+ block_state.control_guidance_end, list
+ ):
+ block_state.control_guidance_start = len(block_state.control_guidance_end) * [
+ block_state.control_guidance_start
+ ]
+ elif not isinstance(block_state.control_guidance_end, list) and isinstance(
+ block_state.control_guidance_start, list
+ ):
+ block_state.control_guidance_end = len(block_state.control_guidance_start) * [
+ block_state.control_guidance_end
+ ]
+ elif not isinstance(block_state.control_guidance_start, list) and not isinstance(
+ block_state.control_guidance_end, list
+ ):
+ mult = (
+ len(block_state.control_image_latents) if isinstance(controlnet, QwenImageMultiControlNetModel) else 1
+ )
+ block_state.control_guidance_start, block_state.control_guidance_end = (
+ mult * [block_state.control_guidance_start],
+ mult * [block_state.control_guidance_end],
+ )
+
+ # controlnet_conditioning_scale (align format)
+ if isinstance(controlnet, QwenImageMultiControlNetModel) and isinstance(
+ block_state.controlnet_conditioning_scale, float
+ ):
+ block_state.controlnet_conditioning_scale = [block_state.controlnet_conditioning_scale] * mult
+
+ # controlnet_keep
+ block_state.controlnet_keep = []
+ for i in range(len(block_state.timesteps)):
+ keeps = [
+ 1.0 - float(i / len(block_state.timesteps) < s or (i + 1) / len(block_state.timesteps) > e)
+ for s, e in zip(block_state.control_guidance_start, block_state.control_guidance_end)
+ ]
+ block_state.controlnet_keep.append(keeps[0] if isinstance(controlnet, QwenImageControlNetModel) else keeps)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..d2723556f4d3ae099e1c7311b4016e8511ad1193
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/decoders.py
@@ -0,0 +1,511 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import InpaintProcessor, VaeImageProcessor
+from ...models import AutoencoderKLQwenImage
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier
+
+
+logger = logging.get_logger(__name__)
+
+
+# after denoising loop (unpack latents)
+
+
+# auto_docstring
+class QwenImageAfterDenoiseStep(ModularPipelineBlocks):
+ """
+ Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size,
+ channels, 1, height, width)
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ latents (`Tensor`):
+ The latents to decode, can be generated in the denoise step.
+
+ Outputs:
+ latents (`Tensor`):
+ The denoisedlatents unpacked to B, C, 1, H, W
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, channels, 1, height, width)"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ components = [
+ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+ ]
+
+ return components
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents to decode, can be generated in the denoise step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="latents", type_hint=torch.Tensor, description="The denoisedlatents unpacked to B, C, 1, H, W"
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ vae_scale_factor = components.vae_scale_factor
+ block_state.latents = components.pachifier.unpack_latents(
+ block_state.latents, block_state.height, block_state.width, vae_scale_factor=vae_scale_factor
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredAfterDenoiseStep(ModularPipelineBlocks):
+ """
+ Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W) after denoising.
+
+ Components:
+ pachifier (`QwenImageLayeredPachifier`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ layers (`int`, *optional*, defaults to 4):
+ Number of layers to extract from the image
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents. (unpacked to B, C, layers+1, H, W)
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W) after denoising."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents to decode, can be generated in the denoise step.",
+ ),
+ InputParam.template("height", required=True),
+ InputParam.template("width", required=True),
+ InputParam.template("layers"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("latents", note="unpacked to B, C, layers+1, H, W"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Unpack: (B, seq, C*4) -> (B, C, layers+1, H, W)
+ block_state.latents = components.pachifier.unpack_latents(
+ block_state.latents,
+ block_state.height,
+ block_state.width,
+ block_state.layers,
+ components.vae_scale_factor,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# decode step
+
+
+# auto_docstring
+class QwenImageDecoderStep(ModularPipelineBlocks):
+ """
+ Step that decodes the latents to images
+
+ Components:
+ vae (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the latents to images"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ components = [
+ ComponentSpec("vae", AutoencoderKLQwenImage),
+ ]
+
+ return components
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images", note="tensor output of the vae decoder.")]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # YiYi Notes: remove support for output_type = "latents', we can just skip decode/encode step in modular
+ if block_state.latents.ndim == 4:
+ block_state.latents = block_state.latents.unsqueeze(dim=1)
+ elif block_state.latents.ndim != 5:
+ raise ValueError(
+ f"expect latents to be a 4D or 5D tensor but got: {block_state.latents.shape}. Please make sure the latents are unpacked before decode step."
+ )
+ block_state.latents = block_state.latents.to(components.vae.dtype)
+
+ latents_mean = (
+ torch.tensor(components.vae.config.latents_mean)
+ .view(1, components.vae.config.z_dim, 1, 1, 1)
+ .to(block_state.latents.device, block_state.latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+ 1, components.vae.config.z_dim, 1, 1, 1
+ ).to(block_state.latents.device, block_state.latents.dtype)
+ block_state.latents = block_state.latents / latents_std + latents_mean
+ block_state.images = components.vae.decode(block_state.latents, return_dict=False)[0][:, :, 0]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredDecoderStep(ModularPipelineBlocks):
+ """
+ Decode unpacked latents (B, C, layers+1, H, W) into layer images.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Decode unpacked latents (B, C, layers+1, H, W) into layer images."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLQwenImage),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise step.",
+ ),
+ InputParam.template("output_type"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latents = block_state.latents
+
+ # 1. VAE normalization
+ latents = latents.to(components.vae.dtype)
+ latents_mean = (
+ torch.tensor(components.vae.config.latents_mean)
+ .view(1, components.vae.config.z_dim, 1, 1, 1)
+ .to(latents.device, latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+ 1, components.vae.config.z_dim, 1, 1, 1
+ ).to(latents.device, latents.dtype)
+ latents = latents / latents_std + latents_mean
+
+ # 2. Reshape for batch decoding: (B, C, layers+1, H, W) -> (B*layers, C, 1, H, W)
+ b, c, f, h, w = latents.shape
+ # 3. Remove first frame (composite), keep layers frames
+ latents = latents[:, :, 1:]
+ latents = latents.permute(0, 2, 1, 3, 4).reshape(-1, c, 1, h, w)
+
+ # 4. Decode: (B*layers, C, 1, H, W) -> (B*layers, C, H, W)
+ image = components.vae.decode(latents, return_dict=False)[0]
+ image = image.squeeze(2)
+
+ # 5. Postprocess - returns flat list of B*layers images
+ image = components.image_processor.postprocess(image, output_type=block_state.output_type)
+
+ # 6. Chunk into list per batch item
+ images = []
+ for bidx in range(b):
+ images.append(image[bidx * f : (bidx + 1) * f])
+
+ block_state.images = images
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# postprocess the decoded images
+
+
+# auto_docstring
+class QwenImageProcessImagesOutputStep(ModularPipelineBlocks):
+ """
+ postprocess the generated image
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ images (`Tensor`):
+ the generated image tensor from decoders step
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "postprocess the generated image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="images",
+ required=True,
+ type_hint=torch.Tensor,
+ description="the generated image tensor from decoders step",
+ ),
+ InputParam.template("output_type"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
+
+ @staticmethod
+ def check_inputs(output_type):
+ if output_type not in ["pil", "np", "pt"]:
+ raise ValueError(f"Invalid output_type: {output_type}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(block_state.output_type)
+
+ block_state.images = components.image_processor.postprocess(
+ image=block_state.images,
+ output_type=block_state.output_type,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageInpaintProcessImagesOutputStep(ModularPipelineBlocks):
+ """
+ postprocess the generated image, optional apply the mask overally to the original image..
+
+ Components:
+ image_mask_processor (`InpaintProcessor`)
+
+ Inputs:
+ images (`Tensor`):
+ the generated image tensor from decoders step
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ mask_overlay_kwargs (`dict`, *optional*):
+ The kwargs for the postprocess step to apply the mask overlay. generated in
+ InpaintProcessImagesInputStep.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "postprocess the generated image, optional apply the mask overally to the original image.."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_mask_processor",
+ InpaintProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="images",
+ required=True,
+ type_hint=torch.Tensor,
+ description="the generated image tensor from decoders step",
+ ),
+ InputParam.template("output_type"),
+ InputParam(
+ name="mask_overlay_kwargs",
+ type_hint=dict[str, Any],
+ description="The kwargs for the postprocess step to apply the mask overlay. generated in InpaintProcessImagesInputStep.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam.template("images")]
+
+ @staticmethod
+ def check_inputs(output_type, mask_overlay_kwargs):
+ if output_type not in ["pil", "np", "pt"]:
+ raise ValueError(f"Invalid output_type: {output_type}")
+
+ if mask_overlay_kwargs and output_type != "pil":
+ raise ValueError("only support output_type 'pil' for mask overlay")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(block_state.output_type, block_state.mask_overlay_kwargs)
+
+ if block_state.mask_overlay_kwargs is None:
+ mask_overlay_kwargs = {}
+ else:
+ mask_overlay_kwargs = block_state.mask_overlay_kwargs
+
+ block_state.images = components.image_mask_processor.postprocess(
+ image=block_state.images,
+ **mask_overlay_kwargs,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef6fd46ca01af4b2bb3c373d56d4dac37f6e6fe2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/denoise.py
@@ -0,0 +1,943 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import QwenImageControlNetModel, QwenImageTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import QwenImageModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+# ====================
+# 1. LOOP STEPS (run at each denoising step)
+# ====================
+
+
+# loop step:before denoiser
+class QwenImageLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # one timestep
+ block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype)
+ block_state.latent_model_input = block_state.latents
+ return components, block_state
+
+
+class QwenImageEditLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "qwenimage-edit"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam.template("image_latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # one timestep
+
+ block_state.latent_model_input = torch.cat([block_state.latents, block_state.image_latents], dim=1)
+ block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype)
+ return components, block_state
+
+
+class QwenImageLoopBeforeDenoiserControlNet(ModularPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("controlnet", QwenImageControlNetModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that runs the controlnet before the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "control_image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ InputParam.template("controlnet_conditioning_scale", note="updated in prepare_controlnet_inputs step."),
+ InputParam(
+ name="controlnet_keep",
+ required=True,
+ type_hint=list[float],
+ description="The controlnet keep values. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: int):
+ # cond_scale for the timestep (controlnet input)
+ if isinstance(block_state.controlnet_keep[i], list):
+ block_state.cond_scale = [
+ c * s for c, s in zip(block_state.controlnet_conditioning_scale, block_state.controlnet_keep[i])
+ ]
+ else:
+ controlnet_cond_scale = block_state.controlnet_conditioning_scale
+ if isinstance(controlnet_cond_scale, list):
+ controlnet_cond_scale = controlnet_cond_scale[0]
+ block_state.cond_scale = controlnet_cond_scale * block_state.controlnet_keep[i]
+
+ # run controlnet for the guidance batch
+ controlnet_block_samples = components.controlnet(
+ hidden_states=block_state.latent_model_input,
+ controlnet_cond=block_state.control_image_latents,
+ conditioning_scale=block_state.cond_scale,
+ timestep=block_state.timestep / 1000,
+ img_shapes=block_state.img_shapes,
+ encoder_hidden_states=block_state.prompt_embeds,
+ encoder_hidden_states_mask=block_state.prompt_embeds_mask,
+ return_dict=False,
+ )
+
+ block_state.additional_cond_kwargs["controlnet_block_samples"] = controlnet_block_samples
+
+ return components, block_state
+
+
+# loop step:denoiser
+class QwenImageLoopDenoiser(ModularPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that denoise the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", QwenImageTransformer2DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ "img_shapes",
+ required=True,
+ type_hint=list[tuple[int, int]],
+ description="The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ guider_inputs = {
+ "encoder_hidden_states": (
+ getattr(block_state, "prompt_embeds", None),
+ getattr(block_state, "negative_prompt_embeds", None),
+ ),
+ "encoder_hidden_states_mask": (
+ getattr(block_state, "prompt_embeds_mask", None),
+ getattr(block_state, "negative_prompt_embeds_mask", None),
+ ),
+ }
+
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+ additional_cond_kwargs = {}
+ for field_name, field_value in block_state.denoiser_input_fields.items():
+ if field_name in transformer_args and field_name not in guider_inputs:
+ additional_cond_kwargs[field_name] = field_value
+ block_state.additional_cond_kwargs.update(additional_cond_kwargs)
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+ # YiYi TODO: add cache context
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=block_state.timestep / 1000,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ **block_state.additional_cond_kwargs,
+ )[0]
+
+ components.guider.cleanup_models(components.transformer)
+
+ guider_output = components.guider(guider_state)
+
+ # apply guidance rescale
+ pred_cond_norm = torch.norm(guider_output.pred_cond, dim=-1, keepdim=True)
+ pred_norm = torch.norm(guider_output.pred, dim=-1, keepdim=True)
+ block_state.noise_pred = guider_output.pred * (pred_cond_norm / pred_norm)
+
+ return components, block_state
+
+
+class QwenImageEditLoopDenoiser(ModularPipelineBlocks):
+ model_name = "qwenimage-edit"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that denoise the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", QwenImageTransformer2DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("attention_kwargs"),
+ InputParam.template("denoiser_input_fields"),
+ InputParam(
+ "img_shapes",
+ required=True,
+ type_hint=list[tuple[int, int]],
+ description="The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ guider_inputs = {
+ "encoder_hidden_states": (
+ getattr(block_state, "prompt_embeds", None),
+ getattr(block_state, "negative_prompt_embeds", None),
+ ),
+ "encoder_hidden_states_mask": (
+ getattr(block_state, "prompt_embeds_mask", None),
+ getattr(block_state, "negative_prompt_embeds_mask", None),
+ ),
+ }
+
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+ additional_cond_kwargs = {}
+ for field_name, field_value in block_state.denoiser_input_fields.items():
+ if field_name in transformer_args and field_name not in guider_inputs:
+ additional_cond_kwargs[field_name] = field_value
+ block_state.additional_cond_kwargs.update(additional_cond_kwargs)
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+ # YiYi TODO: add cache context
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input,
+ timestep=block_state.timestep / 1000,
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ **block_state.additional_cond_kwargs,
+ )[0]
+
+ components.guider.cleanup_models(components.transformer)
+
+ guider_output = components.guider(guider_state)
+
+ pred = guider_output.pred[:, : block_state.latents.size(1)]
+ pred_cond = guider_output.pred_cond[:, : block_state.latents.size(1)]
+
+ # apply guidance rescale
+ pred_cond_norm = torch.norm(pred_cond, dim=-1, keepdim=True)
+ pred_norm = torch.norm(pred, dim=-1, keepdim=True)
+ block_state.noise_pred = pred * (pred_cond_norm / pred_norm)
+
+ return components, block_state
+
+
+# loop step:after denoiser
+class QwenImageLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that updates the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class QwenImageLoopAfterDenoiserInpaint(ModularPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that updates the latents using mask and image_latents for inpainting. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `QwenImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.",
+ ),
+ InputParam.template("image_latents"),
+ InputParam(
+ "initial_noise",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("latents"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.init_latents_proper = block_state.image_latents
+ if i < len(block_state.timesteps) - 1:
+ block_state.noise_timestep = block_state.timesteps[i + 1]
+ block_state.init_latents_proper = components.scheduler.scale_noise(
+ block_state.init_latents_proper, torch.tensor([block_state.noise_timestep]), block_state.initial_noise
+ )
+
+ block_state.latents = (
+ 1 - block_state.mask
+ ) * block_state.init_latents_proper + block_state.mask * block_state.latents
+
+ return components, block_state
+
+
+# ====================
+# 2. DENOISE LOOP WRAPPER: define the denoising loop logic
+# ====================
+class QwenImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoise the latents over `timesteps`. "
+ "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam.template("num_inference_steps", required=True),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ block_state.additional_cond_kwargs = {}
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# ====================
+# 3. DENOISE STEPS: compose the denoising loop with loop wrapper + loop steps
+# ====================
+
+
+# Qwen Image (text2image, image2image)
+
+
+# auto_docstring
+class QwenImageDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageLoopBeforeDenoiser`
+ - `QwenImageLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ This block supports text2image and image2image tasks for QwenImage.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+
+ block_classes = [
+ QwenImageLoopBeforeDenoiser,
+ QwenImageLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents.\n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method\n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageLoopBeforeDenoiser`\n"
+ " - `QwenImageLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ "This block supports text2image and image2image tasks for QwenImage."
+ )
+
+
+# Qwen Image (inpainting)
+# auto_docstring
+class QwenImageInpaintDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageLoopBeforeDenoiser`
+ - `QwenImageLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ - `QwenImageLoopAfterDenoiserInpaint`
+ This block supports inpainting tasks for QwenImage.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
+ mask (`Tensor`):
+ The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ initial_noise (`Tensor`):
+ The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageLoopBeforeDenoiser,
+ QwenImageLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ QwenImageLoopAfterDenoiserInpaint,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser", "after_denoiser_inpaint"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageLoopBeforeDenoiser`\n"
+ " - `QwenImageLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiserInpaint`\n"
+ "This block supports inpainting tasks for QwenImage."
+ )
+
+
+# Qwen Image (text2image, image2image) with controlnet
+# auto_docstring
+class QwenImageControlNetDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageLoopBeforeDenoiser`
+ - `QwenImageLoopBeforeDenoiserControlNet`
+ - `QwenImageLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ This block supports text2img/img2img tasks with controlnet for QwenImage.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) controlnet (`QwenImageControlNetModel`) transformer
+ (`QwenImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ control_image_latents (`Tensor`):
+ The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.)
+ controlnet_keep (`list`):
+ The controlnet keep values. Can be generated in prepare_controlnet_inputs step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageLoopBeforeDenoiser,
+ QwenImageLoopBeforeDenoiserControlNet,
+ QwenImageLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "before_denoiser_controlnet", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageLoopBeforeDenoiser`\n"
+ " - `QwenImageLoopBeforeDenoiserControlNet`\n"
+ " - `QwenImageLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ "This block supports text2img/img2img tasks with controlnet for QwenImage."
+ )
+
+
+# Qwen Image (inpainting) with controlnet
+# auto_docstring
+class QwenImageInpaintControlNetDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageLoopBeforeDenoiser`
+ - `QwenImageLoopBeforeDenoiserControlNet`
+ - `QwenImageLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ - `QwenImageLoopAfterDenoiserInpaint`
+ This block supports inpainting tasks with controlnet for QwenImage.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) controlnet (`QwenImageControlNetModel`) transformer
+ (`QwenImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ control_image_latents (`Tensor`):
+ The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning. (updated in prepare_controlnet_inputs step.)
+ controlnet_keep (`list`):
+ The controlnet keep values. Can be generated in prepare_controlnet_inputs step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. can be generated in prepare_additional_inputs step.
+ mask (`Tensor`):
+ The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ initial_noise (`Tensor`):
+ The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageLoopBeforeDenoiser,
+ QwenImageLoopBeforeDenoiserControlNet,
+ QwenImageLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ QwenImageLoopAfterDenoiserInpaint,
+ ]
+ block_names = [
+ "before_denoiser",
+ "before_denoiser_controlnet",
+ "denoiser",
+ "after_denoiser",
+ "after_denoiser_inpaint",
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageLoopBeforeDenoiser`\n"
+ " - `QwenImageLoopBeforeDenoiserControlNet`\n"
+ " - `QwenImageLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiserInpaint`\n"
+ "This block supports inpainting tasks with controlnet for QwenImage."
+ )
+
+
+# Qwen Image Edit (image2image)
+# auto_docstring
+class QwenImageEditDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageEditLoopBeforeDenoiser`
+ - `QwenImageEditLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ This block supports QwenImage Edit.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditLoopBeforeDenoiser,
+ QwenImageEditLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageEditLoopBeforeDenoiser`\n"
+ " - `QwenImageEditLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ "This block supports QwenImage Edit."
+ )
+
+
+# Qwen Image Edit (inpainting)
+# auto_docstring
+class QwenImageEditInpaintDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageEditLoopBeforeDenoiser`
+ - `QwenImageEditLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ - `QwenImageLoopAfterDenoiserInpaint`
+ This block supports inpainting tasks for QwenImage Edit.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
+ mask (`Tensor`):
+ The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.
+ initial_noise (`Tensor`):
+ The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditLoopBeforeDenoiser,
+ QwenImageEditLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ QwenImageLoopAfterDenoiserInpaint,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser", "after_denoiser_inpaint"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageEditLoopBeforeDenoiser`\n"
+ " - `QwenImageEditLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiserInpaint`\n"
+ "This block supports inpainting tasks for QwenImage Edit."
+ )
+
+
+# Qwen Image Layered (image2image)
+# auto_docstring
+class QwenImageLayeredDenoiseStep(QwenImageDenoiseLoopWrapper):
+ """
+ Denoise step that iteratively denoise the latents.
+ Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method At each iteration, it runs blocks
+ defined in `sub_blocks` sequencially:
+ - `QwenImageEditLoopBeforeDenoiser`
+ - `QwenImageEditLoopDenoiser`
+ - `QwenImageLoopAfterDenoiser`
+ This block supports QwenImage Layered.
+
+ Components:
+ guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ latents (`Tensor`):
+ The initial latents to use for the denoising process. Can be generated in prepare_latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ img_shapes (`list`):
+ The shape of the image latents for RoPE calculation. Can be generated in prepare_additional_inputs step.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = [
+ QwenImageEditLoopBeforeDenoiser,
+ QwenImageEditLoopDenoiser,
+ QwenImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `QwenImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+ " - `QwenImageEditLoopBeforeDenoiser`\n"
+ " - `QwenImageEditLoopDenoiser`\n"
+ " - `QwenImageLoopAfterDenoiser`\n"
+ "This block supports QwenImage Layered."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..aaaafa98e2ac6fd669a04174458c7d2c7b30288e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/encoders.py
@@ -0,0 +1,1780 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Text and VAE encoder blocks for QwenImage pipelines.
+"""
+
+import PIL
+import torch
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import InpaintProcessor, VaeImageProcessor, is_valid_image, is_valid_image_imagelist
+from ...models import AutoencoderKLQwenImage, QwenImageControlNetModel, QwenImageMultiControlNetModel
+from ...pipelines.qwenimage.pipeline_qwenimage_edit import calculate_dimensions
+from ...utils import logging
+from ...utils.torch_utils import unwrap_module
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import QwenImageModularPipeline
+from .prompt_templates import (
+ QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE,
+ QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE,
+ QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX,
+ QWENIMAGE_EDIT_PROMPT_TEMPLATE,
+ QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX,
+ QWENIMAGE_LAYERED_CAPTION_PROMPT_CN,
+ QWENIMAGE_LAYERED_CAPTION_PROMPT_EN,
+ QWENIMAGE_PROMPT_TEMPLATE,
+ QWENIMAGE_PROMPT_TEMPLATE_START_IDX,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
+ bool_mask = mask.bool()
+ valid_lengths = bool_mask.sum(dim=1)
+ selected = hidden_states[bool_mask]
+ split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+ return split_result
+
+
+def get_qwen_prompt_embeds(
+ text_encoder,
+ tokenizer,
+ prompt: str | list[str] = None,
+ prompt_template_encode: str = QWENIMAGE_PROMPT_TEMPLATE,
+ prompt_template_encode_start_idx: int = QWENIMAGE_PROMPT_TEMPLATE_START_IDX,
+ tokenizer_max_length: int = 1024,
+ device: torch.device | None = None,
+):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ template = prompt_template_encode
+ drop_idx = prompt_template_encode_start_idx
+ txt = [template.format(e) for e in prompt]
+ txt_tokens = tokenizer(
+ txt, max_length=tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt"
+ ).to(device)
+ encoder_hidden_states = text_encoder(
+ input_ids=txt_tokens.input_ids,
+ attention_mask=txt_tokens.attention_mask,
+ output_hidden_states=True,
+ )
+ hidden_states = encoder_hidden_states.hidden_states[-1]
+
+ split_hidden_states = _extract_masked_hidden(hidden_states, txt_tokens.attention_mask)
+ split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+ attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+ max_seq_len = max([e.size(0) for e in split_hidden_states])
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+ )
+ encoder_attention_mask = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+ )
+
+ prompt_embeds = prompt_embeds.to(device=device)
+
+ return prompt_embeds, encoder_attention_mask
+
+
+def get_qwen_prompt_embeds_edit(
+ text_encoder,
+ processor,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | None = None,
+ prompt_template_encode: str = QWENIMAGE_EDIT_PROMPT_TEMPLATE,
+ prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX,
+ device: torch.device | None = None,
+):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ template = prompt_template_encode
+ drop_idx = prompt_template_encode_start_idx
+ txt = [template.format(e) for e in prompt]
+
+ model_inputs = processor(
+ text=txt,
+ images=image,
+ padding=True,
+ return_tensors="pt",
+ ).to(device)
+
+ outputs = text_encoder(
+ input_ids=model_inputs["input_ids"],
+ attention_mask=model_inputs["attention_mask"],
+ pixel_values=model_inputs.get("pixel_values"),
+ image_grid_thw=model_inputs.get("image_grid_thw"),
+ output_hidden_states=True,
+ )
+
+ hidden_states = outputs.hidden_states[-1]
+ split_hidden_states = _extract_masked_hidden(hidden_states, model_inputs["attention_mask"])
+ split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+ attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+ max_seq_len = max([e.size(0) for e in split_hidden_states])
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+ )
+ encoder_attention_mask = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+ )
+
+ prompt_embeds = prompt_embeds.to(device=device)
+
+ return prompt_embeds, encoder_attention_mask
+
+
+def get_qwen_prompt_embeds_edit_plus(
+ text_encoder,
+ processor,
+ prompt: str | list[str] = None,
+ image: torch.Tensor | list[PIL.Image.Image, PIL.Image.Image] | None = None,
+ prompt_template_encode: str = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE,
+ img_template_encode: str = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE,
+ prompt_template_encode_start_idx: int = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX,
+ device: torch.device | None = None,
+):
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if isinstance(image, list):
+ base_img_prompt = ""
+ for i, img in enumerate(image):
+ base_img_prompt += img_template_encode.format(i + 1)
+ elif image is not None:
+ base_img_prompt = img_template_encode.format(1)
+ else:
+ base_img_prompt = ""
+
+ template = prompt_template_encode
+
+ drop_idx = prompt_template_encode_start_idx
+ txt = [template.format(base_img_prompt + e) for e in prompt]
+
+ model_inputs = processor(
+ text=txt,
+ images=image,
+ padding=True,
+ return_tensors="pt",
+ ).to(device)
+ outputs = text_encoder(
+ input_ids=model_inputs["input_ids"],
+ attention_mask=model_inputs["attention_mask"],
+ pixel_values=model_inputs.get("pixel_values"),
+ image_grid_thw=model_inputs.get("image_grid_thw"),
+ output_hidden_states=True,
+ )
+
+ hidden_states = outputs.hidden_states[-1]
+ split_hidden_states = _extract_masked_hidden(hidden_states, model_inputs["attention_mask"])
+ split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+ attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+ max_seq_len = max([e.size(0) for e in split_hidden_states])
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+ )
+ encoder_attention_mask = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+ )
+
+ prompt_embeds = prompt_embeds.to(device=device)
+ return prompt_embeds, encoder_attention_mask
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+# Modified from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._encode_vae_image
+def encode_vae_image(
+ image: torch.Tensor,
+ vae: AutoencoderKLQwenImage,
+ generator: torch.Generator,
+ device: torch.device,
+ dtype: torch.dtype,
+ latent_channels: int = 16,
+ sample_mode: str = "argmax",
+):
+ if not isinstance(image, torch.Tensor):
+ raise ValueError(f"Expected image to be a tensor, got {type(image)}.")
+
+ # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width
+ if image.dim() == 4:
+ image = image.unsqueeze(2)
+ elif image.dim() != 5:
+ raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.")
+
+ image = image.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode)
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean)
+ .view(1, latent_channels, 1, 1, 1)
+ .to(image_latents.device, image_latents.dtype)
+ )
+ latents_std = (
+ torch.tensor(vae.config.latents_std)
+ .view(1, latent_channels, 1, 1, 1)
+ .to(image_latents.device, image_latents.dtype)
+ )
+ image_latents = (image_latents - latents_mean) / latents_std
+
+ return image_latents
+
+
+# ====================
+# 1. RESIZE
+# ====================
+# In QwenImage pipelines, resize is a separate step because the resized image is used in VL encoding and vae encoder blocks:
+#
+# image (PIL.Image.Image)
+# │
+# ▼
+# resized_image ([PIL.Image.Image])
+# │
+# ├──► text_encoder ──► prompt_embeds, prompt_embeds_mask
+# │ (VL encoding needs the resized image for vision-language fusion)
+# │
+# └──► image_processor ──► processed_image (torch.Tensor, pixel space)
+# │
+# ▼
+# vae_encoder ──► image_latents (torch.Tensor, latent space)
+#
+# In most of our other pipelines, resizing is done as part of the image preprocessing step.
+# ====================
+
+
+# auto_docstring
+class QwenImageEditResizeStep(ModularPipelineBlocks):
+ """
+ Image Resize step that resize the image to target area while maintaining the aspect ratio.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ """
+
+ model_name = "qwenimage-edit"
+
+ @property
+ def description(self) -> str:
+ return "Image Resize step that resize the image to target area while maintaining the aspect ratio."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_resize_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [InputParam.template("image")]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="resized_image",
+ type_hint=list[PIL.Image.Image],
+ description="The resized images",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ images = block_state.image
+
+ if not is_valid_image_imagelist(images):
+ raise ValueError(f"Images must be image or list of images but are {type(images)}")
+
+ if is_valid_image(images):
+ images = [images]
+
+ image_width, image_height = images[0].size
+ calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height)
+
+ resized_images = [
+ components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width)
+ for image in images
+ ]
+
+ block_state.resized_image = resized_images
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageLayeredResizeStep(ModularPipelineBlocks):
+ """
+ Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while
+ maintaining the aspect ratio.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ resolution (`int`, *optional*, defaults to 640):
+ The target area to resize the image to, can be 1024 or 640
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Image Resize step that resize the image to a target area (defined by the resolution parameter from user) while maintaining the aspect ratio."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_resize_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image"),
+ InputParam(
+ name="resolution",
+ default=640,
+ type_hint=int,
+ description="The target area to resize the image to, can be 1024 or 640",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="resized_image",
+ type_hint=list[PIL.Image.Image],
+ description="The resized images",
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(resolution: int):
+ if resolution not in [1024, 640]:
+ raise ValueError(f"Resolution must be 1024 or 640 but is {resolution}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(resolution=block_state.resolution)
+
+ images = block_state.image
+
+ if not is_valid_image_imagelist(images):
+ raise ValueError(f"Images must be image or list of images but are {type(images)}")
+
+ if is_valid_image(images):
+ images = [images]
+
+ image_width, image_height = images[0].size
+ target_area = block_state.resolution * block_state.resolution
+ calculated_width, calculated_height, _ = calculate_dimensions(target_area, image_width / image_height)
+
+ resized_images = [
+ components.image_resize_processor.resize(image, height=calculated_height, width=calculated_width)
+ for image in images
+ ]
+
+ block_state.resized_image = resized_images
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditPlusResizeStep(ModularPipelineBlocks):
+ """
+ Resize images for QwenImage Edit Plus pipeline.
+ Produces two outputs: resized_image (1024x1024) for VAE encoding, resized_cond_image (384x384) for VL text
+ encoding. Each image is resized independently based on its own aspect ratio.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+
+ Outputs:
+ resized_image (`list`):
+ Images resized to 1024x1024 target area for VAE encoding
+ resized_cond_image (`list`):
+ Images resized to 384x384 target area for VL text encoding
+ """
+
+ model_name = "qwenimage-edit-plus"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Resize images for QwenImage Edit Plus pipeline.\n"
+ "Produces two outputs: resized_image (1024x1024) for VAE encoding, "
+ "resized_cond_image (384x384) for VL text encoding.\n"
+ "Each image is resized independently based on its own aspect ratio."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_resize_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ # image
+ return [InputParam.template("image")]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="resized_image",
+ type_hint=list[PIL.Image.Image],
+ description="Images resized to 1024x1024 target area for VAE encoding",
+ ),
+ OutputParam(
+ name="resized_cond_image",
+ type_hint=list[PIL.Image.Image],
+ description="Images resized to 384x384 target area for VL text encoding",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ images = block_state.image
+
+ if not is_valid_image_imagelist(images):
+ raise ValueError(f"Images must be image or list of images but are {type(images)}")
+
+ if is_valid_image(images):
+ images = [images]
+
+ # Resize each image independently based on its own aspect ratio
+ resized_images = []
+ resized_cond_images = []
+ for image in images:
+ image_width, image_height = image.size
+
+ # For VAE encoder (1024x1024 target area)
+ vae_width, vae_height, _ = calculate_dimensions(1024 * 1024, image_width / image_height)
+ resized_images.append(components.image_resize_processor.resize(image, height=vae_height, width=vae_width))
+
+ # For VL text encoder (384x384 target area)
+ vl_width, vl_height, _ = calculate_dimensions(384 * 384, image_width / image_height)
+ resized_cond_images.append(
+ components.image_resize_processor.resize(image, height=vl_height, width=vl_width)
+ )
+
+ block_state.resized_image = resized_images
+ block_state.resized_cond_image = resized_cond_images
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# ====================
+# 2. GET IMAGE PROMPT
+# ====================
+
+
+# auto_docstring
+class QwenImageLayeredGetImagePromptStep(ModularPipelineBlocks):
+ """
+ Auto-caption step that generates a text prompt from the input image if none is provided.
+ Uses the VL model (text_encoder) to generate a description of the image. If prompt is already provided, this step
+ passes through unchanged.
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ resized_image (`Image`):
+ The image to generate caption from, should be resized use the resize step
+ use_en_prompt (`bool`, *optional*, defaults to False):
+ Whether to use English prompt template
+
+ Outputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation. If not provided, updated using image caption
+ """
+
+ model_name = "qwenimage-layered"
+
+ def __init__(self):
+ self.image_caption_prompt_en = QWENIMAGE_LAYERED_CAPTION_PROMPT_EN
+ self.image_caption_prompt_cn = QWENIMAGE_LAYERED_CAPTION_PROMPT_CN
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Auto-caption step that generates a text prompt from the input image if none is provided.\n"
+ "Uses the VL model (text_encoder) to generate a description of the image.\n"
+ "If prompt is already provided, this step passes through unchanged."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
+ ComponentSpec("processor", Qwen2VLProcessor),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template(
+ "prompt", required=False
+ ), # it is not required for qwenimage-layered, unlike other pipelines
+ InputParam(
+ name="resized_image",
+ required=True,
+ type_hint=PIL.Image.Image,
+ description="The image to generate caption from, should be resized use the resize step",
+ ),
+ InputParam(
+ name="use_en_prompt",
+ default=False,
+ type_hint=bool,
+ description="Whether to use English prompt template",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="prompt",
+ type_hint=str,
+ description="The prompt or prompts to guide image generation. If not provided, updated using image caption",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ # If prompt is empty or None, generate caption from image
+ if block_state.prompt is None or block_state.prompt == "" or block_state.prompt == " ":
+ if block_state.use_en_prompt:
+ caption_prompt = self.image_caption_prompt_en
+ else:
+ caption_prompt = self.image_caption_prompt_cn
+
+ model_inputs = components.processor(
+ text=caption_prompt,
+ images=block_state.resized_image,
+ padding=True,
+ return_tensors="pt",
+ ).to(device)
+
+ generated_ids = components.text_encoder.generate(**model_inputs, max_new_tokens=512)
+ generated_ids_trimmed = [
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
+ ]
+ output_text = components.processor.batch_decode(
+ generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+ )[0]
+
+ block_state.prompt = output_text.strip()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# ====================
+# 3. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageTextEncoderStep(ModularPipelineBlocks):
+ """
+ Text Encoder step that generates text embeddings to guide the image generation.
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`):
+ The tokenizer to use guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 1024):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage"
+
+ def __init__(self):
+ self.prompt_template_encode = QWENIMAGE_PROMPT_TEMPLATE
+ self.prompt_template_encode_start_idx = QWENIMAGE_PROMPT_TEMPLATE_START_IDX
+ self.tokenizer_max_length = 1024
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings to guide the image generation."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration, description="The text encoder to use"),
+ ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam.template("max_sequence_length", default=1024),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("prompt_embeds_mask"),
+ OutputParam.template("negative_prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @staticmethod
+ def check_inputs(prompt, negative_prompt, max_sequence_length):
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if (
+ negative_prompt is not None
+ and not isinstance(negative_prompt, str)
+ and not isinstance(negative_prompt, list)
+ ):
+ raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+ if max_sequence_length is not None and max_sequence_length > 1024:
+ raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ self.check_inputs(block_state.prompt, block_state.negative_prompt, block_state.max_sequence_length)
+
+ block_state.prompt_embeds, block_state.prompt_embeds_mask = get_qwen_prompt_embeds(
+ components.text_encoder,
+ components.tokenizer,
+ prompt=block_state.prompt,
+ prompt_template_encode=self.prompt_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ tokenizer_max_length=self.tokenizer_max_length,
+ device=device,
+ )
+
+ block_state.prompt_embeds = block_state.prompt_embeds[:, : block_state.max_sequence_length]
+ block_state.prompt_embeds_mask = block_state.prompt_embeds_mask[:, : block_state.max_sequence_length]
+
+ block_state.negative_prompt_embeds = None
+ block_state.negative_prompt_embeds_mask = None
+ if components.requires_unconditional_embeds:
+ negative_prompt = block_state.negative_prompt or ""
+ block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_qwen_prompt_embeds(
+ components.text_encoder,
+ components.tokenizer,
+ prompt=negative_prompt,
+ prompt_template_encode=self.prompt_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ tokenizer_max_length=self.tokenizer_max_length,
+ device=device,
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds[
+ :, : block_state.max_sequence_length
+ ]
+ block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask[
+ :, : block_state.max_sequence_length
+ ]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditTextEncoderStep(ModularPipelineBlocks):
+ """
+ Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image
+ generation.
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ resized_image (`Image`):
+ The image prompt to encode, should be resized using resize step
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage"
+
+ def __init__(self):
+ self.prompt_template_encode = QWENIMAGE_EDIT_PROMPT_TEMPLATE
+ self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that processes both prompt and image together to generate text embeddings for guiding image generation."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
+ ComponentSpec("processor", Qwen2VLProcessor),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam(
+ name="resized_image",
+ required=True,
+ type_hint=PIL.Image.Image,
+ description="The image prompt to encode, should be resized using resize step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("prompt_embeds_mask"),
+ OutputParam.template("negative_prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @staticmethod
+ def check_inputs(prompt, negative_prompt):
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if (
+ negative_prompt is not None
+ and not isinstance(negative_prompt, str)
+ and not isinstance(negative_prompt, list)
+ ):
+ raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(block_state.prompt, block_state.negative_prompt)
+
+ device = components._execution_device
+
+ block_state.prompt_embeds, block_state.prompt_embeds_mask = get_qwen_prompt_embeds_edit(
+ components.text_encoder,
+ components.processor,
+ prompt=block_state.prompt,
+ image=block_state.resized_image,
+ prompt_template_encode=self.prompt_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ device=device,
+ )
+
+ block_state.negative_prompt_embeds = None
+ block_state.negative_prompt_embeds_mask = None
+ if components.requires_unconditional_embeds:
+ negative_prompt = block_state.negative_prompt or " "
+ block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_qwen_prompt_embeds_edit(
+ components.text_encoder,
+ components.processor,
+ prompt=negative_prompt,
+ image=block_state.resized_image,
+ prompt_template_encode=self.prompt_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ device=device,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditPlusTextEncoderStep(ModularPipelineBlocks):
+ """
+ Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together to generate text
+ embeddings for guiding image generation.
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor (`Qwen2VLProcessor`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ resized_cond_image (`Tensor`):
+ The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using
+ resize step
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage-edit-plus"
+
+ def __init__(self):
+ self.prompt_template_encode = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE
+ self.img_template_encode = QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE
+ self.prompt_template_encode_start_idx = QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text Encoder step for QwenImage Edit Plus that processes prompt and multiple images together "
+ "to generate text embeddings for guiding image generation."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen2_5_VLForConditionalGeneration),
+ ComponentSpec("processor", Qwen2VLProcessor),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("prompt"),
+ InputParam.template("negative_prompt"),
+ InputParam(
+ name="resized_cond_image",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The image(s) to encode, can be a single image or list of images, should be resized to 384x384 using resize step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam.template("prompt_embeds_mask"),
+ OutputParam.template("negative_prompt_embeds"),
+ OutputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @staticmethod
+ def check_inputs(prompt, negative_prompt):
+ if not isinstance(prompt, str) and not isinstance(prompt, list):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if (
+ negative_prompt is not None
+ and not isinstance(negative_prompt, str)
+ and not isinstance(negative_prompt, list)
+ ):
+ raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(block_state.prompt, block_state.negative_prompt)
+
+ device = components._execution_device
+
+ block_state.prompt_embeds, block_state.prompt_embeds_mask = get_qwen_prompt_embeds_edit_plus(
+ components.text_encoder,
+ components.processor,
+ prompt=block_state.prompt,
+ image=block_state.resized_cond_image,
+ prompt_template_encode=self.prompt_template_encode,
+ img_template_encode=self.img_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ device=device,
+ )
+
+ block_state.negative_prompt_embeds = None
+ block_state.negative_prompt_embeds_mask = None
+ if components.requires_unconditional_embeds:
+ negative_prompt = block_state.negative_prompt or " "
+ block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = (
+ get_qwen_prompt_embeds_edit_plus(
+ components.text_encoder,
+ components.processor,
+ prompt=negative_prompt,
+ image=block_state.resized_cond_image,
+ prompt_template_encode=self.prompt_template_encode,
+ img_template_encode=self.img_template_encode,
+ prompt_template_encode_start_idx=self.prompt_template_encode_start_idx,
+ device=device,
+ )
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# ====================
+# 4. IMAGE PREPROCESS
+# ====================
+
+
+# auto_docstring
+class QwenImageInpaintProcessImagesInputStep(ModularPipelineBlocks):
+ """
+ Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be
+ resized to the given height and width.
+
+ Components:
+ image_mask_processor (`InpaintProcessor`)
+
+ Inputs:
+ mask_image (`Image`):
+ Mask image for inpainting.
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ processed_mask_image (`Tensor`):
+ The processed mask image
+ mask_overlay_kwargs (`dict`):
+ The kwargs for the postprocess step to apply the mask overlay
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be resized to the given height and width."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_mask_processor",
+ InpaintProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("mask_image"),
+ InputParam.template("image"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("padding_mask_crop"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="processed_image",
+ type_hint=torch.Tensor,
+ description="The processed image",
+ ),
+ OutputParam(
+ name="processed_mask_image",
+ type_hint=torch.Tensor,
+ description="The processed mask image",
+ ),
+ OutputParam(
+ name="mask_overlay_kwargs",
+ type_hint=dict,
+ description="The kwargs for the postprocess step to apply the mask overlay",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
+ )
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+
+ block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = (
+ components.image_mask_processor.preprocess(
+ image=block_state.image,
+ mask=block_state.mask_image,
+ height=height,
+ width=width,
+ padding_mask_crop=block_state.padding_mask_crop,
+ )
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditInpaintProcessImagesInputStep(ModularPipelineBlocks):
+ """
+ Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be
+ resized first.
+
+ Components:
+ image_mask_processor (`InpaintProcessor`)
+
+ Inputs:
+ mask_image (`Image`):
+ Mask image for inpainting.
+ resized_image (`Image`):
+ The resized image. should be generated using a resize step
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ processed_mask_image (`Tensor`):
+ The processed mask image
+ mask_overlay_kwargs (`dict`):
+ The kwargs for the postprocess step to apply the mask overlay
+ """
+
+ model_name = "qwenimage-edit"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images should be resized first."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_mask_processor",
+ InpaintProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("mask_image"),
+ InputParam(
+ name="resized_image",
+ required=True,
+ type_hint=PIL.Image.Image,
+ description="The resized image. should be generated using a resize step",
+ ),
+ InputParam.template("padding_mask_crop"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="processed_image", type_hint=torch.Tensor, description="The processed image"),
+ OutputParam(
+ name="processed_mask_image",
+ type_hint=torch.Tensor,
+ description="The processed mask image",
+ ),
+ OutputParam(
+ name="mask_overlay_kwargs",
+ type_hint=dict,
+ description="The kwargs for the postprocess step to apply the mask overlay",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ width, height = block_state.resized_image[0].size
+
+ block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = (
+ components.image_mask_processor.preprocess(
+ image=block_state.resized_image,
+ mask=block_state.mask_image,
+ height=height,
+ width=width,
+ padding_mask_crop=block_state.padding_mask_crop,
+ )
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageProcessImagesInputStep(ModularPipelineBlocks):
+ """
+ Image Preprocess step. will resize the image to the given height and width.
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step. will resize the image to the given height and width."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="processed_image",
+ type_hint=torch.Tensor,
+ description="The processed image",
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
+ )
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+
+ block_state.processed_image = components.image_processor.preprocess(
+ image=block_state.image,
+ height=height,
+ width=width,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditProcessImagesInputStep(ModularPipelineBlocks):
+ """
+ Image Preprocess step. Images needs to be resized first.
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ resized_image (`list`):
+ The resized image. should be generated using a resize step
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ """
+
+ model_name = "qwenimage-edit"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step. Images needs to be resized first."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="resized_image",
+ required=True,
+ type_hint=list[PIL.Image.Image],
+ description="The resized image. should be generated using a resize step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="processed_image",
+ type_hint=torch.Tensor,
+ description="The processed image",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ width, height = block_state.resized_image[0].size
+
+ block_state.processed_image = components.image_processor.preprocess(
+ image=block_state.resized_image,
+ height=height,
+ width=width,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditPlusProcessImagesInputStep(ModularPipelineBlocks):
+ """
+ Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of
+ processed images.
+
+ Components:
+ image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ resized_image (`list`):
+ The resized image. should be generated using a resize step
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ """
+
+ model_name = "qwenimage-edit-plus"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step. Images can be resized first. If a list of images is provided, will return a list of processed images."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="resized_image",
+ required=True,
+ type_hint=list[PIL.Image.Image],
+ description="The resized image. should be generated using a resize step",
+ )
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="processed_image",
+ type_hint=torch.Tensor,
+ description="The processed image",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ image = block_state.resized_image
+
+ is_image_list = isinstance(image, list)
+ if not is_image_list:
+ image = [image]
+
+ processed_images = []
+ for img in image:
+ img_width, img_height = img.size
+ processed_images.append(
+ components.image_processor.preprocess(image=img, height=img_height, width=img_width)
+ )
+
+ if is_image_list:
+ block_state.processed_image = processed_images
+ else:
+ block_state.processed_image = processed_images[0]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# ====================
+# 5. VAE ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageVaeEncoderStep(ModularPipelineBlocks):
+ """
+ VAE Encoder step that converts processed_image into latent representations image_latents.
+ Handles both single images and lists of images with varied resolutions.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ processed_image (`Tensor`):
+ The image tensor to encode
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage"
+
+ def __init__(self, input: InputParam | None = None, output: OutputParam | None = None):
+ """Initialize a VAE encoder step for converting images to latent representations.
+
+ Handles both single images and lists of images. When input is a list, outputs a list of latents. When input is
+ a single tensor, outputs a single latent tensor.
+
+ Args:
+ input (InputParam, optional): Input parameter for the processed image. Defaults to "processed_image".
+ output (OutputParam, optional): Output parameter for the image latents. Defaults to "image_latents".
+ """
+ if input is None:
+ input = InputParam(
+ name="processed_image", required=True, type_hint=torch.Tensor, description="The image tensor to encode"
+ )
+
+ if output is None:
+ output = OutputParam.template("image_latents")
+
+ if not isinstance(input, InputParam):
+ raise ValueError(f"input must be InputParam but is {type(input)}")
+ if not isinstance(output, OutputParam):
+ raise ValueError(f"output must be OutputParam but is {type(output)}")
+
+ self._input = input
+ self._output = output
+ self._image_input_name = input.name
+ self._image_latents_output_name = output.name
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return (
+ f"VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}.\n"
+ "Handles both single images and lists of images with varied resolutions."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKLQwenImage)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ self._input, # default is "processed_image"
+ InputParam.template("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [self._output] # default is "image_latents"
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ image = getattr(block_state, self._image_input_name)
+ is_image_list = isinstance(image, list)
+ if not is_image_list:
+ image = [image]
+
+ # Handle both single image and list of images
+ image_latents = []
+ for img in image:
+ image_latents.append(
+ encode_vae_image(
+ image=img,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=dtype,
+ latent_channels=components.num_channels_latents,
+ )
+ )
+ if not is_image_list:
+ image_latents = image_latents[0]
+
+ setattr(block_state, self._image_latents_output_name, image_latents)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageControlNetVaeEncoderStep(ModularPipelineBlocks):
+ """
+ VAE Encoder step that converts `control_image` into latent representations control_image_latents.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) controlnet (`QwenImageControlNetModel`) control_image_processor
+ (`VaeImageProcessor`)
+
+ Inputs:
+ control_image (`Image`):
+ Control image for ControlNet conditioning.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ control_image_latents (`Tensor`):
+ The latents representing the control image
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "VAE Encoder step that converts `control_image` into latent representations control_image_latents.\n"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ components = [
+ ComponentSpec("vae", AutoencoderKLQwenImage),
+ ComponentSpec("controlnet", QwenImageControlNetModel),
+ ComponentSpec(
+ "control_image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+ return components
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("control_image"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ InputParam.template("generator"),
+ ]
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "control_image_latents",
+ type_hint=torch.Tensor,
+ description="The latents representing the control image",
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor):
+ if height is not None and height % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * 2) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(block_state.height, block_state.width, components.vae_scale_factor)
+
+ device = components._execution_device
+ dtype = components.vae.dtype
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+
+ controlnet = unwrap_module(components.controlnet)
+ if isinstance(controlnet, QwenImageMultiControlNetModel) and not isinstance(block_state.control_image, list):
+ block_state.control_image = [block_state.control_image]
+
+ if isinstance(controlnet, QwenImageMultiControlNetModel):
+ block_state.control_image_latents = []
+ for control_image_ in block_state.control_image:
+ control_image_ = components.control_image_processor.preprocess(
+ image=control_image_,
+ height=height,
+ width=width,
+ )
+
+ control_image_latents_ = encode_vae_image(
+ image=control_image_,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=dtype,
+ latent_channels=components.num_channels_latents,
+ sample_mode="sample",
+ )
+ block_state.control_image_latents.append(control_image_latents_)
+
+ elif isinstance(controlnet, QwenImageControlNetModel):
+ control_image = components.control_image_processor.preprocess(
+ image=block_state.control_image,
+ height=height,
+ width=width,
+ )
+ block_state.control_image_latents = encode_vae_image(
+ image=control_image,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=dtype,
+ latent_channels=components.num_channels_latents,
+ sample_mode="sample",
+ )
+
+ else:
+ raise ValueError(
+ f"Expected controlnet to be a QwenImageControlNetModel or QwenImageMultiControlNetModel, got {type(controlnet)}"
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# ====================
+# 6. PERMUTE LATENTS
+# ====================
+
+
+# auto_docstring
+class QwenImageLayeredPermuteLatentsStep(ModularPipelineBlocks):
+ """
+ Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing.
+
+ Inputs:
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ image_latents (`Tensor`):
+ The latent representation of the input image. (permuted from [B, C, 1, H, W] to [B, 1, C, H, W])
+ """
+
+ model_name = "qwenimage-layered"
+
+ @property
+ def description(self) -> str:
+ return "Permute image latents from (B, C, 1, H, W) to (B, 1, C, H, W) for Layered packing."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("image_latents"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam.template("image_latents", note="permuted from [B, C, 1, H, W] to [B, 1, C, H, W]"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Permute: (B, C, 1, H, W) -> (B, 1, C, H, W)
+ latents = block_state.image_latents
+ block_state.image_latents = latents.permute(0, 2, 1, 3, 4)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/inputs.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/inputs.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d1ba74b06ace9cb540dcedde2a9b925babd5e3b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/inputs.py
@@ -0,0 +1,1024 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...models import QwenImageMultiControlNetModel
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import QwenImageLayeredPachifier, QwenImageModularPipeline, QwenImagePachifier
+
+
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_images_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_images_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_images_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_images_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_images_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> tuple[int, int]:
+ """Calculate image dimensions from latent tensor dimensions.
+
+ This function converts latent space dimensions to image space dimensions by multiplying the latent height and width
+ by the VAE scale factor.
+
+ Args:
+ latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions.
+ Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width]
+ vae_scale_factor (int): The scale factor used by the VAE to compress images.
+ Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
+
+ Returns:
+ tuple[int, int]: The calculated image dimensions as (height, width)
+
+ Raises:
+ ValueError: If latents tensor doesn't have 4 or 5 dimensions
+
+ """
+ # make sure the latents are not packed
+ if latents.ndim != 4 and latents.ndim != 5:
+ raise ValueError(f"unpacked latents must have 4 or 5 dimensions, but got {latents.ndim}")
+
+ latent_height, latent_width = latents.shape[-2:]
+
+ height = latent_height * vae_scale_factor
+ width = latent_width * vae_scale_factor
+
+ return height, width
+
+
+# auto_docstring
+class QwenImageTextInputsStep(ModularPipelineBlocks):
+ """
+ Text input processing step that standardizes text embeddings for the pipeline.
+ This step:
+ 1. Determines `batch_size` and `dtype` based on `prompt_embeds`
+ 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)
+
+ This block should be placed after all encoder steps to process the text embeddings before they are used in
+ subsequent pipeline steps.
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ summary_section = (
+ "Text input processing step that standardizes text embeddings for the pipeline.\n"
+ "This step:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
+ )
+
+ # Placement guidance
+ placement_section = "\n\nThis block should be placed after all encoder steps to process the text embeddings before they are used in subsequent pipeline steps."
+
+ return summary_section + placement_section
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("prompt_embeds"),
+ InputParam.template("prompt_embeds_mask"),
+ InputParam.template("negative_prompt_embeds"),
+ InputParam.template("negative_prompt_embeds_mask"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(name="batch_size", type_hint=int, description="The batch size of the prompt embeddings"),
+ OutputParam(name="dtype", type_hint=torch.dtype, description="The data type of the prompt embeddings"),
+ OutputParam.template("prompt_embeds", note="batch-expanded"),
+ OutputParam.template("prompt_embeds_mask", note="batch-expanded"),
+ OutputParam.template("negative_prompt_embeds", note="batch-expanded"),
+ OutputParam.template("negative_prompt_embeds_mask", note="batch-expanded"),
+ ]
+
+ @staticmethod
+ def check_inputs(
+ prompt_embeds,
+ prompt_embeds_mask,
+ negative_prompt_embeds,
+ negative_prompt_embeds_mask,
+ ):
+ if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+ raise ValueError("`negative_prompt_embeds_mask` is required when `negative_prompt_embeds` is not None")
+
+ if negative_prompt_embeds is None and negative_prompt_embeds_mask is not None:
+ raise ValueError("cannot pass `negative_prompt_embeds_mask` without `negative_prompt_embeds`")
+
+ if prompt_embeds_mask.shape[0] != prompt_embeds.shape[0]:
+ raise ValueError("`prompt_embeds_mask` must have the same batch size as `prompt_embeds`")
+
+ elif negative_prompt_embeds is not None and negative_prompt_embeds.shape[0] != prompt_embeds.shape[0]:
+ raise ValueError("`negative_prompt_embeds` must have the same batch size as `prompt_embeds`")
+
+ elif (
+ negative_prompt_embeds_mask is not None and negative_prompt_embeds_mask.shape[0] != prompt_embeds.shape[0]
+ ):
+ raise ValueError("`negative_prompt_embeds_mask` must have the same batch size as `prompt_embeds`")
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ self.check_inputs(
+ prompt_embeds=block_state.prompt_embeds,
+ prompt_embeds_mask=block_state.prompt_embeds_mask,
+ negative_prompt_embeds=block_state.negative_prompt_embeds,
+ negative_prompt_embeds_mask=block_state.negative_prompt_embeds_mask,
+ )
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# auto_docstring
+class QwenImageAdditionalInputsStep(ModularPipelineBlocks):
+ """
+ Input processing step that:
+ 1. For image latent inputs: Updates height/width if None, patchifies, and expands batch size
+ 2. For additional batch inputs: Expands batch dimensions to match final batch size
+
+ Configured inputs:
+ - Image latent inputs: ['image_latents']
+
+ This block should be placed after the encoder steps and the text input step.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+ batch-expanded)
+ """
+
+ model_name = "qwenimage"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
+ ):
+ # by default, process `image_latents`
+ if image_latent_inputs is None:
+ image_latent_inputs = [InputParam.template("image_latents")]
+ if additional_batch_inputs is None:
+ additional_batch_inputs = []
+
+ if not isinstance(image_latent_inputs, list):
+ raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}")
+ else:
+ for input_param in image_latent_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}")
+
+ if not isinstance(additional_batch_inputs, list):
+ raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}")
+ else:
+ for input_param in additional_batch_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(
+ f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}"
+ )
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ summary_section = (
+ "Input processing step that:\n"
+ " 1. For image latent inputs: Updates height/width if None, patchifies, and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}"
+
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("batch_size"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+ # default is `image_latents`
+ inputs += self._image_latent_inputs + self._additional_batch_inputs
+
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ outputs = [
+ OutputParam(
+ name="image_height",
+ type_hint=int,
+ description="The image height calculated from the image latents dimension",
+ ),
+ OutputParam(
+ name="image_width",
+ type_hint=int,
+ description="The image width calculated from the image latents dimension",
+ ),
+ ]
+
+ # `height`/`width` are not new outputs, but they will be updated if any image latent inputs are provided
+ if len(self._image_latent_inputs) > 0:
+ outputs.append(
+ OutputParam(name="height", type_hint=int, description="if not provided, updated to image height")
+ )
+ outputs.append(
+ OutputParam(name="width", type_hint=int, description="if not provided, updated to image width")
+ )
+
+ # image latent inputs are modified in place (patchified and batch-expanded)
+ for input_param in self._image_latent_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (patchified and batch-expanded)",
+ )
+ )
+
+ # additional batch inputs (batch-expanded only)
+ for input_param in self._additional_batch_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (batch-expanded)",
+ )
+ )
+
+ return outputs
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs
+ for input_param in self._image_latent_inputs:
+ image_latent_input_name = input_param.name
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate height/width from latents and update if not provided
+ height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor)
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ if not hasattr(block_state, "image_height"):
+ block_state.image_height = height
+ if not hasattr(block_state, "image_width"):
+ block_state.image_width = width
+
+ # 2. Patchify
+ image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor)
+
+ # 3. Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=image_latent_input_name,
+ input_tensor=image_latent_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_param in self._additional_batch_inputs:
+ input_name = input_param.name
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageEditPlusAdditionalInputsStep(ModularPipelineBlocks):
+ """
+ Input processing step for Edit Plus that:
+ 1. For image latent inputs (list): Collects heights/widths, patchifies each, concatenates, expands batch
+ 2. For additional batch inputs: Expands batch dimensions to match final batch size
+ Height/width defaults to last image in the list.
+
+ Configured inputs:
+ - Image latent inputs: ['image_latents']
+
+ This block should be placed after the encoder steps and the text input step.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ image_height (`list`):
+ The image heights calculated from the image latents dimension
+ image_width (`list`):
+ The image widths calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified,
+ concatenated, and batch-expanded)
+ """
+
+ model_name = "qwenimage-edit-plus"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
+ ):
+ if image_latent_inputs is None:
+ image_latent_inputs = [InputParam.template("image_latents")]
+ if additional_batch_inputs is None:
+ additional_batch_inputs = []
+
+ if not isinstance(image_latent_inputs, list):
+ raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}")
+ else:
+ for input_param in image_latent_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}")
+
+ if not isinstance(additional_batch_inputs, list):
+ raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}")
+ else:
+ for input_param in additional_batch_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(
+ f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}"
+ )
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ summary_section = (
+ "Input processing step for Edit Plus that:\n"
+ " 1. For image latent inputs (list): Collects heights/widths, patchifies each, concatenates, expands batch\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size\n"
+ " Height/width defaults to last image in the list."
+ )
+
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}"
+
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImagePachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("batch_size"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ # default is `image_latents`
+ inputs += self._image_latent_inputs + self._additional_batch_inputs
+
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ outputs = [
+ OutputParam(
+ name="image_height",
+ type_hint=list[int],
+ description="The image heights calculated from the image latents dimension",
+ ),
+ OutputParam(
+ name="image_width",
+ type_hint=list[int],
+ description="The image widths calculated from the image latents dimension",
+ ),
+ ]
+
+ # `height`/`width` are updated if any image latent inputs are provided
+ if len(self._image_latent_inputs) > 0:
+ outputs.append(
+ OutputParam(name="height", type_hint=int, description="if not provided, updated to image height")
+ )
+ outputs.append(
+ OutputParam(name="width", type_hint=int, description="if not provided, updated to image width")
+ )
+
+ # image latent inputs are modified in place (patchified, concatenated, and batch-expanded)
+ for input_param in self._image_latent_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (patchified, concatenated, and batch-expanded)",
+ )
+ )
+
+ # additional batch inputs (batch-expanded only)
+ for input_param in self._additional_batch_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (batch-expanded)",
+ )
+ )
+
+ return outputs
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs
+ for input_param in self._image_latent_inputs:
+ image_latent_input_name = input_param.name
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ is_list = isinstance(image_latent_tensor, list)
+ if not is_list:
+ image_latent_tensor = [image_latent_tensor]
+
+ image_heights = []
+ image_widths = []
+ packed_image_latent_tensors = []
+
+ for i, img_latent_tensor in enumerate(image_latent_tensor):
+ # 1. Calculate height/width from latents
+ height, width = calculate_dimension_from_latents(img_latent_tensor, components.vae_scale_factor)
+ image_heights.append(height)
+ image_widths.append(width)
+
+ # 2. Patchify
+ img_latent_tensor = components.pachifier.pack_latents(img_latent_tensor)
+
+ # 3. Expand batch size
+ img_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=f"{image_latent_input_name}[{i}]",
+ input_tensor=img_latent_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+ packed_image_latent_tensors.append(img_latent_tensor)
+
+ # Concatenate all packed latents along dim=1
+ packed_image_latent_tensors = torch.cat(packed_image_latent_tensors, dim=1)
+
+ # Output lists of heights/widths
+ block_state.image_height = image_heights
+ block_state.image_width = image_widths
+
+ # Default height/width from last image
+ block_state.height = block_state.height or image_heights[-1]
+ block_state.width = block_state.width or image_widths[-1]
+
+ setattr(block_state, image_latent_input_name, packed_image_latent_tensors)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_param in self._additional_batch_inputs:
+ input_name = input_param.name
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# same as QwenImageAdditionalInputsStep, but with layered pachifier.
+
+
+# auto_docstring
+class QwenImageLayeredAdditionalInputsStep(ModularPipelineBlocks):
+ """
+ Input processing step for Layered that:
+ 1. For image latent inputs: Updates height/width if None, patchifies with layered pachifier, and expands batch
+ size
+ 2. For additional batch inputs: Expands batch dimensions to match final batch size
+
+ Configured inputs:
+ - Image latent inputs: ['image_latents']
+
+ This block should be placed after the encoder steps and the text input step.
+
+ Components:
+ pachifier (`QwenImageLayeredPachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified
+ with layered pachifier and batch-expanded)
+ """
+
+ model_name = "qwenimage-layered"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[InputParam] | None = None,
+ additional_batch_inputs: list[InputParam] | None = None,
+ ):
+ if image_latent_inputs is None:
+ image_latent_inputs = [InputParam.template("image_latents")]
+ if additional_batch_inputs is None:
+ additional_batch_inputs = []
+
+ if not isinstance(image_latent_inputs, list):
+ raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}")
+ else:
+ for input_param in image_latent_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}")
+
+ if not isinstance(additional_batch_inputs, list):
+ raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}")
+ else:
+ for input_param in additional_batch_inputs:
+ if not isinstance(input_param, InputParam):
+ raise ValueError(
+ f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}"
+ )
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ summary_section = (
+ "Input processing step for Layered that:\n"
+ " 1. For image latent inputs: Updates height/width if None, patchifies with layered pachifier, and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}"
+
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("pachifier", QwenImageLayeredPachifier, default_creation_method="from_config"),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("batch_size"),
+ ]
+ # default is `image_latents`
+
+ inputs += self._image_latent_inputs + self._additional_batch_inputs
+
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ outputs = [
+ OutputParam(
+ name="image_height",
+ type_hint=int,
+ description="The image height calculated from the image latents dimension",
+ ),
+ OutputParam(
+ name="image_width",
+ type_hint=int,
+ description="The image width calculated from the image latents dimension",
+ ),
+ ]
+
+ if len(self._image_latent_inputs) > 0:
+ outputs.append(
+ OutputParam(name="height", type_hint=int, description="if not provided, updated to image height")
+ )
+ outputs.append(
+ OutputParam(name="width", type_hint=int, description="if not provided, updated to image width")
+ )
+
+ # Add outputs for image latent inputs (patchified with layered pachifier and batch-expanded)
+ for input_param in self._image_latent_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (patchified with layered pachifier and batch-expanded)",
+ )
+ )
+
+ # Add outputs for additional batch inputs (batch-expanded only)
+ for input_param in self._additional_batch_inputs:
+ outputs.append(
+ OutputParam(
+ name=input_param.name,
+ type_hint=input_param.type_hint,
+ description=input_param.description + " (batch-expanded)",
+ )
+ )
+
+ return outputs
+
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs
+ for input_param in self._image_latent_inputs:
+ image_latent_input_name = input_param.name
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate height/width from latents and update if not provided
+ # Layered latents are (B, layers, C, H, W)
+ height = image_latent_tensor.shape[3] * components.vae_scale_factor
+ width = image_latent_tensor.shape[4] * components.vae_scale_factor
+ block_state.height = height
+ block_state.width = width
+
+ if not hasattr(block_state, "image_height"):
+ block_state.image_height = height
+ if not hasattr(block_state, "image_width"):
+ block_state.image_width = width
+
+ # 2. Patchify with layered pachifier
+ image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor)
+
+ # 3. Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=image_latent_input_name,
+ input_tensor=image_latent_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_param in self._additional_batch_inputs:
+ input_name = input_param.name
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+# auto_docstring
+class QwenImageControlNetInputsStep(ModularPipelineBlocks):
+ """
+ prepare the `control_image_latents` for controlnet. Insert after all the other inputs steps.
+
+ Inputs:
+ control_image_latents (`Tensor`):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ batch_size (`int`, *optional*, defaults to 1):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can
+ be generated in input step.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+
+ Outputs:
+ control_image_latents (`Tensor`):
+ The control image latents (patchified and batch-expanded).
+ height (`int`):
+ if not provided, updated to control image height
+ width (`int`):
+ if not provided, updated to control image width
+ """
+
+ model_name = "qwenimage"
+
+ @property
+ def description(self) -> str:
+ return "prepare the `control_image_latents` for controlnet. Insert after all the other inputs steps."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ name="control_image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The control image latents to use for the denoising process. Can be generated in controlnet vae encoder step.",
+ ),
+ InputParam.template("batch_size"),
+ InputParam.template("num_images_per_prompt"),
+ InputParam.template("height"),
+ InputParam.template("width"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ name="control_image_latents",
+ type_hint=torch.Tensor,
+ description="The control image latents (patchified and batch-expanded).",
+ ),
+ OutputParam(name="height", type_hint=int, description="if not provided, updated to control image height"),
+ OutputParam(name="width", type_hint=int, description="if not provided, updated to control image width"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: QwenImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ if isinstance(components.controlnet, QwenImageMultiControlNetModel):
+ control_image_latents = []
+ # loop through each control_image_latents
+ for i, control_image_latents_ in enumerate(block_state.control_image_latents):
+ # 1. update height/width if not provided
+ height, width = calculate_dimension_from_latents(control_image_latents_, components.vae_scale_factor)
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ # 2. pack
+ control_image_latents_ = components.pachifier.pack_latents(control_image_latents_)
+
+ # 3. repeat to match the batch size
+ control_image_latents_ = repeat_tensor_to_batch_size(
+ input_name=f"control_image_latents[{i}]",
+ input_tensor=control_image_latents_,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ control_image_latents.append(control_image_latents_)
+
+ block_state.control_image_latents = control_image_latents
+
+ else:
+ # 1. update height/width if not provided
+ height, width = calculate_dimension_from_latents(
+ block_state.control_image_latents, components.vae_scale_factor
+ )
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ # 2. pack
+ block_state.control_image_latents = components.pachifier.pack_latents(block_state.control_image_latents)
+
+ # 3. repeat to match the batch size
+ block_state.control_image_latents = repeat_tensor_to_batch_size(
+ input_name="control_image_latents",
+ input_tensor=block_state.control_image_latents,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ block_state.control_image_latents = block_state.control_image_latents
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4756ca53b5acb698ff18c3bf66f0f755c3a4360
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage.py
@@ -0,0 +1,1224 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+ QwenImageControlNetBeforeDenoiserStep,
+ QwenImageCreateMaskLatentsStep,
+ QwenImagePrepareLatentsStep,
+ QwenImagePrepareLatentsWithStrengthStep,
+ QwenImageRoPEInputsStep,
+ QwenImageSetTimestepsStep,
+ QwenImageSetTimestepsWithStrengthStep,
+)
+from .decoders import (
+ QwenImageAfterDenoiseStep,
+ QwenImageDecoderStep,
+ QwenImageInpaintProcessImagesOutputStep,
+ QwenImageProcessImagesOutputStep,
+)
+from .denoise import (
+ QwenImageControlNetDenoiseStep,
+ QwenImageDenoiseStep,
+ QwenImageInpaintControlNetDenoiseStep,
+ QwenImageInpaintDenoiseStep,
+)
+from .encoders import (
+ QwenImageControlNetVaeEncoderStep,
+ QwenImageInpaintProcessImagesInputStep,
+ QwenImageProcessImagesInputStep,
+ QwenImageTextEncoderStep,
+ QwenImageVaeEncoderStep,
+)
+from .inputs import (
+ QwenImageAdditionalInputsStep,
+ QwenImageControlNetInputsStep,
+ QwenImageTextInputsStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# ====================
+# 1. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageAutoTextEncoderStep(AutoPipelineBlocks):
+ """
+ Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block.
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`):
+ The tokenizer to use guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 1024):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImageTextEncoderStep()]
+ block_names = ["text_encoder"]
+ block_trigger_inputs = ["prompt"]
+
+ @property
+ def description(self) -> str:
+ return "Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block."
+ " - `QwenImageTextEncoderStep` (text_encoder) is used when `prompt` is provided."
+ " - if `prompt` is not provided, step will be skipped."
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageInpaintVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ This step is used for processing image and mask inputs for inpainting tasks. It:
+ - Resizes the image to the target size, based on `height` and `width`.
+ - Processes and updates `image` and `mask_image`.
+ - Creates `image_latents`.
+
+ Components:
+ image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ mask_image (`Image`):
+ Mask image for inpainting.
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ processed_mask_image (`Tensor`):
+ The processed mask image
+ mask_overlay_kwargs (`dict`):
+ The kwargs for the postprocess step to apply the mask overlay
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImageInpaintProcessImagesInputStep(), QwenImageVaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step is used for processing image and mask inputs for inpainting tasks. It:\n"
+ " - Resizes the image to the target size, based on `height` and `width`.\n"
+ " - Processes and updates `image` and `mask_image`.\n"
+ " - Creates `image_latents`."
+ )
+
+
+# auto_docstring
+class QwenImageImg2ImgVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that preprocess andencode the image inputs into their latent representations.
+
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ processed_image (`Tensor`):
+ The processed image
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage"
+
+ block_classes = [QwenImageProcessImagesInputStep(), QwenImageVaeEncoderStep()]
+ block_names = ["preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that preprocess andencode the image inputs into their latent representations."
+
+
+class QwenImageAutoVaeEncoderStep(AutoPipelineBlocks):
+ block_classes = [QwenImageInpaintVaeEncoderStep, QwenImageImg2ImgVaeEncoderStep]
+ block_names = ["inpaint", "img2img"]
+ block_trigger_inputs = ["mask_image", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block.\n"
+ + " - `QwenImageInpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n"
+ + " - `QwenImageImg2ImgVaeEncoderStep` (img2img) is used when `image` is provided.\n"
+ + " - if `mask_image` or `image` is not provided, step will be skipped."
+ )
+
+
+# optional controlnet vae encoder
+# auto_docstring
+class QwenImageOptionalControlNetVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+ This is an auto pipeline block.
+ - `QwenImageControlNetVaeEncoderStep` (controlnet) is used when `control_image` is provided.
+ - if `control_image` is not provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) controlnet (`QwenImageControlNetModel`) control_image_processor
+ (`VaeImageProcessor`)
+
+ Inputs:
+ control_image (`Image`, *optional*):
+ Control image for ControlNet conditioning.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ control_image_latents (`Tensor`):
+ The latents representing the control image
+ """
+
+ block_classes = [QwenImageControlNetVaeEncoderStep]
+ block_names = ["controlnet"]
+ block_trigger_inputs = ["control_image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block.\n"
+ + " - `QwenImageControlNetVaeEncoderStep` (controlnet) is used when `control_image` is provided.\n"
+ + " - if `control_image` is not provided, step will be skipped."
+ )
+
+
+# ====================
+# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
+# ====================
+
+
+# assemble input steps
+# auto_docstring
+class QwenImageImg2ImgInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the img2img denoising step. It:
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+ batch-expanded)
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImageTextInputsStep(), QwenImageAdditionalInputsStep()]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return "Input step that prepares the inputs for the img2img denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+
+
+# auto_docstring
+class QwenImageInpaintInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the inpainting denoising step. It:
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+ batch-expanded)
+ processed_mask_image (`Tensor`):
+ The processed mask image (batch-expanded)
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageAdditionalInputsStep(
+ additional_batch_inputs=[
+ InputParam(name="processed_mask_image", type_hint=torch.Tensor, description="The processed mask image")
+ ]
+ ),
+ ]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return "Input step that prepares the inputs for the inpainting denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents` and `processed_mask_image`).\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+
+
+# assemble prepare latents steps
+# auto_docstring
+class QwenImageInpaintPrepareLatentsStep(SequentialPipelineBlocks):
+ """
+ This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:
+ - Add noise to the image latents to create the latents input for the denoiser.
+ - Create the pachified latents `mask` based on the processedmask image.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ latents (`Tensor`):
+ The initial random noised, can be generated in prepare latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be
+ generated from vae encoder and updated in input step.)
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ processed_mask_image (`Tensor`):
+ The processed mask to use for the inpainting process.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ dtype (`dtype`, *optional*, defaults to torch.float32):
+ The dtype of the model inputs, can be generated in input step.
+
+ Outputs:
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ latents (`Tensor`):
+ The scaled noisy latents to use for inpainting/image-to-image denoising.
+ mask (`Tensor`):
+ The mask to use for the inpainting process.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImagePrepareLatentsWithStrengthStep(), QwenImageCreateMaskLatentsStep()]
+ block_names = ["add_noise_to_latents", "create_mask_latents"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:\n"
+ " - Add noise to the image latents to create the latents input for the denoiser.\n"
+ " - Create the pachified latents `mask` based on the processedmask image.\n"
+ )
+
+
+# assemble denoising steps
+
+
+# Qwen Image (text2image)
+# auto_docstring
+class QwenImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs
+ (timesteps, latents, rope inputs etc.).
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image (inpainting)
+# auto_docstring
+class QwenImageInpaintCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint
+ task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageInpaintInputStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsWithStrengthStep(),
+ QwenImageInpaintPrepareLatentsStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageInpaintDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_inpaint_latents",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image (image2image)
+# auto_docstring
+class QwenImageImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img
+ task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageImg2ImgInputStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsWithStrengthStep(),
+ QwenImagePrepareLatentsWithStrengthStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_img2img_latents",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image (text2image) with controlnet
+# auto_docstring
+class QwenImageControlNetCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs
+ (timesteps, latents, rope inputs etc.).
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet
+ (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ control_image_latents (`Tensor`):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ control_guidance_start (`float`, *optional*, defaults to 0.0):
+ When to start applying ControlNet.
+ control_guidance_end (`float`, *optional*, defaults to 1.0):
+ When to stop applying ControlNet.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageControlNetInputsStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageControlNetBeforeDenoiserStep(),
+ QwenImageControlNetDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "controlnet_input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_rope_inputs",
+ "controlnet_before_denoise",
+ "controlnet_denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image (inpainting) with controlnet
+# auto_docstring
+class QwenImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint
+ task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet
+ (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+ control_image_latents (`Tensor`):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ control_guidance_start (`float`, *optional*, defaults to 0.0):
+ When to start applying ControlNet.
+ control_guidance_end (`float`, *optional*, defaults to 1.0):
+ When to stop applying ControlNet.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageInpaintInputStep(),
+ QwenImageControlNetInputsStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsWithStrengthStep(),
+ QwenImageInpaintPrepareLatentsStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageControlNetBeforeDenoiserStep(),
+ QwenImageInpaintControlNetDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "controlnet_input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_inpaint_latents",
+ "prepare_rope_inputs",
+ "controlnet_before_denoise",
+ "controlnet_denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for inpaint task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image (image2image) with controlnet
+# auto_docstring
+class QwenImageControlNetImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img
+ task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) controlnet
+ (`QwenImageControlNetModel`) guider (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ control_image_latents (`Tensor`):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ control_guidance_start (`float`, *optional*, defaults to 0.0):
+ When to start applying ControlNet.
+ control_guidance_end (`float`, *optional*, defaults to 1.0):
+ When to stop applying ControlNet.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage"
+ block_classes = [
+ QwenImageImg2ImgInputStep(),
+ QwenImageControlNetInputsStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsWithStrengthStep(),
+ QwenImagePrepareLatentsWithStrengthStep(),
+ QwenImageRoPEInputsStep(),
+ QwenImageControlNetBeforeDenoiserStep(),
+ QwenImageControlNetDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "controlnet_input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_img2img_latents",
+ "prepare_rope_inputs",
+ "controlnet_before_denoise",
+ "controlnet_denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Before denoise step that prepare the inputs (timesteps, latents, rope inputs etc.) for the denoise step for img2img task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Auto denoise step for QwenImage
+class QwenImageAutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ block_classes = [
+ QwenImageCoreDenoiseStep,
+ QwenImageInpaintCoreDenoiseStep,
+ QwenImageImg2ImgCoreDenoiseStep,
+ QwenImageControlNetCoreDenoiseStep,
+ QwenImageControlNetInpaintCoreDenoiseStep,
+ QwenImageControlNetImg2ImgCoreDenoiseStep,
+ ]
+ block_names = [
+ "text2image",
+ "inpaint",
+ "img2img",
+ "controlnet_text2image",
+ "controlnet_inpaint",
+ "controlnet_img2img",
+ ]
+ block_trigger_inputs = ["control_image_latents", "processed_mask_image", "image_latents"]
+ default_block_name = "text2image"
+
+ def select_block(self, control_image_latents=None, processed_mask_image=None, image_latents=None):
+ if control_image_latents is not None:
+ if processed_mask_image is not None:
+ return "controlnet_inpaint"
+ elif image_latents is not None:
+ return "controlnet_img2img"
+ else:
+ return "controlnet_text2image"
+ else:
+ if processed_mask_image is not None:
+ return "inpaint"
+ elif image_latents is not None:
+ return "img2img"
+ else:
+ return "text2image"
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process. \n"
+ + " - `QwenImageCoreDenoiseStep` (text2image) for text2image tasks.\n"
+ + " - `QwenImageInpaintCoreDenoiseStep` (inpaint) for inpaint tasks.\n"
+ + " - `QwenImageImg2ImgCoreDenoiseStep` (img2img) for img2img tasks.\n"
+ + " - `QwenImageControlNetCoreDenoiseStep` (controlnet_text2image) for text2image tasks with controlnet.\n"
+ + " - `QwenImageControlNetInpaintCoreDenoiseStep` (controlnet_inpaint) for inpaint tasks with controlnet.\n"
+ + " - `QwenImageControlNetImg2ImgCoreDenoiseStep` (controlnet_img2img) for img2img tasks with controlnet.\n"
+ + "This step support text-to-image, image-to-image, inpainting, and controlnet tasks for QwenImage:\n"
+ + " - for image-to-image generation, you need to provide `image_latents`\n"
+ + " - for inpainting, you need to provide `processed_mask_image` and `image_latents`\n"
+ + " - to run the controlnet workflow, you need to provide `control_image_latents`\n"
+ + " - for text-to-image generation, all you need to provide is prompt embeddings"
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# ====================
+# 4. DECODE
+# ====================
+
+
+# standard decode step works for most tasks except for inpaint
+# auto_docstring
+class QwenImageDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode step that decodes the latents to images and postprocess the generated image.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self):
+ return "Decode step that decodes the latents to images and postprocess the generated image."
+
+
+# Inpaint decode step
+# auto_docstring
+class QwenImageInpaintDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode step that decodes the latents to images and postprocess the generated image, optional apply the mask
+ overally to the original image.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ mask_overlay_kwargs (`dict`, *optional*):
+ The kwargs for the postprocess step to apply the mask overlay. generated in
+ InpaintProcessImagesInputStep.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage"
+ block_classes = [QwenImageDecoderStep(), QwenImageInpaintProcessImagesOutputStep()]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self):
+ return "Decode step that decodes the latents to images and postprocess the generated image, optional apply the mask overally to the original image."
+
+
+# Auto decode step for QwenImage
+class QwenImageAutoDecodeStep(AutoPipelineBlocks):
+ block_classes = [QwenImageInpaintDecodeStep, QwenImageDecodeStep]
+ block_names = ["inpaint_decode", "decode"]
+ block_trigger_inputs = ["mask", None]
+
+ @property
+ def description(self):
+ return (
+ "Decode step that decode the latents into images. \n"
+ " This is an auto pipeline block that works for inpaint/text2image/img2img tasks, for both QwenImage and QwenImage-Edit.\n"
+ + " - `QwenImageInpaintDecodeStep` (inpaint) is used when `mask` is provided.\n"
+ + " - `QwenImageDecodeStep` (text2image/img2img) is used when `mask` is not provided.\n"
+ )
+
+
+# ====================
+# 5. AUTO BLOCKS & PRESETS
+# ====================
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", QwenImageAutoTextEncoderStep()),
+ ("vae_encoder", QwenImageAutoVaeEncoderStep()),
+ ("controlnet_vae_encoder", QwenImageOptionalControlNetVaeEncoderStep()),
+ ("denoise", QwenImageAutoCoreDenoiseStep()),
+ ("decode", QwenImageAutoDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class QwenImageAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `prompt`, `image`
+ - `inpainting`: requires `prompt`, `mask_image`, `image`
+ - `controlnet_text2image`: requires `prompt`, `control_image`
+ - `controlnet_image2image`: requires `prompt`, `image`, `control_image`
+ - `controlnet_inpainting`: requires `prompt`, `mask_image`, `image`, `control_image`
+
+ Components:
+ text_encoder (`Qwen2_5_VLForConditionalGeneration`): The text encoder to use tokenizer (`Qwen2Tokenizer`):
+ The tokenizer to use guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae
+ (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) controlnet (`QwenImageControlNetModel`)
+ control_image_processor (`VaeImageProcessor`) pachifier (`QwenImagePachifier`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 1024):
+ Maximum sequence length for prompt encoding.
+ mask_image (`Image`, *optional*):
+ Mask image for inpainting.
+ image (`Image | list`, *optional*):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ control_image (`Image`, *optional*):
+ Control image for ControlNet conditioning.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ image_latents (`Tensor`, *optional*):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ control_image_latents (`Tensor`, *optional*):
+ The control image latents to use for the denoising process. Can be generated in controlnet vae encoder
+ step.
+ control_guidance_start (`float`, *optional*, defaults to 0.0):
+ When to start applying ControlNet.
+ control_guidance_end (`float`, *optional*, defaults to 1.0):
+ When to stop applying ControlNet.
+ controlnet_conditioning_scale (`float`, *optional*, defaults to 1.0):
+ Scale for ControlNet conditioning.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ mask_overlay_kwargs (`dict`, *optional*):
+ The kwargs for the postprocess step to apply the mask overlay. generated in
+ InpaintProcessImagesInputStep.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage"
+
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+
+ # Workflow map defines the trigger conditions for each workflow.
+ # How to define:
+ # - Only include required inputs and trigger inputs (inputs that determine which blocks run)
+ # - currently, only supports `True` means the workflow triggers when the input is not None
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"prompt": True, "image": True},
+ "inpainting": {"prompt": True, "mask_image": True, "image": True},
+ "controlnet_text2image": {"prompt": True, "control_image": True},
+ "controlnet_image2image": {"prompt": True, "image": True, "control_image": True},
+ "controlnet_inpainting": {"prompt": True, "mask_image": True, "image": True, "control_image": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using QwenImage."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py
new file mode 100644
index 0000000000000000000000000000000000000000..0eff10432ebec540906d234b540398ffe825bbf2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit.py
@@ -0,0 +1,796 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+ QwenImageCreateMaskLatentsStep,
+ QwenImageEditRoPEInputsStep,
+ QwenImagePrepareLatentsStep,
+ QwenImagePrepareLatentsWithStrengthStep,
+ QwenImageSetTimestepsStep,
+ QwenImageSetTimestepsWithStrengthStep,
+)
+from .decoders import (
+ QwenImageAfterDenoiseStep,
+ QwenImageDecoderStep,
+ QwenImageInpaintProcessImagesOutputStep,
+ QwenImageProcessImagesOutputStep,
+)
+from .denoise import (
+ QwenImageEditDenoiseStep,
+ QwenImageEditInpaintDenoiseStep,
+)
+from .encoders import (
+ QwenImageEditInpaintProcessImagesInputStep,
+ QwenImageEditProcessImagesInputStep,
+ QwenImageEditResizeStep,
+ QwenImageEditTextEncoderStep,
+ QwenImageVaeEncoderStep,
+)
+from .inputs import (
+ QwenImageAdditionalInputsStep,
+ QwenImageTextInputsStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# ====================
+# 1. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageEditVLEncoderStep(SequentialPipelineBlocks):
+ """
+ QwenImage-Edit VL encoder step that encode the image and text prompts together.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditResizeStep(),
+ QwenImageEditTextEncoderStep(),
+ ]
+ block_names = ["resize", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "QwenImage-Edit VL encoder step that encode the image and text prompts together."
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# Edit VAE encoder
+# auto_docstring
+class QwenImageEditVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae
+ (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ processed_image (`Tensor`):
+ The processed image
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditResizeStep(),
+ QwenImageEditProcessImagesInputStep(),
+ QwenImageVaeEncoderStep(),
+ ]
+ block_names = ["resize", "preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that encode the image inputs into their latent representations."
+
+
+# Edit Inpaint VAE encoder
+# auto_docstring
+class QwenImageEditInpaintVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ This step is used for processing image and mask inputs for QwenImage-Edit inpaint tasks. It:
+ - resize the image for target area (1024 * 1024) while maintaining the aspect ratio.
+ - process the resized image and mask image.
+ - create image latents.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) vae
+ (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ mask_image (`Image`):
+ Mask image for inpainting.
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ processed_image (`Tensor`):
+ The processed image
+ processed_mask_image (`Tensor`):
+ The processed mask image
+ mask_overlay_kwargs (`dict`):
+ The kwargs for the postprocess step to apply the mask overlay
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditResizeStep(),
+ QwenImageEditInpaintProcessImagesInputStep(),
+ QwenImageVaeEncoderStep(),
+ ]
+ block_names = ["resize", "preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step is used for processing image and mask inputs for QwenImage-Edit inpaint tasks. It:\n"
+ " - resize the image for target area (1024 * 1024) while maintaining the aspect ratio.\n"
+ " - process the resized image and mask image.\n"
+ " - create image latents."
+ )
+
+
+# Auto VAE encoder
+class QwenImageEditAutoVaeEncoderStep(AutoPipelineBlocks):
+ block_classes = [QwenImageEditInpaintVaeEncoderStep, QwenImageEditVaeEncoderStep]
+ block_names = ["edit_inpaint", "edit"]
+ block_trigger_inputs = ["mask_image", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ "This is an auto pipeline block.\n"
+ " - `QwenImageEditInpaintVaeEncoderStep` (edit_inpaint) is used when `mask_image` is provided.\n"
+ " - `QwenImageEditVaeEncoderStep` (edit) is used when `image` is provided.\n"
+ " - if `mask_image` or `image` is not provided, step will be skipped."
+ )
+
+
+# ====================
+# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
+# ====================
+
+
+# assemble input steps
+# auto_docstring
+class QwenImageEditInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the edit denoising step. It:
+ - make sure the text embeddings have consistent batch size as well as the additional inputs.
+ - update height/width based `image_latents`, patchify `image_latents`.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+ batch-expanded)
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageAdditionalInputsStep(),
+ ]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the edit denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+ )
+
+
+# auto_docstring
+class QwenImageEditInpaintInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the edit inpaint denoising step. It:
+ - make sure the text embeddings have consistent batch size as well as the additional inputs.
+ - update height/width based `image_latents`, patchify `image_latents`.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+ batch-expanded)
+ processed_mask_image (`Tensor`):
+ The processed mask image (batch-expanded)
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageAdditionalInputsStep(
+ additional_batch_inputs=[
+ InputParam(name="processed_mask_image", type_hint=torch.Tensor, description="The processed mask image")
+ ]
+ ),
+ ]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the edit inpaint denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+ )
+
+
+# assemble prepare latents steps
+# auto_docstring
+class QwenImageEditInpaintPrepareLatentsStep(SequentialPipelineBlocks):
+ """
+ This step prepares the latents/image_latents and mask inputs for the edit inpainting denoising step. It:
+ - Add noise to the image latents to create the latents input for the denoiser.
+ - Create the patchified latents `mask` based on the processed mask image.
+
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ latents (`Tensor`):
+ The initial random noised, can be generated in prepare latent step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be
+ generated from vae encoder and updated in input step.)
+ timesteps (`Tensor`):
+ The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+ processed_mask_image (`Tensor`):
+ The processed mask to use for the inpainting process.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ dtype (`dtype`, *optional*, defaults to torch.float32):
+ The dtype of the model inputs, can be generated in input step.
+
+ Outputs:
+ initial_noise (`Tensor`):
+ The initial random noised used for inpainting denoising.
+ latents (`Tensor`):
+ The scaled noisy latents to use for inpainting/image-to-image denoising.
+ mask (`Tensor`):
+ The mask to use for the inpainting process.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [QwenImagePrepareLatentsWithStrengthStep(), QwenImageCreateMaskLatentsStep()]
+ block_names = ["add_noise_to_latents", "create_mask_latents"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "This step prepares the latents/image_latents and mask inputs for the edit inpainting denoising step. It:\n"
+ " - Add noise to the image latents to create the latents input for the denoiser.\n"
+ " - Create the patchified latents `mask` based on the processed mask image.\n"
+ )
+
+
+# Qwen Image Edit (image2image) core denoise step
+# auto_docstring
+class QwenImageEditCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for QwenImage-Edit edit (img2img) task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditInputStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsStep(),
+ QwenImageEditRoPEInputsStep(),
+ QwenImageEditDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Core denoising workflow for QwenImage-Edit edit (img2img) task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Qwen Image Edit (inpainting) core denoise step
+# auto_docstring
+class QwenImageEditInpaintCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for QwenImage-Edit edit inpaint task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditInpaintInputStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsWithStrengthStep(),
+ QwenImageEditInpaintPrepareLatentsStep(),
+ QwenImageEditRoPEInputsStep(),
+ QwenImageEditInpaintDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_inpaint_latents",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Core denoising workflow for QwenImage-Edit edit inpaint task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# Auto core denoise step for QwenImage Edit
+class QwenImageEditAutoCoreDenoiseStep(ConditionalPipelineBlocks):
+ model_name = "qwenimage-edit"
+ block_classes = [
+ QwenImageEditInpaintCoreDenoiseStep,
+ QwenImageEditCoreDenoiseStep,
+ ]
+ block_names = ["edit_inpaint", "edit"]
+ block_trigger_inputs = ["processed_mask_image", "image_latents"]
+ default_block_name = "edit"
+
+ def select_block(self, processed_mask_image=None, image_latents=None) -> str | None:
+ if processed_mask_image is not None:
+ return "edit_inpaint"
+ elif image_latents is not None:
+ return "edit"
+ return None
+
+ @property
+ def description(self):
+ return (
+ "Auto core denoising step that selects the appropriate workflow based on inputs.\n"
+ " - `QwenImageEditInpaintCoreDenoiseStep` when `processed_mask_image` is provided\n"
+ " - `QwenImageEditCoreDenoiseStep` when `image_latents` is provided\n"
+ "Supports edit (img2img) and edit inpainting tasks for QwenImage-Edit."
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# ====================
+# 4. DECODE
+# ====================
+
+
+# Decode step (standard)
+# auto_docstring
+class QwenImageEditDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode step that decodes the latents to images and postprocess the generated image.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self):
+ return "Decode step that decodes the latents to images and postprocess the generated image."
+
+
+# Inpaint decode step
+# auto_docstring
+class QwenImageEditInpaintDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask
+ overlay to the original image.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ mask_overlay_kwargs (`dict`, *optional*):
+ The kwargs for the postprocess step to apply the mask overlay. generated in
+ InpaintProcessImagesInputStep.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = [QwenImageDecoderStep(), QwenImageInpaintProcessImagesOutputStep()]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self):
+ return "Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask overlay to the original image."
+
+
+# Auto decode step
+class QwenImageEditAutoDecodeStep(AutoPipelineBlocks):
+ block_classes = [QwenImageEditInpaintDecodeStep, QwenImageEditDecodeStep]
+ block_names = ["inpaint_decode", "decode"]
+ block_trigger_inputs = ["mask", None]
+
+ @property
+ def description(self):
+ return (
+ "Decode step that decode the latents into images.\n"
+ "This is an auto pipeline block.\n"
+ " - `QwenImageEditInpaintDecodeStep` (inpaint) is used when `mask` is provided.\n"
+ " - `QwenImageEditDecodeStep` (edit) is used when `mask` is not provided.\n"
+ )
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# ====================
+# 5. AUTO BLOCKS & PRESETS
+# ====================
+
+EDIT_AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", QwenImageEditVLEncoderStep()),
+ ("vae_encoder", QwenImageEditAutoVaeEncoderStep()),
+ ("denoise", QwenImageEditAutoCoreDenoiseStep()),
+ ("decode", QwenImageEditAutoDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class QwenImageEditAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for edit (img2img) and edit inpaint tasks using QwenImage-Edit.
+ - for edit (img2img) generation, you need to provide `image`
+ - for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide
+ `padding_mask_crop`
+
+
+ Supported workflows:
+ - `image_conditioned`: requires `prompt`, `image`
+ - `image_conditioned_inpainting`: requires `prompt`, `mask_image`, `image`
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae
+ (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) pachifier (`QwenImagePachifier`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ mask_image (`Image`, *optional*):
+ Mask image for inpainting.
+ padding_mask_crop (`int`, *optional*):
+ Padding for mask cropping in inpainting.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ processed_mask_image (`Tensor`, *optional*):
+ The processed mask image
+ latents (`Tensor`):
+ Pre-generated noisy latents for image generation.
+ num_inference_steps (`int`):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ strength (`float`, *optional*, defaults to 0.9):
+ Strength for img2img/inpainting.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+ mask_overlay_kwargs (`dict`, *optional*):
+ The kwargs for the postprocess step to apply the mask overlay. generated in
+ InpaintProcessImagesInputStep.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage-edit"
+ block_classes = EDIT_AUTO_BLOCKS.values()
+ block_names = EDIT_AUTO_BLOCKS.keys()
+ _workflow_map = {
+ "image_conditioned": {"prompt": True, "image": True},
+ "image_conditioned_inpainting": {"prompt": True, "mask_image": True, "image": True},
+ }
+
+ @property
+ def description(self):
+ return (
+ "Auto Modular pipeline for edit (img2img) and edit inpaint tasks using QwenImage-Edit.\n"
+ "- for edit (img2img) generation, you need to provide `image`\n"
+ "- for edit inpainting, you need to provide `mask_image` and `image`, optionally you can provide `padding_mask_crop`\n"
+ )
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0e1eeaa1d012fa11ca6124724e54d0a4ca1c13a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_edit_plus.py
@@ -0,0 +1,407 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ QwenImageEditPlusRoPEInputsStep,
+ QwenImagePrepareLatentsStep,
+ QwenImageSetTimestepsStep,
+)
+from .decoders import (
+ QwenImageAfterDenoiseStep,
+ QwenImageDecoderStep,
+ QwenImageProcessImagesOutputStep,
+)
+from .denoise import (
+ QwenImageEditDenoiseStep,
+)
+from .encoders import (
+ QwenImageEditPlusProcessImagesInputStep,
+ QwenImageEditPlusResizeStep,
+ QwenImageEditPlusTextEncoderStep,
+ QwenImageVaeEncoderStep,
+)
+from .inputs import (
+ QwenImageEditPlusAdditionalInputsStep,
+ QwenImageTextInputsStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# ====================
+# 1. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks):
+ """
+ QwenImage-Edit Plus VL encoder step that encodes the image and text prompts together.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+
+ Outputs:
+ resized_image (`list`):
+ Images resized to 1024x1024 target area for VAE encoding
+ resized_cond_image (`list`):
+ Images resized to 384x384 target area for VL text encoding
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = [
+ QwenImageEditPlusResizeStep(),
+ QwenImageEditPlusTextEncoderStep(),
+ ]
+ block_names = ["resize", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "QwenImage-Edit Plus VL encoder step that encodes the image and text prompts together."
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ VAE encoder step that encodes image inputs into latent representations.
+ Each image is resized independently based on its own aspect ratio to 1024x1024 target area.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae
+ (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ resized_image (`list`):
+ Images resized to 1024x1024 target area for VAE encoding
+ resized_cond_image (`list`):
+ Images resized to 384x384 target area for VL text encoding
+ processed_image (`Tensor`):
+ The processed image
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = [
+ QwenImageEditPlusResizeStep(),
+ QwenImageEditPlusProcessImagesInputStep(),
+ QwenImageVaeEncoderStep(),
+ ]
+ block_names = ["resize", "preprocess", "encode"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "VAE encoder step that encodes image inputs into latent representations.\n"
+ "Each image is resized independently based on its own aspect ratio to 1024x1024 target area."
+ )
+
+
+# ====================
+# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
+# ====================
+
+
+# assemble input steps
+# auto_docstring
+class QwenImageEditPlusInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the Edit Plus denoising step. It:
+ - Standardizes text embeddings batch size.
+ - Processes list of image latents: patchifies, concatenates along dim=1, expands batch.
+ - Outputs lists of image_height/image_width for RoPE calculation.
+ - Defaults height/width from last image in the list.
+
+ Components:
+ pachifier (`QwenImagePachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`list`):
+ The image heights calculated from the image latents dimension
+ image_width (`list`):
+ The image widths calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified,
+ concatenated, and batch-expanded)
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageEditPlusAdditionalInputsStep(),
+ ]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the Edit Plus denoising step. It:\n"
+ " - Standardizes text embeddings batch size.\n"
+ " - Processes list of image latents: patchifies, concatenates along dim=1, expands batch.\n"
+ " - Outputs lists of image_height/image_width for RoPE calculation.\n"
+ " - Defaults height/width from last image in the list."
+ )
+
+
+# Qwen Image Edit Plus (image2image) core denoise step
+# auto_docstring
+class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for QwenImage-Edit Plus edit (img2img) task.
+
+ Components:
+ pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = [
+ QwenImageEditPlusInputStep(),
+ QwenImagePrepareLatentsStep(),
+ QwenImageSetTimestepsStep(),
+ QwenImageEditPlusRoPEInputsStep(),
+ QwenImageEditDenoiseStep(),
+ QwenImageAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Core denoising workflow for QwenImage-Edit Plus edit (img2img) task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# ====================
+# 4. DECODE
+# ====================
+
+
+# auto_docstring
+class QwenImageEditPlusDecodeStep(SequentialPipelineBlocks):
+ """
+ Decode step that decodes the latents to images and postprocesses the generated image.
+
+ Components:
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ latents (`Tensor`):
+ The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+ step.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images. (tensor output of the vae decoder.)
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()]
+ block_names = ["decode", "postprocess"]
+
+ @property
+ def description(self):
+ return "Decode step that decodes the latents to images and postprocesses the generated image."
+
+
+# ====================
+# 5. AUTO BLOCKS & PRESETS
+# ====================
+
+EDIT_PLUS_AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", QwenImageEditPlusVLEncoderStep()),
+ ("vae_encoder", QwenImageEditPlusVaeEncoderStep()),
+ ("denoise", QwenImageEditPlusCoreDenoiseStep()),
+ ("decode", QwenImageEditPlusDecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for edit (img2img) tasks using QwenImage-Edit Plus.
+ - `image` is required input (can be single image or list of images).
+ - Each image is resized independently based on its own aspect ratio.
+ - VL encoder uses 384x384 target area, VAE encoder uses 1024x1024 target area.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) guider (`ClassifierFreeGuidance`) image_processor (`VaeImageProcessor`) vae
+ (`AutoencoderKLQwenImage`) pachifier (`QwenImagePachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+ transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ prompt (`str`):
+ The prompt or prompts to guide image generation.
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage-edit-plus"
+ block_classes = EDIT_PLUS_AUTO_BLOCKS.values()
+ block_names = EDIT_PLUS_AUTO_BLOCKS.keys()
+
+ @property
+ def description(self):
+ return (
+ "Auto Modular pipeline for edit (img2img) tasks using QwenImage-Edit Plus.\n"
+ "- `image` is required input (can be single image or list of images).\n"
+ "- Each image is resized independently based on its own aspect ratio.\n"
+ "- VL encoder uses 384x384 target area, VAE encoder uses 1024x1024 target area."
+ )
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py
new file mode 100644
index 0000000000000000000000000000000000000000..f380667d2dd4b0dd3a8bc26c837a8b99d9a9c425
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_blocks_qwenimage_layered.py
@@ -0,0 +1,366 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ QwenImageLayeredPrepareLatentsStep,
+ QwenImageLayeredRoPEInputsStep,
+ QwenImageLayeredSetTimestepsStep,
+)
+from .decoders import (
+ QwenImageLayeredAfterDenoiseStep,
+ QwenImageLayeredDecoderStep,
+)
+from .denoise import (
+ QwenImageLayeredDenoiseStep,
+)
+from .encoders import (
+ QwenImageEditProcessImagesInputStep,
+ QwenImageLayeredGetImagePromptStep,
+ QwenImageLayeredPermuteLatentsStep,
+ QwenImageLayeredResizeStep,
+ QwenImageTextEncoderStep,
+ QwenImageVaeEncoderStep,
+)
+from .inputs import (
+ QwenImageLayeredAdditionalInputsStep,
+ QwenImageTextInputsStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# ====================
+# 1. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class QwenImageLayeredTextEncoderStep(SequentialPipelineBlocks):
+ """
+ QwenImage-Layered Text encoder step that encode the text prompt, will generate a prompt based on image if not
+ provided.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) tokenizer (`Qwen2Tokenizer`): The tokenizer to use guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ resolution (`int`, *optional*, defaults to 640):
+ The target area to resize the image to, can be 1024 or 640
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ use_en_prompt (`bool`, *optional*, defaults to False):
+ Whether to use English prompt template
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 1024):
+ Maximum sequence length for prompt encoding.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ prompt (`str`):
+ The prompt or prompts to guide image generation. If not provided, updated using image caption
+ prompt_embeds (`Tensor`):
+ The prompt embeddings.
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask.
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings.
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask.
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = [
+ QwenImageLayeredResizeStep(),
+ QwenImageLayeredGetImagePromptStep(),
+ QwenImageTextEncoderStep(),
+ ]
+ block_names = ["resize", "get_image_prompt", "encode"]
+
+ @property
+ def description(self) -> str:
+ return "QwenImage-Layered Text encoder step that encode the text prompt, will generate a prompt based on image if not provided."
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# Edit VAE encoder
+# auto_docstring
+class QwenImageLayeredVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Vae encoder step that encode the image inputs into their latent representations.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) image_processor (`VaeImageProcessor`) vae
+ (`AutoencoderKLQwenImage`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ resolution (`int`, *optional*, defaults to 640):
+ The target area to resize the image to, can be 1024 or 640
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+
+ Outputs:
+ resized_image (`list`):
+ The resized images
+ processed_image (`Tensor`):
+ The processed image
+ image_latents (`Tensor`):
+ The latent representation of the input image.
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = [
+ QwenImageLayeredResizeStep(),
+ QwenImageEditProcessImagesInputStep(),
+ QwenImageVaeEncoderStep(),
+ QwenImageLayeredPermuteLatentsStep(),
+ ]
+ block_names = ["resize", "preprocess", "encode", "permute"]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that encode the image inputs into their latent representations."
+
+
+# ====================
+# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
+# ====================
+
+
+# assemble input steps
+# auto_docstring
+class QwenImageLayeredInputStep(SequentialPipelineBlocks):
+ """
+ Input step that prepares the inputs for the layered denoising step. It:
+ - make sure the text embeddings have consistent batch size as well as the additional inputs.
+ - update height/width based `image_latents`, patchify `image_latents`.
+
+ Components:
+ pachifier (`QwenImageLayeredPachifier`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+ Outputs:
+ batch_size (`int`):
+ The batch size of the prompt embeddings
+ dtype (`dtype`):
+ The data type of the prompt embeddings
+ prompt_embeds (`Tensor`):
+ The prompt embeddings. (batch-expanded)
+ prompt_embeds_mask (`Tensor`):
+ The encoder attention mask. (batch-expanded)
+ negative_prompt_embeds (`Tensor`):
+ The negative prompt embeddings. (batch-expanded)
+ negative_prompt_embeds_mask (`Tensor`):
+ The negative prompt embeddings mask. (batch-expanded)
+ image_height (`int`):
+ The image height calculated from the image latents dimension
+ image_width (`int`):
+ The image width calculated from the image latents dimension
+ height (`int`):
+ if not provided, updated to image height
+ width (`int`):
+ if not provided, updated to image width
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified
+ with layered pachifier and batch-expanded)
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = [
+ QwenImageTextInputsStep(),
+ QwenImageLayeredAdditionalInputsStep(),
+ ]
+ block_names = ["text_inputs", "additional_inputs"]
+
+ @property
+ def description(self):
+ return (
+ "Input step that prepares the inputs for the layered denoising step. It:\n"
+ " - make sure the text embeddings have consistent batch size as well as the additional inputs.\n"
+ " - update height/width based `image_latents`, patchify `image_latents`."
+ )
+
+
+# Qwen Image Layered (image2image) core denoise step
+# auto_docstring
+class QwenImageLayeredCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Core denoising workflow for QwenImage-Layered img2img task.
+
+ Components:
+ pachifier (`QwenImageLayeredPachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ prompt_embeds_mask (`Tensor`):
+ mask for the text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+ negative_prompt_embeds_mask (`Tensor`, *optional*):
+ mask for the negative text embeddings. Can be generated from text_encoder step.
+ image_latents (`Tensor`):
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ layers (`int`, *optional*, defaults to 4):
+ Number of layers to extract from the image
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = [
+ QwenImageLayeredInputStep(),
+ QwenImageLayeredPrepareLatentsStep(),
+ QwenImageLayeredSetTimestepsStep(),
+ QwenImageLayeredRoPEInputsStep(),
+ QwenImageLayeredDenoiseStep(),
+ QwenImageLayeredAfterDenoiseStep(),
+ ]
+ block_names = [
+ "input",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_rope_inputs",
+ "denoise",
+ "after_denoise",
+ ]
+
+ @property
+ def description(self):
+ return "Core denoising workflow for QwenImage-Layered img2img task."
+
+ @property
+ def outputs(self):
+ return [
+ OutputParam.template("latents"),
+ ]
+
+
+# ====================
+# 4. AUTO BLOCKS & PRESETS
+# ====================
+
+LAYERED_AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", QwenImageLayeredTextEncoderStep()),
+ ("vae_encoder", QwenImageLayeredVaeEncoderStep()),
+ ("denoise", QwenImageLayeredCoreDenoiseStep()),
+ ("decode", QwenImageLayeredDecoderStep()),
+ ]
+)
+
+
+# auto_docstring
+class QwenImageLayeredAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for layered denoising tasks using QwenImage-Layered.
+
+ Components:
+ image_resize_processor (`VaeImageProcessor`) text_encoder (`Qwen2_5_VLForConditionalGeneration`) processor
+ (`Qwen2VLProcessor`) tokenizer (`Qwen2Tokenizer`): The tokenizer to use guider (`ClassifierFreeGuidance`)
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) pachifier (`QwenImageLayeredPachifier`)
+ scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`QwenImageTransformer2DModel`)
+
+ Inputs:
+ image (`Image | list`):
+ Reference image(s) for denoising. Can be a single image or list of images.
+ resolution (`int`, *optional*, defaults to 640):
+ The target area to resize the image to, can be 1024 or 640
+ prompt (`str`, *optional*):
+ The prompt or prompts to guide image generation.
+ use_en_prompt (`bool`, *optional*, defaults to False):
+ Whether to use English prompt template
+ negative_prompt (`str`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ max_sequence_length (`int`, *optional*, defaults to 1024):
+ Maximum sequence length for prompt encoding.
+ generator (`Generator`, *optional*):
+ Torch generator for deterministic generation.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ latents (`Tensor`, *optional*):
+ Pre-generated noisy latents for image generation.
+ layers (`int`, *optional*, defaults to 4):
+ Number of layers to extract from the image
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps.
+ sigmas (`list`, *optional*):
+ Custom sigmas for the denoising process.
+ attention_kwargs (`dict`, *optional*):
+ Additional kwargs for attention processors.
+ **denoiser_input_fields (`None`, *optional*):
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ Output format: 'pil', 'np', 'pt'.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "qwenimage-layered"
+ block_classes = LAYERED_AUTO_BLOCKS.values()
+ block_names = LAYERED_AUTO_BLOCKS.keys()
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for layered denoising tasks using QwenImage-Layered."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..f80ea5d1721f1c07bcbb0fe1a37deae8b52ea15b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/modular_pipeline.py
@@ -0,0 +1,297 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...configuration_utils import ConfigMixin, register_to_config
+from ...loaders import QwenImageLoraLoaderMixin
+from ..modular_pipeline import ModularPipeline
+
+
+class QwenImagePachifier(ConfigMixin):
+ """
+ A class to pack and unpack latents for QwenImage.
+ """
+
+ config_name = "config.json"
+
+ @register_to_config
+ def __init__(self, patch_size: int = 2):
+ super().__init__()
+
+ def pack_latents(self, latents):
+ if latents.ndim != 4 and latents.ndim != 5:
+ raise ValueError(f"Latents must have 4 or 5 dimensions, but got {latents.ndim}")
+
+ if latents.ndim == 4:
+ latents = latents.unsqueeze(2)
+
+ batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width = latents.shape
+ patch_size = self.config.patch_size
+
+ if latent_height % patch_size != 0 or latent_width % patch_size != 0:
+ raise ValueError(
+ f"Latent height and width must be divisible by {patch_size}, but got {latent_height} and {latent_width}"
+ )
+
+ latents = latents.view(
+ batch_size,
+ num_channels_latents,
+ latent_height // patch_size,
+ patch_size,
+ latent_width // patch_size,
+ patch_size,
+ )
+ latents = latents.permute(
+ 0, 2, 4, 1, 3, 5
+ ) # Batch_size, num_patches_height, num_patches_width, num_channels_latents, patch_size, patch_size
+ latents = latents.reshape(
+ batch_size,
+ (latent_height // patch_size) * (latent_width // patch_size),
+ num_channels_latents * patch_size * patch_size,
+ )
+
+ return latents
+
+ def unpack_latents(self, latents, height, width, vae_scale_factor=8):
+ if latents.ndim != 3:
+ raise ValueError(f"Latents must have 3 dimensions, but got {latents.ndim}")
+
+ batch_size, num_patches, channels = latents.shape
+ patch_size = self.config.patch_size
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = patch_size * (int(height) // (vae_scale_factor * patch_size))
+ width = patch_size * (int(width) // (vae_scale_factor * patch_size))
+
+ latents = latents.view(
+ batch_size,
+ height // patch_size,
+ width // patch_size,
+ channels // (patch_size * patch_size),
+ patch_size,
+ patch_size,
+ )
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+ latents = latents.reshape(batch_size, channels // (patch_size * patch_size), 1, height, width)
+
+ return latents
+
+
+class QwenImageLayeredPachifier(ConfigMixin):
+ """
+ A class to pack and unpack latents for QwenImage Layered.
+
+ Unlike QwenImagePachifier, this handles 5D latents with shape (B, layers+1, C, H, W).
+ """
+
+ config_name = "config.json"
+
+ @register_to_config
+ def __init__(self, patch_size: int = 2):
+ super().__init__()
+
+ def pack_latents(self, latents):
+ """
+ Pack latents from (B, layers, C, H, W) to (B, layers * H/2 * W/2, C*4).
+ """
+
+ if latents.ndim != 5:
+ raise ValueError(f"Latents must have 5 dimensions (B, layers, C, H, W), but got {latents.ndim}")
+
+ batch_size, layers, num_channels_latents, latent_height, latent_width = latents.shape
+ patch_size = self.config.patch_size
+
+ if latent_height % patch_size != 0 or latent_width % patch_size != 0:
+ raise ValueError(
+ f"Latent height and width must be divisible by {patch_size}, but got {latent_height} and {latent_width}"
+ )
+
+ latents = latents.view(
+ batch_size,
+ layers,
+ num_channels_latents,
+ latent_height // patch_size,
+ patch_size,
+ latent_width // patch_size,
+ patch_size,
+ )
+ latents = latents.permute(0, 1, 3, 5, 2, 4, 6)
+ latents = latents.reshape(
+ batch_size,
+ layers * (latent_height // patch_size) * (latent_width // patch_size),
+ num_channels_latents * patch_size * patch_size,
+ )
+ return latents
+
+ def unpack_latents(self, latents, height, width, layers, vae_scale_factor=8):
+ """
+ Unpack latents from (B, seq, C*4) to (B, C, layers+1, H, W).
+ """
+
+ if latents.ndim != 3:
+ raise ValueError(f"Latents must have 3 dimensions, but got {latents.ndim}")
+
+ batch_size, _, channels = latents.shape
+ patch_size = self.config.patch_size
+
+ height = patch_size * (int(height) // (vae_scale_factor * patch_size))
+ width = patch_size * (int(width) // (vae_scale_factor * patch_size))
+
+ latents = latents.view(
+ batch_size,
+ layers + 1,
+ height // patch_size,
+ width // patch_size,
+ channels // (patch_size * patch_size),
+ patch_size,
+ patch_size,
+ )
+ latents = latents.permute(0, 1, 4, 2, 5, 3, 6)
+ latents = latents.reshape(
+ batch_size,
+ layers + 1,
+ channels // (patch_size * patch_size),
+ height,
+ width,
+ )
+ latents = latents.permute(0, 2, 1, 3, 4) # (b, c, f, h, w)
+
+ return latents
+
+
+class QwenImageModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin):
+ """
+ A ModularPipeline for QwenImage.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "QwenImageAutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ return 128
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels // 4
+ return num_channels_latents
+
+ @property
+ def is_guidance_distilled(self):
+ is_guidance_distilled = False
+ if hasattr(self, "transformer") and self.transformer is not None:
+ is_guidance_distilled = self.transformer.config.guidance_embeds
+ return is_guidance_distilled
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
+
+
+class QwenImageEditModularPipeline(ModularPipeline, QwenImageLoraLoaderMixin):
+ """
+ A ModularPipeline for QwenImage-Edit.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "QwenImageEditAutoBlocks"
+
+ # YiYi TODO: qwen edit should not provide default height/width, should be derived from the resized input image (after adjustment) produced by the resize step.
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ return 128
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels // 4
+ return num_channels_latents
+
+ @property
+ def is_guidance_distilled(self):
+ is_guidance_distilled = False
+ if hasattr(self, "transformer") and self.transformer is not None:
+ is_guidance_distilled = self.transformer.config.guidance_embeds
+ return is_guidance_distilled
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
+
+
+class QwenImageEditPlusModularPipeline(QwenImageEditModularPipeline):
+ """
+ A ModularPipeline for QwenImage-Edit Plus.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "QwenImageEditPlusAutoBlocks"
+
+
+class QwenImageLayeredModularPipeline(QwenImageModularPipeline):
+ """
+ A ModularPipeline for QwenImage-Layered.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "QwenImageLayeredAutoBlocks"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a87fde024fc13ea3938715522a24d00580a29bf
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/qwenimage/prompt_templates.py
@@ -0,0 +1,121 @@
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Prompt templates for QwenImage pipelines.
+
+This module centralizes all prompt templates used across different QwenImage pipeline variants:
+- QwenImage (base): Text-only encoding for text-to-image generation
+- QwenImage Edit: VL encoding with single image for image editing
+- QwenImage Edit Plus: VL encoding with multiple images for multi-reference editing
+- QwenImage Layered: Auto-captioning for image decomposition
+"""
+
+# ============================================
+# QwenImage Base (text-only encoding)
+# ============================================
+# Used for text-to-image generation where only text prompt is encoded
+
+QWENIMAGE_PROMPT_TEMPLATE = (
+ "<|im_start|>system\n"
+ "Describe the image by detailing the color, shape, size, texture, quantity, text, "
+ "spatial relationships of the objects and background:<|im_end|>\n"
+ "<|im_start|>user\n{}<|im_end|>\n"
+ "<|im_start|>assistant\n"
+)
+QWENIMAGE_PROMPT_TEMPLATE_START_IDX = 34
+
+
+# ============================================
+# QwenImage Edit (VL encoding with single image)
+# ============================================
+# Used for single-image editing where both image and text are encoded together
+
+QWENIMAGE_EDIT_PROMPT_TEMPLATE = (
+ "<|im_start|>system\n"
+ "Describe the key features of the input image (color, shape, size, texture, objects, background), "
+ "then explain how the user's text instruction should alter or modify the image. "
+ "Generate a new image that meets the user's requirements while maintaining consistency "
+ "with the original input where appropriate.<|im_end|>\n"
+ "<|im_start|>user\n"
+ "<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n"
+ "<|im_start|>assistant\n"
+)
+QWENIMAGE_EDIT_PROMPT_TEMPLATE_START_IDX = 64
+
+
+# ============================================
+# QwenImage Edit Plus (VL encoding with multiple images)
+# ============================================
+# Used for multi-reference editing where multiple images and text are encoded together
+# The img_template is used to format each image in the prompt
+
+QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE = (
+ "<|im_start|>system\n"
+ "Describe the key features of the input image (color, shape, size, texture, objects, background), "
+ "then explain how the user's text instruction should alter or modify the image. "
+ "Generate a new image that meets the user's requirements while maintaining consistency "
+ "with the original input where appropriate.<|im_end|>\n"
+ "<|im_start|>user\n{}<|im_end|>\n"
+ "<|im_start|>assistant\n"
+)
+QWENIMAGE_EDIT_PLUS_IMG_TEMPLATE = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+QWENIMAGE_EDIT_PLUS_PROMPT_TEMPLATE_START_IDX = 64
+
+
+# ============================================
+# QwenImage Layered (auto-captioning)
+# ============================================
+# Used for image decomposition where the VL model generates a caption from the input image
+# if no prompt is provided. These prompts instruct the model to describe the image in detail.
+
+QWENIMAGE_LAYERED_CAPTION_PROMPT_EN = (
+ "<|im_start|>system\n"
+ "You are a helpful assistant.<|im_end|>\n"
+ "<|im_start|>user\n"
+ "# Image Annotator\n"
+ "You are a professional image annotator. Please write an image caption based on the input image:\n"
+ "1. Write the caption using natural, descriptive language without structured formats or rich text.\n"
+ "2. Enrich caption details by including:\n"
+ " - Object attributes, such as quantity, color, shape, size, material, state, position, actions, and so on\n"
+ " - Vision Relations between objects, such as spatial relations, functional relations, possessive relations, "
+ "attachment relations, action relations, comparative relations, causal relations, and so on\n"
+ " - Environmental details, such as weather, lighting, colors, textures, atmosphere, and so on\n"
+ " - Identify the text clearly visible in the image, without translation or explanation, "
+ "and highlight it in the caption with quotation marks\n"
+ "3. Maintain authenticity and accuracy:\n"
+ " - Avoid generalizations\n"
+ " - Describe all visible information in the image, while do not add information not explicitly shown in the image\n"
+ "<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n"
+ "<|im_start|>assistant\n"
+)
+
+QWENIMAGE_LAYERED_CAPTION_PROMPT_CN = (
+ "<|im_start|>system\n"
+ "You are a helpful assistant.<|im_end|>\n"
+ "<|im_start|>user\n"
+ "# 图像标注器\n"
+ "你是一个专业的图像标注器。请基于输入图像,撰写图注:\n"
+ "1. 使用自然、描述性的语言撰写图注,不要使用结构化形式或富文本形式。\n"
+ "2. 通过加入以下内容,丰富图注细节:\n"
+ " - 对象的属性:如数量、颜色、形状、大小、位置、材质、状态、动作等\n"
+ " - 对象间的视觉关系:如空间关系、功能关系、动作关系、从属关系、比较关系、因果关系等\n"
+ " - 环境细节:例如天气、光照、颜色、纹理、气氛等\n"
+ " - 文字内容:识别图像中清晰可见的文字,不做翻译和解释,用引号在图注中强调\n"
+ "3. 保持真实性与准确性:\n"
+ " - 不要使用笼统的描述\n"
+ " - 描述图像中所有可见的信息,但不要加入没有在图像中出现的内容\n"
+ "<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n"
+ "<|im_start|>assistant\n"
+)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..95b0d42aa38b1359ed6b097f9a03687ccf914ea5
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_stable_diffusion_3"] = ["StableDiffusion3AutoBlocks"]
+ _import_structure["modular_pipeline"] = ["StableDiffusion3ModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_stable_diffusion_3 import StableDiffusion3AutoBlocks
+ from .modular_pipeline import StableDiffusion3ModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..6efcaf4fc880e785637e662da5f66f5f1036c2ff
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/before_denoise.py
@@ -0,0 +1,457 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusion3ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+def _get_initial_timesteps_and_optionals(
+ transformer,
+ scheduler,
+ height,
+ width,
+ patch_size,
+ vae_scale_factor,
+ num_inference_steps,
+ sigmas,
+ device,
+ mu=None,
+):
+ scheduler_kwargs = {}
+ if scheduler.config.get("use_dynamic_shifting", None) and mu is None:
+ image_seq_len = (height // vae_scale_factor // patch_size) * (width // vae_scale_factor // patch_size)
+ mu = calculate_shift(
+ image_seq_len,
+ scheduler.config.get("base_image_seq_len", 256),
+ scheduler.config.get("max_image_seq_len", 4096),
+ scheduler.config.get("base_shift", 0.5),
+ scheduler.config.get("max_shift", 1.16),
+ )
+ scheduler_kwargs["mu"] = mu
+ elif mu is not None:
+ scheduler_kwargs["mu"] = mu
+
+ timesteps, num_inference_steps = retrieve_timesteps(
+ scheduler, num_inference_steps, device, sigmas=sigmas, **scheduler_kwargs
+ )
+ return timesteps, num_inference_steps
+
+
+class StableDiffusion3SetTimestepsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "num_inference_steps",
+ default=50,
+ description="The number of denoising steps.",
+ ),
+ InputParam(
+ "timesteps",
+ description="Custom timesteps to use for the denoising process.",
+ ),
+ InputParam("sigmas", description="Custom sigmas to use for the denoising process."),
+ InputParam(
+ "height",
+ type_hint=int,
+ description="The height in pixels of the generated image.",
+ ),
+ InputParam(
+ "width",
+ type_hint=int,
+ description="The width in pixels of the generated image.",
+ ),
+ InputParam(
+ "mu",
+ type_hint=float,
+ description="The mu value used for dynamic shifting. If not provided, it is dynamically calculated.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "timesteps",
+ type_hint=torch.Tensor,
+ description="The timesteps schedule for the denoising process.",
+ ),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The final number of inference steps.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ timesteps, num_inference_steps = _get_initial_timesteps_and_optionals(
+ components.transformer,
+ components.scheduler,
+ block_state.height,
+ block_state.width,
+ components.patch_size,
+ components.vae_scale_factor,
+ block_state.num_inference_steps,
+ block_state.sigmas,
+ block_state.device,
+ getattr(block_state, "mu", None),
+ )
+
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = num_inference_steps
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3Img2ImgSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for img2img inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "num_inference_steps",
+ default=50,
+ description="The number of denoising steps.",
+ ),
+ InputParam(
+ "timesteps",
+ description="Custom timesteps to use for the denoising process.",
+ ),
+ InputParam("sigmas", description="Custom sigmas to use for the denoising process."),
+ InputParam(
+ "strength",
+ default=0.6,
+ description="Indicates extent to transform the reference image.",
+ ),
+ InputParam(
+ "height",
+ type_hint=int,
+ description="The height in pixels of the generated image.",
+ ),
+ InputParam(
+ "width",
+ type_hint=int,
+ description="The width in pixels of the generated image.",
+ ),
+ InputParam(
+ "mu",
+ type_hint=float,
+ description="The mu value used for dynamic shifting. If not provided, it is dynamically calculated.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "timesteps",
+ type_hint=torch.Tensor,
+ description="The timesteps schedule for the denoising process.",
+ ),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The final number of inference steps.",
+ ),
+ ]
+
+ @staticmethod
+ def get_timesteps(scheduler, num_inference_steps, strength):
+ init_timestep = min(num_inference_steps * strength, num_inference_steps)
+ t_start = int(max(num_inference_steps - init_timestep, 0))
+ timesteps = scheduler.timesteps[t_start * scheduler.order :]
+ if hasattr(scheduler, "set_begin_index"):
+ scheduler.set_begin_index(t_start * scheduler.order)
+ return timesteps, num_inference_steps - t_start
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ timesteps, num_inference_steps = _get_initial_timesteps_and_optionals(
+ components.transformer,
+ components.scheduler,
+ block_state.height,
+ block_state.width,
+ components.patch_size,
+ components.vae_scale_factor,
+ block_state.num_inference_steps,
+ block_state.sigmas,
+ block_state.device,
+ getattr(block_state, "mu", None),
+ )
+
+ timesteps, num_inference_steps = self.get_timesteps(
+ components.scheduler, num_inference_steps, block_state.strength
+ )
+
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = num_inference_steps
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3PrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step for Text-to-Image"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "height",
+ type_hint=int,
+ description="The height in pixels of the generated image.",
+ ),
+ InputParam(
+ "width",
+ type_hint=int,
+ description="The width in pixels of the generated image.",
+ ),
+ InputParam(
+ "latents",
+ type_hint=torch.Tensor | None,
+ description="Pre-generated noisy latents to be used as inputs for image generation.",
+ ),
+ InputParam(
+ "num_images_per_prompt",
+ type_hint=int,
+ default=1,
+ description="The number of images to generate per prompt.",
+ ),
+ InputParam(
+ "generator",
+ description="One or a list of torch generator(s) to make generation deterministic.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="The batch size for latent generation.",
+ ),
+ InputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="The data type for the latents.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The prepared latent tensors to be denoised.",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+ batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ if block_state.latents is not None:
+ block_state.latents = block_state.latents.to(device=block_state.device, dtype=block_state.dtype)
+ else:
+ shape = (
+ batch_size,
+ components.num_channels_latents,
+ int(block_state.height) // components.vae_scale_factor,
+ int(block_state.width) // components.vae_scale_factor,
+ )
+ block_state.latents = randn_tensor(
+ shape,
+ generator=block_state.generator,
+ device=block_state.device,
+ dtype=block_state.dtype,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3Img2ImgPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to be scaled by the scheduler.",
+ ),
+ InputParam(
+ "image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The image latents encoded by the VAE.",
+ ),
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps schedule.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The noised latents prepared for denoising.",
+ ),
+ OutputParam(
+ "initial_noise",
+ type_hint=torch.Tensor,
+ description="The initial noise applied to the image latents.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
+ block_state.initial_noise = block_state.latents
+ block_state.latents = components.scheduler.scale_noise(
+ block_state.image_latents, latent_timestep, block_state.latents
+ )
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2ac1eab6fd7a45e4d405a57f2f1a2df7e4a33b2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/decoders.py
@@ -0,0 +1,79 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKL
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)
+
+
+class StableDiffusion3DecodeStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8, "vae_latent_channels": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "output_type",
+ default="pil",
+ description="The output format of the generated image (e.g., 'pil', 'pt', 'np').",
+ ),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents to be decoded.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("images", type_hint=list[PIL.Image.Image] | torch.Tensor)]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae = components.vae
+
+ if not block_state.output_type == "latent":
+ latents = (block_state.latents / vae.config.scaling_factor) + vae.config.shift_factor
+ block_state.images = vae.decode(latents, return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(
+ block_state.images, output_type=block_state.output_type
+ )
+ else:
+ block_state.images = block_state.latents
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7f9cba1a03494c070dd4d3cf62897f8dc9a28c0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/denoise.py
@@ -0,0 +1,231 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models.transformers import SD3Transformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusion3ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class StableDiffusion3LoopDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", SD3Transformer2DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step within the denoising loop that denoises the latents."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "joint_attention_kwargs",
+ type_hint=dict,
+ description="A kwargs dictionary passed along to the AttentionProcessor.",
+ ),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process.",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Text embeddings for guidance.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pooled text embeddings for guidance.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Negative text embeddings for guidance.",
+ ),
+ InputParam(
+ "negative_pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Negative pooled text embeddings for guidance.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ components: StableDiffusion3ModularPipeline,
+ block_state: BlockState,
+ i: int,
+ t: torch.Tensor,
+ ) -> PipelineState:
+ do_cfg = block_state.negative_prompt_embeds is not None
+
+ guider_inputs = {
+ "hidden_states": (block_state.latents, block_state.latents) if do_cfg else block_state.latents,
+ "encoder_hidden_states": (
+ block_state.prompt_embeds,
+ block_state.negative_prompt_embeds,
+ )
+ if do_cfg
+ else block_state.prompt_embeds,
+ "text_embeds": (
+ block_state.pooled_prompt_embeds,
+ block_state.negative_pooled_prompt_embeds,
+ )
+ if do_cfg
+ else block_state.pooled_prompt_embeds,
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+
+ latent_model_input = guider_state_batch.hidden_states
+ prompt_embeds = guider_state_batch.encoder_hidden_states
+ pooled_projections = getattr(guider_state_batch, "text_embeds", None)
+
+ timestep = t.expand(latent_model_input.shape[0])
+
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep,
+ encoder_hidden_states=prompt_embeds,
+ pooled_projections=pooled_projections,
+ joint_attention_kwargs=block_state.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ components.guider.cleanup_models(components.transformer)
+
+ guider_output = components.guider(guider_state)
+ block_state.noise_pred = guider_output.pred
+
+ return components, block_state
+
+
+class StableDiffusion3LoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents",
+ type_hint=torch.Tensor,
+ description="The denoised latent tensors.",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ components: StableDiffusion3ModularPipeline,
+ block_state: BlockState,
+ i: int,
+ t: torch.Tensor,
+ ):
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class StableDiffusion3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ComponentSpec("transformer", SD3Transformer2DModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam("timesteps", required=True, type_hint=torch.Tensor),
+ InputParam("num_inference_steps", required=True, type_hint=int),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order,
+ 0,
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3DenoiseStep(StableDiffusion3DenoiseLoopWrapper):
+ block_classes = [StableDiffusion3LoopDenoiser, StableDiffusion3LoopAfterDenoiser]
+ block_names = ["denoiser", "after_denoiser"]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..d6715448074df84a2ec55516aa144112a80930ee
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/encoders.py
@@ -0,0 +1,562 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from transformers import (
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ T5EncoderModel,
+ T5TokenizerFast,
+)
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...loaders import SD3LoraLoaderMixin
+from ...models import AutoencoderKL
+from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusion3ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def encode_vae_image(
+ vae: AutoencoderKL,
+ image: torch.Tensor,
+ generator: torch.Generator,
+ sample_mode="sample",
+):
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(
+ vae.encode(image[i : i + 1]),
+ generator=generator[i],
+ sample_mode=sample_mode,
+ )
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
+
+ image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor
+ return image_latents
+
+
+def _get_t5_prompt_embeds(
+ text_encoder: T5EncoderModel | None,
+ tokenizer: T5TokenizerFast | None,
+ prompt: str | list[str] = None,
+ max_sequence_length: int = 256,
+ device: torch.device | None = None,
+ joint_attention_dim: int = 4096,
+ dtype: torch.dtype | None = None,
+):
+ device = device or (text_encoder.device if text_encoder is not None else torch.device("cpu"))
+ dtype = dtype or (text_encoder.dtype if text_encoder is not None else torch.float32)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ if text_encoder is None or tokenizer is None:
+ return torch.zeros(
+ (batch_size, max_sequence_length, joint_attention_dim),
+ device=device,
+ dtype=dtype,
+ )
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+ logger.warning(
+ f"The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ prompt_embeds = text_encoder(text_input_ids.to(device))[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ return prompt_embeds
+
+
+def _get_clip_prompt_embeds(
+ text_encoder: CLIPTextModelWithProjection | None,
+ tokenizer: CLIPTokenizer | None,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ clip_skip: int | None = None,
+ hidden_size: int = 768,
+ dtype: torch.dtype | None = None,
+):
+ device = device or (text_encoder.device if text_encoder is not None else torch.device("cpu"))
+ dtype = dtype or (text_encoder.dtype if text_encoder is not None else torch.float32)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ if text_encoder is None or tokenizer is None:
+ prompt_embeds = torch.zeros((batch_size, 77, hidden_size), device=device, dtype=dtype)
+ pooled_prompt_embeds = torch.zeros((batch_size, hidden_size), device=device, dtype=dtype)
+ return prompt_embeds, pooled_prompt_embeds
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+ logger.warning(
+ f"The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {tokenizer.model_max_length} tokens: {removed_text}"
+ )
+ prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+ pooled_prompt_embeds = prompt_embeds[0]
+
+ if clip_skip is None:
+ prompt_embeds = prompt_embeds.hidden_states[-2]
+ else:
+ prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ return prompt_embeds, pooled_prompt_embeds
+
+
+def encode_prompt(
+ components,
+ prompt: str | list[str],
+ prompt_2: str | list[str] | None = None,
+ prompt_3: str | list[str] | None = None,
+ device: torch.device | None = None,
+ negative_prompt: str | list[str] | None = None,
+ negative_prompt_2: str | list[str] | None = None,
+ negative_prompt_3: str | list[str] | None = None,
+ clip_skip: int | None = None,
+ max_sequence_length: int = 256,
+ lora_scale: float | None = None,
+):
+ device = device or components._execution_device
+
+ expected_dtype = None
+ if components.text_encoder is not None:
+ expected_dtype = components.text_encoder.dtype
+ elif components.text_encoder_2 is not None:
+ expected_dtype = components.text_encoder_2.dtype
+ elif getattr(components, "transformer", None) is not None:
+ expected_dtype = components.transformer.dtype
+ else:
+ expected_dtype = torch.float32
+
+ if lora_scale is not None and isinstance(components, SD3LoraLoaderMixin):
+ components._lora_scale = lora_scale
+ if components.text_encoder is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(components.text_encoder, lora_scale)
+ if components.text_encoder_2 is not None and USE_PEFT_BACKEND:
+ scale_lora_layers(components.text_encoder_2, lora_scale)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ prompt_2 = prompt_2 or prompt
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+ prompt_3 = prompt_3 or prompt
+ prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+ prompt_embed, pooled_prompt_embed = _get_clip_prompt_embeds(
+ components.text_encoder,
+ components.tokenizer,
+ prompt=prompt,
+ device=device,
+ clip_skip=clip_skip,
+ hidden_size=768,
+ dtype=expected_dtype,
+ )
+ prompt_2_embed, pooled_prompt_2_embed = _get_clip_prompt_embeds(
+ components.text_encoder_2,
+ components.tokenizer_2,
+ prompt=prompt_2,
+ device=device,
+ clip_skip=clip_skip,
+ hidden_size=1280,
+ dtype=expected_dtype,
+ )
+ clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
+
+ t5_prompt_embed = _get_t5_prompt_embeds(
+ components.text_encoder_3,
+ components.tokenizer_3,
+ prompt=prompt_3,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ joint_attention_dim=(
+ components.transformer.config.joint_attention_dim
+ if getattr(components, "transformer", None) is not None
+ else 4096
+ ),
+ dtype=expected_dtype,
+ )
+
+ clip_prompt_embeds = torch.nn.functional.pad(
+ clip_prompt_embeds,
+ (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1]),
+ )
+ prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
+ pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+ negative_prompt_embeds = None
+ negative_pooled_prompt_embeds = None
+
+ if negative_prompt is not None or negative_prompt_2 is not None or negative_prompt_3 is not None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt_2 = negative_prompt_2 or negative_prompt
+ negative_prompt_3 = negative_prompt_3 or negative_prompt
+
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+ negative_prompt_2 = (
+ batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+ )
+ negative_prompt_3 = (
+ batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+ )
+
+ negative_prompt_embed, negative_pooled_prompt_embed = _get_clip_prompt_embeds(
+ components.text_encoder,
+ components.tokenizer,
+ prompt=negative_prompt,
+ device=device,
+ clip_skip=None,
+ hidden_size=768,
+ dtype=expected_dtype,
+ )
+ negative_prompt_2_embed, negative_pooled_prompt_2_embed = _get_clip_prompt_embeds(
+ components.text_encoder_2,
+ components.tokenizer_2,
+ prompt=negative_prompt_2,
+ device=device,
+ clip_skip=None,
+ hidden_size=1280,
+ dtype=expected_dtype,
+ )
+ negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
+
+ t5_negative_prompt_embed = _get_t5_prompt_embeds(
+ components.text_encoder_3,
+ components.tokenizer_3,
+ prompt=negative_prompt_3,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ joint_attention_dim=(
+ components.transformer.config.joint_attention_dim
+ if getattr(components, "transformer", None) is not None
+ else 4096
+ ),
+ dtype=expected_dtype,
+ )
+
+ negative_clip_prompt_embeds = torch.nn.functional.pad(
+ negative_clip_prompt_embeds,
+ (
+ 0,
+ t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1],
+ ),
+ )
+ negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
+ negative_pooled_prompt_embeds = torch.cat(
+ [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+ )
+
+ if components.text_encoder is not None and isinstance(components, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+ unscale_lora_layers(components.text_encoder, lora_scale)
+ if components.text_encoder_2 is not None and isinstance(components, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+ unscale_lora_layers(components.text_encoder_2, lora_scale)
+
+ return (
+ prompt_embeds,
+ negative_prompt_embeds,
+ pooled_prompt_embeds,
+ negative_pooled_prompt_embeds,
+ )
+
+
+class StableDiffusion3ProcessImagesInputStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def description(self) -> str:
+ return "Image Preprocess step for SD3."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8, "vae_latent_channels": 16}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "image",
+ description="The input image to be used as the starting point for the image-to-image process.",
+ ),
+ InputParam("height", description="The height in pixels of the generated image."),
+ InputParam("width", description="The width in pixels of the generated image."),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam(name="processed_image", description="The pre-processed image tensor.")]
+
+ @staticmethod
+ def check_inputs(height, width, vae_scale_factor, patch_size):
+ if height is not None and height % (vae_scale_factor * patch_size) != 0:
+ raise ValueError(f"Height must be divisible by {vae_scale_factor * patch_size} but is {height}")
+
+ if width is not None and width % (vae_scale_factor * patch_size) != 0:
+ raise ValueError(f"Width must be divisible by {vae_scale_factor * patch_size} but is {width}")
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState):
+ block_state = self.get_block_state(state)
+
+ if block_state.image is None:
+ raise ValueError("`image` cannot be None")
+
+ image = block_state.image
+ self.check_inputs(
+ height=block_state.height,
+ width=block_state.width,
+ vae_scale_factor=components.vae_scale_factor,
+ patch_size=components.patch_size,
+ )
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+
+ block_state.processed_image = components.image_processor.preprocess(image=image, height=height, width=width)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3VaeEncoderStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ def __init__(
+ self,
+ input_name: str = "processed_image",
+ output_name: str = "image_latents",
+ sample_mode: str = "sample",
+ ):
+ self._image_input_name = input_name
+ self._image_latents_output_name = output_name
+ self.sample_mode = sample_mode
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return f"Dynamic VAE Encoder step that converts {self._image_input_name} into latent representations {self._image_latents_output_name}."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [ComponentSpec("vae", AutoencoderKL)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ self._image_input_name,
+ description="The processed image input to be encoded.",
+ ),
+ InputParam(
+ "generator",
+ description="One or a list of torch generator(s) to make generation deterministic.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ self._image_latents_output_name,
+ type_hint=torch.Tensor,
+ description="The latents representing the reference image",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ image = getattr(block_state, self._image_input_name)
+
+ if image is None:
+ setattr(block_state, self._image_latents_output_name, None)
+ else:
+ device = components._execution_device
+ dtype = components.vae.dtype
+ image = image.to(device=device, dtype=dtype)
+ image_latents = encode_vae_image(
+ image=image,
+ vae=components.vae,
+ generator=block_state.generator,
+ sample_mode=self.sample_mode,
+ )
+ setattr(block_state, self._image_latents_output_name, image_latents)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3TextEncoderStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generates text embeddings to guide the image generation for SD3."
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", CLIPTextModelWithProjection),
+ ComponentSpec("tokenizer", CLIPTokenizer),
+ ComponentSpec("text_encoder_2", CLIPTextModelWithProjection),
+ ComponentSpec("tokenizer_2", CLIPTokenizer),
+ ComponentSpec("text_encoder_3", T5EncoderModel),
+ ComponentSpec("tokenizer_3", T5TokenizerFast),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "prompt",
+ description="The prompt or prompts to guide the image generation.",
+ ),
+ InputParam(
+ "prompt_2",
+ description="The prompt or prompts to be sent to tokenizer_2 and text_encoder_2.",
+ ),
+ InputParam(
+ "prompt_3",
+ description="The prompt or prompts to be sent to tokenizer_3 and text_encoder_3.",
+ ),
+ InputParam(
+ "negative_prompt",
+ description="The prompt or prompts not to guide the image generation.",
+ ),
+ InputParam(
+ "negative_prompt_2",
+ description="The prompt or prompts not to guide the image generation for tokenizer_2.",
+ ),
+ InputParam(
+ "negative_prompt_3",
+ description="The prompt or prompts not to guide the image generation for tokenizer_3.",
+ ),
+ InputParam(
+ "clip_skip",
+ type_hint=int,
+ description="Number of layers to be skipped from CLIP while computing the prompt embeddings.",
+ ),
+ InputParam(
+ "max_sequence_length",
+ type_hint=int,
+ default=256,
+ description="Maximum sequence length to use with the prompt.",
+ ),
+ InputParam(
+ "joint_attention_kwargs",
+ description="A kwargs dictionary passed along to the AttentionProcessor.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("prompt_embeds", type_hint=torch.Tensor),
+ OutputParam("negative_prompt_embeds", type_hint=torch.Tensor),
+ OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor),
+ OutputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ lora_scale = (
+ block_state.joint_attention_kwargs.get("scale", None) if block_state.joint_attention_kwargs else None
+ )
+
+ (
+ prompt_embeds,
+ negative_prompt_embeds,
+ pooled_prompt_embeds,
+ negative_pooled_prompt_embeds,
+ ) = encode_prompt(
+ components=components,
+ prompt=block_state.prompt,
+ prompt_2=block_state.prompt_2,
+ prompt_3=block_state.prompt_3,
+ device=block_state.device,
+ negative_prompt=block_state.negative_prompt,
+ negative_prompt_2=block_state.negative_prompt_2,
+ negative_prompt_3=block_state.negative_prompt_3,
+ clip_skip=block_state.clip_skip,
+ max_sequence_length=block_state.max_sequence_length,
+ lora_scale=lora_scale,
+ )
+
+ block_state.prompt_embeds = prompt_embeds
+ block_state.negative_prompt_embeds = negative_prompt_embeds
+ block_state.pooled_prompt_embeds = pooled_prompt_embeds
+ block_state.negative_pooled_prompt_embeds = negative_pooled_prompt_embeds
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/inputs.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/inputs.py
new file mode 100644
index 0000000000000000000000000000000000000000..581528a0351e4a88e3fec6fdfc3cfa7a8577881a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/inputs.py
@@ -0,0 +1,325 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import InputParam, OutputParam
+from .modular_pipeline import StableDiffusion3ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.inputs.repeat_tensor_to_batch_size
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_images_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_images_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_images_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_images_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_images_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.inputs.calculate_dimension_from_latents
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> tuple[int, int]:
+ """Calculate image dimensions from latent tensor dimensions.
+
+ This function converts latent space dimensions to image space dimensions by multiplying the latent height and width
+ by the VAE scale factor.
+
+ Args:
+ latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions.
+ Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width]
+ vae_scale_factor (int): The scale factor used by the VAE to compress images.
+ Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
+
+ Returns:
+ tuple[int, int]: The calculated image dimensions as (height, width)
+
+ Raises:
+ ValueError: If latents tensor doesn't have 4 or 5 dimensions
+
+ """
+ # make sure the latents are not packed
+ if latents.ndim != 4 and latents.ndim != 5:
+ raise ValueError(f"unpacked latents must have 4 or 5 dimensions, but got {latents.ndim}")
+
+ latent_height, latent_width = latents.shape[-2:]
+
+ height = latent_height * vae_scale_factor
+ width = latent_width * vae_scale_factor
+
+ return height, width
+
+
+class StableDiffusion3TextInputStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Text input processing step that standardizes text embeddings for SD3, applying CFG duplication if needed."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "num_images_per_prompt",
+ default=1,
+ description="The number of images to generate per prompt.",
+ ),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pre-generated pooled text embeddings.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Pre-generated negative text embeddings.",
+ ),
+ InputParam(
+ "negative_pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Pre-generated negative pooled text embeddings.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="The batch size for the inference.",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="The expected data type for latents.",
+ ),
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ description="The processed text embeddings.",
+ ),
+ OutputParam(
+ "pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="The processed pooled text embeddings.",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="The processed negative text embeddings.",
+ ),
+ OutputParam(
+ "negative_pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="The processed negative pooled text embeddings.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1)
+
+ pooled_prompt_embeds = block_state.pooled_prompt_embeds.repeat(1, block_state.num_images_per_prompt)
+ pooled_prompt_embeds = pooled_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, -1
+ )
+
+ if getattr(block_state, "negative_prompt_embeds", None) is not None:
+ _, neg_seq_len, _ = block_state.negative_prompt_embeds.shape
+ negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt,
+ neg_seq_len,
+ -1,
+ )
+
+ negative_pooled_prompt_embeds = block_state.negative_pooled_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt
+ )
+ negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, -1
+ )
+
+ block_state.negative_prompt_embeds = negative_prompt_embeds
+ block_state.negative_pooled_prompt_embeds = negative_pooled_prompt_embeds
+ else:
+ block_state.negative_prompt_embeds = None
+ block_state.negative_pooled_prompt_embeds = None
+
+ block_state.prompt_embeds = prompt_embeds
+ block_state.pooled_prompt_embeds = pooled_prompt_embeds
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusion3AdditionalInputsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-3"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[str] = ["image_latents"],
+ additional_batch_inputs: list[str] = [],
+ ):
+ self._image_latent_inputs = (
+ image_latent_inputs if isinstance(image_latent_inputs, list) else [image_latent_inputs]
+ )
+ self._additional_batch_inputs = (
+ additional_batch_inputs if isinstance(additional_batch_inputs, list) else [additional_batch_inputs]
+ )
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ return "Updates height/width if None, and expands batch size. SD3 does not pack latents on pipeline level."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(
+ "num_images_per_prompt",
+ default=1,
+ description="The number of images to generate per prompt.",
+ ),
+ InputParam("batch_size", required=True, description="The batch size."),
+ InputParam("height", description="The height in pixels of the generated image."),
+ InputParam("width", description="The width in pixels of the generated image."),
+ ]
+ for name in self._image_latent_inputs + self._additional_batch_inputs:
+ inputs.append(InputParam(name, description=f"Latent input {name} to be processed."))
+ return inputs
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_height",
+ type_hint=int,
+ description="The height of the generated image.",
+ ),
+ OutputParam(
+ "image_width",
+ type_hint=int,
+ description="The width of the generated image.",
+ ),
+ ]
+
+ def __call__(self, components: StableDiffusion3ModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ for input_name in self._image_latent_inputs:
+ tensor = getattr(block_state, input_name)
+ if tensor is None:
+ continue
+
+ height, width = calculate_dimension_from_latents(tensor, components.vae_scale_factor)
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ if not hasattr(block_state, "image_height"):
+ block_state.image_height = height
+ if not hasattr(block_state, "image_width"):
+ block_state.image_width = width
+
+ tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+ setattr(block_state, input_name, tensor)
+
+ for input_name in self._additional_batch_inputs:
+ tensor = getattr(block_state, input_name)
+ if tensor is None:
+ continue
+ tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+ setattr(block_state, input_name, tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py
new file mode 100644
index 0000000000000000000000000000000000000000..7141bdcede1c328c8840c496eadcf58c7f90108a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_blocks_stable_diffusion_3.py
@@ -0,0 +1,366 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+ StableDiffusion3Img2ImgPrepareLatentsStep,
+ StableDiffusion3Img2ImgSetTimestepsStep,
+ StableDiffusion3PrepareLatentsStep,
+ StableDiffusion3SetTimestepsStep,
+)
+from .decoders import StableDiffusion3DecodeStep
+from .denoise import StableDiffusion3DenoiseStep
+from .encoders import (
+ StableDiffusion3ProcessImagesInputStep,
+ StableDiffusion3TextEncoderStep,
+ StableDiffusion3VaeEncoderStep,
+)
+from .inputs import StableDiffusion3AdditionalInputsStep, StableDiffusion3TextInputStep
+
+
+logger = logging.get_logger(__name__)
+
+
+# auto_docstring
+class StableDiffusion3Img2ImgVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ The input image to be used as the starting point for the image-to-image process.
+ height (`None`, *optional*):
+ The height in pixels of the generated image.
+ width (`None`, *optional*):
+ The width in pixels of the generated image.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+
+ Outputs:
+ processed_image (`None`):
+ The pre-processed image tensor.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = [
+ StableDiffusion3ProcessImagesInputStep(),
+ StableDiffusion3VaeEncoderStep(),
+ ]
+ block_names = ["preprocess", "encode"]
+
+
+# auto_docstring
+class StableDiffusion3AutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Components:
+ image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`)
+
+ Inputs:
+ image (`None`, *optional*):
+ The input image to be used as the starting point for the image-to-image process.
+ height (`None`, *optional*):
+ The height in pixels of the generated image.
+ width (`None`, *optional*):
+ The width in pixels of the generated image.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+
+ Outputs:
+ processed_image (`None`):
+ The pre-processed image tensor.
+ image_latents (`Tensor`):
+ The latents representing the reference image
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = [StableDiffusion3Img2ImgVaeEncoderStep]
+ block_names = ["img2img"]
+ block_trigger_inputs = ["image"]
+
+
+# auto_docstring
+class StableDiffusion3T2ICoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer
+ (`SD3Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings.
+ pooled_prompt_embeds (`Tensor`):
+ Pre-generated pooled text embeddings.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings.
+ negative_pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings.
+ height (`int`, *optional*):
+ The height in pixels of the generated image.
+ width (`int`, *optional*):
+ The width in pixels of the generated image.
+ latents (`Tensor | NoneType`, *optional*):
+ Pre-generated noisy latents to be used as inputs for image generation.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`None`, *optional*):
+ Custom timesteps to use for the denoising process.
+ sigmas (`None`, *optional*):
+ Custom sigmas to use for the denoising process.
+ mu (`float`, *optional*):
+ The mu value used for dynamic shifting. If not provided, it is dynamically calculated.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary passed along to the AttentionProcessor.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = [
+ StableDiffusion3TextInputStep(),
+ StableDiffusion3PrepareLatentsStep(),
+ StableDiffusion3SetTimestepsStep(),
+ StableDiffusion3DenoiseStep(),
+ ]
+ block_names = ["text_inputs", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class StableDiffusion3I2ICoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer
+ (`SD3Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings.
+ pooled_prompt_embeds (`Tensor`):
+ Pre-generated pooled text embeddings.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings.
+ negative_pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings.
+ height (`None`, *optional*):
+ The height in pixels of the generated image.
+ width (`None`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`None`, *optional*):
+ Latent input image_latents to be processed.
+ latents (`Tensor | NoneType`, *optional*):
+ Pre-generated noisy latents to be used as inputs for image generation.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ The number of denoising steps.
+ timesteps (`None`, *optional*):
+ Custom timesteps to use for the denoising process.
+ sigmas (`None`, *optional*):
+ Custom sigmas to use for the denoising process.
+ strength (`None`, *optional*, defaults to 0.6):
+ Indicates extent to transform the reference image.
+ mu (`float`, *optional*):
+ The mu value used for dynamic shifting. If not provided, it is dynamically calculated.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary passed along to the AttentionProcessor.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = [
+ StableDiffusion3TextInputStep(),
+ StableDiffusion3AdditionalInputsStep(),
+ StableDiffusion3PrepareLatentsStep(),
+ StableDiffusion3Img2ImgSetTimestepsStep(),
+ StableDiffusion3Img2ImgPrepareLatentsStep(),
+ StableDiffusion3DenoiseStep(),
+ ]
+ block_names = [
+ "text_inputs",
+ "additional_inputs",
+ "prepare_latents",
+ "set_timesteps",
+ "prepare_img2img_latents",
+ "denoise",
+ ]
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class StableDiffusion3AutoCoreDenoiseStep(AutoPipelineBlocks):
+ """
+ Components:
+ scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer
+ (`SD3Transformer2DModel`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings.
+ pooled_prompt_embeds (`Tensor`):
+ Pre-generated pooled text embeddings.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings.
+ negative_pooled_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings.
+ height (`None`, *optional*):
+ The height in pixels of the generated image.
+ width (`None`, *optional*):
+ The width in pixels of the generated image.
+ image_latents (`None`, *optional*):
+ Latent input image_latents to be processed.
+ latents (`Tensor | NoneType`):
+ Pre-generated noisy latents to be used as inputs for image generation.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+ num_inference_steps (`None`):
+ The number of denoising steps.
+ timesteps (`None`):
+ Custom timesteps to use for the denoising process.
+ sigmas (`None`, *optional*):
+ Custom sigmas to use for the denoising process.
+ strength (`None`, *optional*, defaults to 0.6):
+ Indicates extent to transform the reference image.
+ mu (`float`, *optional*):
+ The mu value used for dynamic shifting. If not provided, it is dynamically calculated.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary passed along to the AttentionProcessor.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = [
+ StableDiffusion3I2ICoreDenoiseStep,
+ StableDiffusion3T2ICoreDenoiseStep,
+ ]
+ block_names = ["img2img", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+AUTO_BLOCKS = InsertableDict(
+ [
+ ("text_encoder", StableDiffusion3TextEncoderStep()),
+ ("vae_encoder", StableDiffusion3AutoVaeEncoderStep()),
+ ("denoise", StableDiffusion3AutoCoreDenoiseStep()),
+ ("decode", StableDiffusion3DecodeStep()),
+ ]
+)
+
+
+# auto_docstring
+class StableDiffusion3AutoBlocks(SequentialPipelineBlocks):
+ """
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModelWithProjection`) tokenizer (`CLIPTokenizer`) text_encoder_2
+ (`CLIPTextModelWithProjection`) tokenizer_2 (`CLIPTokenizer`) text_encoder_3 (`T5EncoderModel`) tokenizer_3
+ (`T5Tokenizer`) image_processor (`VaeImageProcessor`) vae (`AutoencoderKL`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer (`SD3Transformer2DModel`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ The prompt or prompts to guide the image generation.
+ prompt_2 (`None`, *optional*):
+ The prompt or prompts to be sent to tokenizer_2 and text_encoder_2.
+ prompt_3 (`None`, *optional*):
+ The prompt or prompts to be sent to tokenizer_3 and text_encoder_3.
+ negative_prompt (`None`, *optional*):
+ The prompt or prompts not to guide the image generation.
+ negative_prompt_2 (`None`, *optional*):
+ The prompt or prompts not to guide the image generation for tokenizer_2.
+ negative_prompt_3 (`None`, *optional*):
+ The prompt or prompts not to guide the image generation for tokenizer_3.
+ clip_skip (`int`, *optional*):
+ Number of layers to be skipped from CLIP while computing the prompt embeddings.
+ max_sequence_length (`int`, *optional*, defaults to 256):
+ Maximum sequence length to use with the prompt.
+ joint_attention_kwargs (`None`, *optional*):
+ A kwargs dictionary passed along to the AttentionProcessor.
+ image (`None`, *optional*):
+ The input image to be used as the starting point for the image-to-image process.
+ height (`None`, *optional*):
+ The height in pixels of the generated image.
+ width (`None`, *optional*):
+ The width in pixels of the generated image.
+ generator (`None`, *optional*):
+ One or a list of torch generator(s) to make generation deterministic.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ image_latents (`None`, *optional*):
+ Latent input image_latents to be processed.
+ latents (`Tensor | NoneType`):
+ Pre-generated noisy latents to be used as inputs for image generation.
+ num_inference_steps (`None`):
+ The number of denoising steps.
+ timesteps (`None`):
+ Custom timesteps to use for the denoising process.
+ sigmas (`None`, *optional*):
+ Custom sigmas to use for the denoising process.
+ strength (`None`, *optional*, defaults to 0.6):
+ Indicates extent to transform the reference image.
+ mu (`float`, *optional*):
+ The mu value used for dynamic shifting. If not provided, it is dynamically calculated.
+ output_type (`None`, *optional*, defaults to pil):
+ The output format of the generated image (e.g., 'pil', 'pt', 'np').
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ model_name = "stable-diffusion-3"
+ block_classes = AUTO_BLOCKS.values()
+ block_names = AUTO_BLOCKS.keys()
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ }
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..428b20aa6f258b3f9e030482cd032f73ff3f73c5
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_3/modular_pipeline.py
@@ -0,0 +1,69 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+
+class StableDiffusion3ModularPipeline(ModularPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin):
+ """
+ A ModularPipeline for Stable Diffusion 3.
+
+ >[!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "StableDiffusion3AutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ if getattr(self, "transformer", None) is not None:
+ return self.transformer.config.sample_size
+ return 128
+
+ @property
+ def patch_size(self):
+ if getattr(self, "transformer", None) is not None:
+ return self.transformer.config.patch_size
+ return 2
+
+ @property
+ def tokenizer_max_length(self):
+ if getattr(self, "tokenizer", None) is not None:
+ return self.tokenizer.model_max_length
+ return 77
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if getattr(self, "vae", None) is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ if getattr(self, "transformer", None) is not None:
+ return self.transformer.config.in_channels
+ return 16
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2a21cfb50ed3b79574cb7a5fb92e2833c18fa11
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["StableDiffusionXLModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_stable_diffusion_xl import StableDiffusionXLAutoBlocks
+ from .modular_pipeline import StableDiffusionXLModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..a76a88d2fc0da3a0ba39f836f45bf99cc01920a3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/before_denoise.py
@@ -0,0 +1,1874 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any
+
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKL, ControlNetModel, ControlNetUnionModel, UNet2DConditionModel
+from ...models.controlnets.multicontrolnet import MultiControlNetModel
+from ...schedulers import EulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor, unwrap_module
+from ..modular_pipeline import (
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusionXLModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# TODO(yiyi, aryan): We need another step before text encoder to set the `num_inference_steps` attribute for guider so that
+# things like when to do guidance and how many conditions to be prepared can be determined. Currently, this is done by
+# always assuming you want to do guidance in the Guiders. So, negative embeddings are prepared regardless of what the
+# configuration of guider is.
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def prepare_latents_img2img(
+ vae, scheduler, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True
+):
+ if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+ raise ValueError(f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}")
+
+ image = image.to(device=device, dtype=dtype)
+
+ batch_size = batch_size * num_images_per_prompt
+
+ if image.shape[1] == 4:
+ init_latents = image
+
+ else:
+ latents_mean = latents_std = None
+ if hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None:
+ latents_mean = torch.tensor(vae.config.latents_mean).view(1, 4, 1, 1)
+ if hasattr(vae.config, "latents_std") and vae.config.latents_std is not None:
+ latents_std = torch.tensor(vae.config.latents_std).view(1, 4, 1, 1)
+ # make sure the VAE is in float32 mode, as it overflows in float16
+ if vae.config.force_upcast:
+ image = image.float()
+ vae.to(dtype=torch.float32)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ elif isinstance(generator, list):
+ if image.shape[0] < batch_size and batch_size % image.shape[0] == 0:
+ image = torch.cat([image] * (batch_size // image.shape[0]), dim=0)
+ elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0:
+ raise ValueError(
+ f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} "
+ )
+
+ init_latents = [
+ retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i]) for i in range(batch_size)
+ ]
+ init_latents = torch.cat(init_latents, dim=0)
+ else:
+ init_latents = retrieve_latents(vae.encode(image), generator=generator)
+
+ if vae.config.force_upcast:
+ vae.to(dtype)
+
+ init_latents = init_latents.to(dtype)
+ if latents_mean is not None and latents_std is not None:
+ latents_mean = latents_mean.to(device=device, dtype=dtype)
+ latents_std = latents_std.to(device=device, dtype=dtype)
+ init_latents = (init_latents - latents_mean) * vae.config.scaling_factor / latents_std
+ else:
+ init_latents = vae.config.scaling_factor * init_latents
+
+ if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+ # expand init_latents for batch_size
+ additional_image_per_prompt = batch_size // init_latents.shape[0]
+ init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+ elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+ raise ValueError(
+ f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+ )
+ else:
+ init_latents = torch.cat([init_latents], dim=0)
+
+ if add_noise:
+ shape = init_latents.shape
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ # get latents
+ init_latents = scheduler.add_noise(init_latents, noise, timestep)
+
+ latents = init_latents
+
+ return latents
+
+
+class StableDiffusionXLInputStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_images_per_prompt`\n\n"
+ "All input tensors are expected to have either batch_size=1 or match the batch_size\n"
+ "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n"
+ "have a final batch_size of batch_size * num_images_per_prompt."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pre-generated pooled text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "negative_pooled_prompt_embeds",
+ description="Pre-generated negative pooled text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ description="Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step.",
+ ),
+ InputParam(
+ "negative_ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ description="Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `prompt_embeds`)",
+ ),
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="negative text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="pooled text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="negative pooled text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="image embeddings for IP-Adapter",
+ ),
+ OutputParam(
+ "negative_ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ kwargs_type="denoiser_input_fields", # already in intermedites state but declare here again for denoiser_input_fields
+ description="negative image embeddings for IP-Adapter",
+ ),
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.prompt_embeds is not None and block_state.negative_prompt_embeds is not None:
+ if block_state.prompt_embeds.shape != block_state.negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {block_state.negative_prompt_embeds.shape}."
+ )
+
+ if block_state.prompt_embeds is not None and block_state.pooled_prompt_embeds is None:
+ raise ValueError(
+ "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+ )
+
+ if block_state.negative_prompt_embeds is not None and block_state.negative_pooled_prompt_embeds is None:
+ raise ValueError(
+ "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+ )
+
+ if block_state.ip_adapter_embeds is not None and not isinstance(block_state.ip_adapter_embeds, list):
+ raise ValueError("`ip_adapter_embeds` must be a list")
+
+ if block_state.negative_ip_adapter_embeds is not None and not isinstance(
+ block_state.negative_ip_adapter_embeds, list
+ ):
+ raise ValueError("`negative_ip_adapter_embeds` must be a list")
+
+ if block_state.ip_adapter_embeds is not None and block_state.negative_ip_adapter_embeds is not None:
+ for i, ip_adapter_embed in enumerate(block_state.ip_adapter_embeds):
+ if ip_adapter_embed.shape != block_state.negative_ip_adapter_embeds[i].shape:
+ raise ValueError(
+ "`ip_adapter_embeds` and `negative_ip_adapter_embeds` must have the same shape when passed directly, but"
+ f" got: `ip_adapter_embeds` {ip_adapter_embed.shape} != `negative_ip_adapter_embeds`"
+ f" {block_state.negative_ip_adapter_embeds[i].shape}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
+ )
+
+ block_state.pooled_prompt_embeds = block_state.pooled_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.pooled_prompt_embeds = block_state.pooled_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, -1
+ )
+
+ if block_state.negative_pooled_prompt_embeds is not None:
+ block_state.negative_pooled_prompt_embeds = block_state.negative_pooled_prompt_embeds.repeat(
+ 1, block_state.num_images_per_prompt, 1
+ )
+ block_state.negative_pooled_prompt_embeds = block_state.negative_pooled_prompt_embeds.view(
+ block_state.batch_size * block_state.num_images_per_prompt, -1
+ )
+
+ if block_state.ip_adapter_embeds is not None:
+ for i, ip_adapter_embed in enumerate(block_state.ip_adapter_embeds):
+ block_state.ip_adapter_embeds[i] = torch.cat(
+ [ip_adapter_embed] * block_state.num_images_per_prompt, dim=0
+ )
+
+ if block_state.negative_ip_adapter_embeds is not None:
+ for i, negative_ip_adapter_embed in enumerate(block_state.negative_ip_adapter_embeds):
+ block_state.negative_ip_adapter_embeds[i] = torch.cat(
+ [negative_ip_adapter_embed] * block_state.num_images_per_prompt, dim=0
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLImg2ImgSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step that sets the timesteps for the scheduler and determines the initial noise level (latent_timestep) for image-to-image/inpainting generation.\n"
+ + "The latent_timestep is calculated from the `strength` parameter - higher strength means starting from a noisier version of the input image."
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ InputParam("denoising_end"),
+ InputParam("strength", default=0.3),
+ InputParam("denoising_start"),
+ # YiYi TODO: do we need num_images_per_prompt here?
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time",
+ ),
+ OutputParam(
+ "latent_timestep",
+ type_hint=torch.Tensor,
+ description="The timestep that represents the initial noise level for image-to-image generation",
+ ),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps with self->components
+ def get_timesteps(components, num_inference_steps, strength, device, denoising_start=None):
+ # get the original timestep using init_timestep
+ if denoising_start is None:
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+ t_start = max(num_inference_steps - init_timestep, 0)
+
+ timesteps = components.scheduler.timesteps[t_start * components.scheduler.order :]
+ if hasattr(components.scheduler, "set_begin_index"):
+ components.scheduler.set_begin_index(t_start * components.scheduler.order)
+
+ return timesteps, num_inference_steps - t_start
+
+ else:
+ # Strength is irrelevant if we directly request a timestep to start at;
+ # that is, strength is determined by the denoising_start instead.
+ discrete_timestep_cutoff = int(
+ round(
+ components.scheduler.config.num_train_timesteps
+ - (denoising_start * components.scheduler.config.num_train_timesteps)
+ )
+ )
+
+ num_inference_steps = (components.scheduler.timesteps < discrete_timestep_cutoff).sum().item()
+ if components.scheduler.order == 2 and num_inference_steps % 2 == 0:
+ # if the scheduler is a 2nd order scheduler we might have to do +1
+ # because `num_inference_steps` might be even given that every timestep
+ # (except the highest one) is duplicated. If `num_inference_steps` is even it would
+ # mean that we cut the timesteps in the middle of the denoising step
+ # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1
+ # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
+ num_inference_steps = num_inference_steps + 1
+
+ # because t_n+1 >= t_n, we slice the timesteps starting from the end
+ t_start = len(components.scheduler.timesteps) - num_inference_steps
+ timesteps = components.scheduler.timesteps[t_start:]
+ if hasattr(components.scheduler, "set_begin_index"):
+ components.scheduler.set_begin_index(t_start)
+ return timesteps, num_inference_steps
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.device = components._execution_device
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ block_state.device,
+ block_state.timesteps,
+ block_state.sigmas,
+ )
+
+ def denoising_value_valid(dnv):
+ return isinstance(dnv, float) and 0 < dnv < 1
+
+ block_state.timesteps, block_state.num_inference_steps = self.get_timesteps(
+ components,
+ block_state.num_inference_steps,
+ block_state.strength,
+ block_state.device,
+ denoising_start=block_state.denoising_start
+ if denoising_value_valid(block_state.denoising_start)
+ else None,
+ )
+ block_state.latent_timestep = block_state.timesteps[:1].repeat(
+ block_state.batch_size * block_state.num_images_per_prompt
+ )
+
+ if (
+ block_state.denoising_end is not None
+ and isinstance(block_state.denoising_end, float)
+ and block_state.denoising_end > 0
+ and block_state.denoising_end < 1
+ ):
+ block_state.discrete_timestep_cutoff = int(
+ round(
+ components.scheduler.config.num_train_timesteps
+ - (block_state.denoising_end * components.scheduler.config.num_train_timesteps)
+ )
+ )
+ block_state.num_inference_steps = len(
+ list(filter(lambda ts: ts >= block_state.discrete_timestep_cutoff, block_state.timesteps))
+ )
+ block_state.timesteps = block_state.timesteps[: block_state.num_inference_steps]
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ InputParam("denoising_end"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"),
+ OutputParam(
+ "num_inference_steps",
+ type_hint=int,
+ description="The number of denoising steps to perform at inference time",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.device = components._execution_device
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ block_state.device,
+ block_state.timesteps,
+ block_state.sigmas,
+ )
+
+ if (
+ block_state.denoising_end is not None
+ and isinstance(block_state.denoising_end, float)
+ and block_state.denoising_end > 0
+ and block_state.denoising_end < 1
+ ):
+ block_state.discrete_timestep_cutoff = int(
+ round(
+ components.scheduler.config.num_train_timesteps
+ - (block_state.denoising_end * components.scheduler.config.num_train_timesteps)
+ )
+ )
+ block_state.num_inference_steps = len(
+ list(filter(lambda ts: ts >= block_state.discrete_timestep_cutoff, block_state.timesteps))
+ )
+ block_state.timesteps = block_state.timesteps[: block_state.num_inference_steps]
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusionXLInpaintPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the latents for the inpainting process"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("latents"),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("denoising_start"),
+ InputParam(
+ "strength",
+ default=0.9999,
+ description="Conceptually, indicates how much to transform the reference `image` (the masked portion of image for inpainting). Must be between 0 and 1. `image` "
+ "will be used as a starting point, adding more noise to it the larger the `strength`. The number of "
+ "denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will "
+ "be maximum and the denoising process will run for the full number of iterations specified in "
+ "`num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of "
+ "`denoising_start` being declared as an integer, the value of `strength` will be ignored.",
+ ),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ InputParam(
+ "latent_timestep",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timestep that represents the initial noise level for image-to-image/inpainting generation. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents representing the reference image for image-to-image/inpainting generation. Can be generated in vae_encode step.",
+ ),
+ InputParam(
+ "mask",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The mask for the inpainting generation. Can be generated in vae_encode step.",
+ ),
+ InputParam(
+ "masked_image_latents",
+ type_hint=torch.Tensor,
+ description="The masked image latents for the inpainting generation (only for inpainting-specific unet). Can be generated in vae_encode step.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ ),
+ OutputParam(
+ "noise",
+ type_hint=torch.Tensor,
+ description="The noise added to the image latents, used for inpainting generation",
+ ),
+ ]
+
+ # Modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline._encode_vae_image with self->components
+ # YiYi TODO: update the _encode_vae_image so that we can use #Coped from
+ @staticmethod
+ def _encode_vae_image(components, image: torch.Tensor, generator: torch.Generator):
+ latents_mean = latents_std = None
+ if hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None:
+ latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1)
+ if hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None:
+ latents_std = torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1)
+
+ dtype = image.dtype
+ if components.vae.config.force_upcast:
+ image = image.float()
+ components.vae.to(dtype=torch.float32)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(components.vae.encode(image[i : i + 1]), generator=generator[i])
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(components.vae.encode(image), generator=generator)
+
+ if components.vae.config.force_upcast:
+ components.vae.to(dtype)
+
+ image_latents = image_latents.to(dtype)
+ if latents_mean is not None and latents_std is not None:
+ latents_mean = latents_mean.to(device=image_latents.device, dtype=dtype)
+ latents_std = latents_std.to(device=image_latents.device, dtype=dtype)
+ image_latents = (image_latents - latents_mean) * components.vae.config.scaling_factor / latents_std
+ else:
+ image_latents = components.vae.config.scaling_factor * image_latents
+
+ return image_latents
+
+ # Modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_latents adding components as first argument
+ def prepare_latents_inpaint(
+ self,
+ components,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ image=None,
+ timestep=None,
+ is_strength_max=True,
+ add_noise=True,
+ ):
+ shape = (
+ batch_size,
+ num_channels_latents,
+ int(height) // components.vae_scale_factor,
+ int(width) // components.vae_scale_factor,
+ )
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ if (image is None or timestep is None) and not is_strength_max:
+ raise ValueError(
+ "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise."
+ "However, either the image or the noise timestep has not been provided."
+ )
+
+ if image.shape[1] == 4:
+ image_latents = image.to(device=device, dtype=dtype)
+ image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)
+ elif latents is None and not is_strength_max:
+ image = image.to(device=device, dtype=dtype)
+ image_latents = self._encode_vae_image(components, image=image, generator=generator)
+ image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)
+
+ if latents is None and add_noise:
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ # if strength is 1. then initialise the latents to noise, else initial to image + noise
+ latents = noise if is_strength_max else components.scheduler.add_noise(image_latents, noise, timestep)
+ # if pure noise then scale the initial latents by the Scheduler's init sigma
+ latents = latents * components.scheduler.init_noise_sigma if is_strength_max else latents
+ elif add_noise:
+ noise = latents.to(device)
+ latents = noise * components.scheduler.init_noise_sigma
+ else:
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ latents = image_latents.to(device)
+
+ outputs = (latents, noise, image_latents)
+
+ return outputs
+
+ # modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_mask_latents
+ # do not accept do_classifier_free_guidance
+ def prepare_mask_latents(
+ self, components, mask, masked_image, batch_size, height, width, dtype, device, generator
+ ):
+ # resize the mask to latents shape as we concatenate the mask to the latents
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+ # and half precision
+ mask = torch.nn.functional.interpolate(
+ mask, size=(height // components.vae_scale_factor, width // components.vae_scale_factor)
+ )
+ mask = mask.to(device=device, dtype=dtype)
+
+ # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
+ if mask.shape[0] < batch_size:
+ if not batch_size % mask.shape[0] == 0:
+ raise ValueError(
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
+ f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
+ " of masks that you pass is divisible by the total requested batch size."
+ )
+ mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
+
+ if masked_image is not None and masked_image.shape[1] == 4:
+ masked_image_latents = masked_image
+ else:
+ masked_image_latents = None
+
+ if masked_image is not None:
+ if masked_image_latents is None:
+ masked_image = masked_image.to(device=device, dtype=dtype)
+ masked_image_latents = self._encode_vae_image(components, masked_image, generator=generator)
+
+ if masked_image_latents.shape[0] < batch_size:
+ if not batch_size % masked_image_latents.shape[0] == 0:
+ raise ValueError(
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
+ )
+ masked_image_latents = masked_image_latents.repeat(
+ batch_size // masked_image_latents.shape[0], 1, 1, 1
+ )
+
+ # aligning device to prevent device errors when concating it with the latent model input
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+
+ return mask, masked_image_latents
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype
+ block_state.device = components._execution_device
+
+ block_state.is_strength_max = block_state.strength == 1.0
+
+ # for non-inpainting specific unet, we do not need masked_image_latents
+ if hasattr(components, "unet") and components.unet is not None:
+ if components.unet.config.in_channels == 4:
+ block_state.masked_image_latents = None
+
+ block_state.add_noise = True if block_state.denoising_start is None else False
+
+ block_state.height = block_state.image_latents.shape[-2] * components.vae_scale_factor
+ block_state.width = block_state.image_latents.shape[-1] * components.vae_scale_factor
+
+ block_state.latents, block_state.noise, block_state.image_latents = self.prepare_latents_inpaint(
+ components,
+ block_state.batch_size * block_state.num_images_per_prompt,
+ components.num_channels_latents,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ block_state.latents,
+ image=block_state.image_latents,
+ timestep=block_state.latent_timestep,
+ is_strength_max=block_state.is_strength_max,
+ add_noise=block_state.add_noise,
+ )
+
+ # 7. Prepare mask latent variables
+ block_state.mask, block_state.masked_image_latents = self.prepare_mask_latents(
+ components,
+ block_state.mask,
+ block_state.masked_image_latents,
+ block_state.batch_size * block_state.num_images_per_prompt,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLImg2ImgPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the latents for the image-to-image generation process"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("latents"),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("denoising_start"),
+ InputParam("generator"),
+ InputParam(
+ "latent_timestep",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timestep that represents the initial noise level for image-to-image/inpainting generation. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "image_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The latents representing the reference image for image-to-image/inpainting generation. Can be generated in vae_encode step.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ InputParam("dtype", required=True, type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype
+ block_state.device = components._execution_device
+ block_state.add_noise = True if block_state.denoising_start is None else False
+ if block_state.latents is None:
+ block_state.latents = prepare_latents_img2img(
+ components.vae,
+ components.scheduler,
+ block_state.image_latents,
+ block_state.latent_timestep,
+ block_state.batch_size,
+ block_state.num_images_per_prompt,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ block_state.add_noise,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ComponentSpec("vae", AutoencoderKL),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the latents for the text-to-image generation process"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("latents"),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (
+ block_state.height is not None
+ and block_state.height % components.vae_scale_factor != 0
+ or block_state.width is not None
+ and block_state.width % components.vae_scale_factor != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor} but are {block_state.height} and {block_state.width}."
+ )
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents with self->comp
+ def prepare_latents(comp, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+ shape = (
+ batch_size,
+ num_channels_latents,
+ int(height) // comp.vae_scale_factor,
+ int(width) // comp.vae_scale_factor,
+ )
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ latents = latents.to(device)
+
+ # scale the initial noise by the standard deviation required by the scheduler
+ latents = latents * comp.scheduler.init_noise_sigma
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ if block_state.dtype is None:
+ block_state.dtype = components.vae.dtype
+
+ block_state.device = components._execution_device
+
+ self.check_inputs(components, block_state)
+
+ block_state.height = block_state.height or components.default_sample_size * components.vae_scale_factor
+ block_state.width = block_state.width or components.default_sample_size * components.vae_scale_factor
+ block_state.num_channels_latents = components.num_channels_latents
+ block_state.latents = self.prepare_latents(
+ components,
+ block_state.batch_size * block_state.num_images_per_prompt,
+ block_state.num_channels_latents,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ block_state.latents,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec("requires_aesthetics_score", False),
+ ]
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("unet", UNet2DConditionModel),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the additional conditioning for the image-to-image/inpainting generation process"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("original_size"),
+ InputParam("target_size"),
+ InputParam("negative_original_size"),
+ InputParam("negative_target_size"),
+ InputParam("crops_coords_top_left", default=(0, 0)),
+ InputParam("negative_crops_coords_top_left", default=(0, 0)),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam("aesthetic_score", default=6.0),
+ InputParam("negative_aesthetic_score", default=2.0),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The pooled prompt embeddings to use for the denoising process (used to determine shapes and dtypes for other additional conditioning inputs). Can be generated in text_encoder step.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "add_time_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The time ids to condition the denoising process",
+ ),
+ OutputParam(
+ "negative_add_time_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The negative time ids to condition the denoising process",
+ ),
+ OutputParam("timestep_cond", type_hint=torch.Tensor, description="The timestep cond to use for LCM"),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids with self->components
+ def _get_add_time_ids(
+ components,
+ original_size,
+ crops_coords_top_left,
+ target_size,
+ aesthetic_score,
+ negative_aesthetic_score,
+ negative_original_size,
+ negative_crops_coords_top_left,
+ negative_target_size,
+ dtype,
+ text_encoder_projection_dim=None,
+ ):
+ if components.config.requires_aesthetics_score:
+ add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
+ add_neg_time_ids = list(
+ negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)
+ )
+ else:
+ add_time_ids = list(original_size + crops_coords_top_left + target_size)
+ add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)
+
+ passed_add_embed_dim = (
+ components.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+ )
+ expected_add_embed_dim = components.unet.add_embedding.linear_1.in_features
+
+ if (
+ expected_add_embed_dim > passed_add_embed_dim
+ and (expected_add_embed_dim - passed_add_embed_dim) == components.unet.config.addition_time_embed_dim
+ ):
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."
+ )
+ elif (
+ expected_add_embed_dim < passed_add_embed_dim
+ and (passed_add_embed_dim - expected_add_embed_dim) == components.unet.config.addition_time_embed_dim
+ ):
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."
+ )
+ elif expected_add_embed_dim != passed_add_embed_dim:
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+ )
+
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+ add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)
+
+ return add_time_ids, add_neg_time_ids
+
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+ def get_guidance_scale_embedding(
+ self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
+ ) -> torch.Tensor:
+ """
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+ Args:
+ w (`torch.Tensor`):
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+ embedding_dim (`int`, *optional*, defaults to 512):
+ Dimension of the embeddings to generate.
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+ Data type of the generated embeddings.
+
+ Returns:
+ `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+ """
+ assert len(w.shape) == 1
+ w = w * 1000.0
+
+ half_dim = embedding_dim // 2
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+ emb = w.to(dtype)[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1))
+ assert emb.shape == (w.shape[0], embedding_dim)
+ return emb
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ block_state.vae_scale_factor = components.vae_scale_factor
+
+ block_state.height, block_state.width = block_state.latents.shape[-2:]
+ block_state.height = block_state.height * block_state.vae_scale_factor
+ block_state.width = block_state.width * block_state.vae_scale_factor
+
+ block_state.original_size = block_state.original_size or (block_state.height, block_state.width)
+ block_state.target_size = block_state.target_size or (block_state.height, block_state.width)
+
+ block_state.text_encoder_projection_dim = int(block_state.pooled_prompt_embeds.shape[-1])
+
+ if block_state.negative_original_size is None:
+ block_state.negative_original_size = block_state.original_size
+ if block_state.negative_target_size is None:
+ block_state.negative_target_size = block_state.target_size
+
+ block_state.add_time_ids, block_state.negative_add_time_ids = self._get_add_time_ids(
+ components,
+ block_state.original_size,
+ block_state.crops_coords_top_left,
+ block_state.target_size,
+ block_state.aesthetic_score,
+ block_state.negative_aesthetic_score,
+ block_state.negative_original_size,
+ block_state.negative_crops_coords_top_left,
+ block_state.negative_target_size,
+ dtype=block_state.pooled_prompt_embeds.dtype,
+ text_encoder_projection_dim=block_state.text_encoder_projection_dim,
+ )
+ block_state.add_time_ids = block_state.add_time_ids.repeat(
+ block_state.batch_size * block_state.num_images_per_prompt, 1
+ ).to(device=block_state.device)
+ block_state.negative_add_time_ids = block_state.negative_add_time_ids.repeat(
+ block_state.batch_size * block_state.num_images_per_prompt, 1
+ ).to(device=block_state.device)
+
+ # Optionally get Guidance Scale Embedding for LCM
+ block_state.timestep_cond = None
+ if (
+ hasattr(components, "unet")
+ and components.unet is not None
+ and components.unet.config.time_cond_proj_dim is not None
+ ):
+ # TODO(yiyi, aryan): Ideally, this should be `embedded_guidance_scale` instead of pulling from guider. Guider scales should be different from this!
+ block_state.guidance_scale_tensor = torch.tensor(components.guider.guidance_scale - 1).repeat(
+ block_state.batch_size * block_state.num_images_per_prompt
+ )
+ block_state.timestep_cond = self.get_guidance_scale_embedding(
+ block_state.guidance_scale_tensor, embedding_dim=components.unet.config.time_cond_proj_dim
+ ).to(device=block_state.device, dtype=block_state.latents.dtype)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusionXLPrepareAdditionalConditioningStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return "Step that prepares the additional conditioning for the text-to-image generation process"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("unet", UNet2DConditionModel),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("original_size"),
+ InputParam("target_size"),
+ InputParam("negative_original_size"),
+ InputParam("negative_target_size"),
+ InputParam("crops_coords_top_left", default=(0, 0)),
+ InputParam("negative_crops_coords_top_left", default=(0, 0)),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "pooled_prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The pooled prompt embeddings to use for the denoising process (used to determine shapes and dtypes for other additional conditioning inputs). Can be generated in text_encoder step.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "add_time_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The time ids to condition the denoising process",
+ ),
+ OutputParam(
+ "negative_add_time_ids",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="The negative time ids to condition the denoising process",
+ ),
+ OutputParam("timestep_cond", type_hint=torch.Tensor, description="The timestep cond to use for LCM"),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids with self->components
+ def _get_add_time_ids(
+ components, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None
+ ):
+ add_time_ids = list(original_size + crops_coords_top_left + target_size)
+
+ passed_add_embed_dim = (
+ components.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+ )
+ expected_add_embed_dim = components.unet.add_embedding.linear_1.in_features
+
+ if expected_add_embed_dim != passed_add_embed_dim:
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+ )
+
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+ return add_time_ids
+
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+ def get_guidance_scale_embedding(
+ self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
+ ) -> torch.Tensor:
+ """
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+ Args:
+ w (`torch.Tensor`):
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+ embedding_dim (`int`, *optional*, defaults to 512):
+ Dimension of the embeddings to generate.
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+ Data type of the generated embeddings.
+
+ Returns:
+ `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+ """
+ assert len(w.shape) == 1
+ w = w * 1000.0
+
+ half_dim = embedding_dim // 2
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+ emb = w.to(dtype)[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1))
+ assert emb.shape == (w.shape[0], embedding_dim)
+ return emb
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.device = components._execution_device
+
+ block_state.height, block_state.width = block_state.latents.shape[-2:]
+ block_state.height = block_state.height * components.vae_scale_factor
+ block_state.width = block_state.width * components.vae_scale_factor
+
+ block_state.original_size = block_state.original_size or (block_state.height, block_state.width)
+ block_state.target_size = block_state.target_size or (block_state.height, block_state.width)
+
+ block_state.text_encoder_projection_dim = int(block_state.pooled_prompt_embeds.shape[-1])
+
+ block_state.add_time_ids = self._get_add_time_ids(
+ components,
+ block_state.original_size,
+ block_state.crops_coords_top_left,
+ block_state.target_size,
+ block_state.pooled_prompt_embeds.dtype,
+ text_encoder_projection_dim=block_state.text_encoder_projection_dim,
+ )
+ if block_state.negative_original_size is not None and block_state.negative_target_size is not None:
+ block_state.negative_add_time_ids = self._get_add_time_ids(
+ components,
+ block_state.negative_original_size,
+ block_state.negative_crops_coords_top_left,
+ block_state.negative_target_size,
+ block_state.pooled_prompt_embeds.dtype,
+ text_encoder_projection_dim=block_state.text_encoder_projection_dim,
+ )
+ else:
+ block_state.negative_add_time_ids = block_state.add_time_ids
+
+ block_state.add_time_ids = block_state.add_time_ids.repeat(
+ block_state.batch_size * block_state.num_images_per_prompt, 1
+ ).to(device=block_state.device)
+ block_state.negative_add_time_ids = block_state.negative_add_time_ids.repeat(
+ block_state.batch_size * block_state.num_images_per_prompt, 1
+ ).to(device=block_state.device)
+
+ # Optionally get Guidance Scale Embedding for LCM
+ block_state.timestep_cond = None
+ if (
+ hasattr(components, "unet")
+ and components.unet is not None
+ and components.unet.config.time_cond_proj_dim is not None
+ ):
+ # TODO(yiyi, aryan): Ideally, this should be `embedded_guidance_scale` instead of pulling from guider. Guider scales should be different from this!
+ block_state.guidance_scale_tensor = torch.tensor(components.guider.guidance_scale - 1).repeat(
+ block_state.batch_size * block_state.num_images_per_prompt
+ )
+ block_state.timestep_cond = self.get_guidance_scale_embedding(
+ block_state.guidance_scale_tensor, embedding_dim=components.unet.config.time_cond_proj_dim
+ ).to(device=block_state.device, dtype=block_state.latents.dtype)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusionXLControlNetInputStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("controlnet", ControlNetModel),
+ ComponentSpec(
+ "control_image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"do_convert_rgb": True, "do_normalize": False}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "step that prepare inputs for controlnet"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("control_image", required=True),
+ InputParam("control_guidance_start", default=0.0),
+ InputParam("control_guidance_end", default=1.0),
+ InputParam("controlnet_conditioning_scale", default=1.0),
+ InputParam("guess_mode", default=False),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "crops_coords",
+ type_hint=tuple[int] | None,
+ description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("controlnet_cond", type_hint=torch.Tensor, description="The processed control image"),
+ OutputParam(
+ "control_guidance_start", type_hint=list[float], description="The controlnet guidance start values"
+ ),
+ OutputParam(
+ "control_guidance_end", type_hint=list[float], description="The controlnet guidance end values"
+ ),
+ OutputParam(
+ "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values"
+ ),
+ OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"),
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
+ ]
+
+ # Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
+ # 1. return image without apply any guidance
+ # 2. add crops_coords and resize_mode to preprocess()
+ @staticmethod
+ def prepare_control_image(
+ components,
+ image,
+ width,
+ height,
+ batch_size,
+ num_images_per_prompt,
+ device,
+ dtype,
+ crops_coords=None,
+ ):
+ if crops_coords is not None:
+ image = components.control_image_processor.preprocess(
+ image, height=height, width=width, crops_coords=crops_coords, resize_mode="fill"
+ ).to(dtype=torch.float32)
+ else:
+ image = components.control_image_processor.preprocess(image, height=height, width=width).to(
+ dtype=torch.float32
+ )
+
+ image_batch_size = image.shape[0]
+ if image_batch_size == 1:
+ repeat_by = batch_size
+ else:
+ # image batch size is the same as prompt batch size
+ repeat_by = num_images_per_prompt
+
+ image = image.repeat_interleave(repeat_by, dim=0)
+ image = image.to(device=device, dtype=dtype)
+ return image
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # (1) prepare controlnet inputs
+ block_state.device = components._execution_device
+ block_state.height, block_state.width = block_state.latents.shape[-2:]
+ block_state.height = block_state.height * components.vae_scale_factor
+ block_state.width = block_state.width * components.vae_scale_factor
+
+ controlnet = unwrap_module(components.controlnet)
+
+ # (1.1)
+ # control_guidance_start/control_guidance_end (align format)
+ if not isinstance(block_state.control_guidance_start, list) and isinstance(
+ block_state.control_guidance_end, list
+ ):
+ block_state.control_guidance_start = len(block_state.control_guidance_end) * [
+ block_state.control_guidance_start
+ ]
+ elif not isinstance(block_state.control_guidance_end, list) and isinstance(
+ block_state.control_guidance_start, list
+ ):
+ block_state.control_guidance_end = len(block_state.control_guidance_start) * [
+ block_state.control_guidance_end
+ ]
+ elif not isinstance(block_state.control_guidance_start, list) and not isinstance(
+ block_state.control_guidance_end, list
+ ):
+ mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1
+ block_state.control_guidance_start, block_state.control_guidance_end = (
+ mult * [block_state.control_guidance_start],
+ mult * [block_state.control_guidance_end],
+ )
+
+ # (1.2)
+ # controlnet_conditioning_scale (align format)
+ if isinstance(controlnet, MultiControlNetModel) and isinstance(
+ block_state.controlnet_conditioning_scale, float
+ ):
+ block_state.controlnet_conditioning_scale = [block_state.controlnet_conditioning_scale] * len(
+ controlnet.nets
+ )
+
+ # (1.3)
+ # global_pool_conditions
+ block_state.global_pool_conditions = (
+ controlnet.config.global_pool_conditions
+ if isinstance(controlnet, ControlNetModel)
+ else controlnet.nets[0].config.global_pool_conditions
+ )
+ # (1.4)
+ # guess_mode
+ block_state.guess_mode = block_state.guess_mode or block_state.global_pool_conditions
+
+ # (1.5)
+ # control_image
+ if isinstance(controlnet, ControlNetModel):
+ block_state.control_image = self.prepare_control_image(
+ components,
+ image=block_state.control_image,
+ width=block_state.width,
+ height=block_state.height,
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ device=block_state.device,
+ dtype=controlnet.dtype,
+ crops_coords=block_state.crops_coords,
+ )
+ elif isinstance(controlnet, MultiControlNetModel):
+ control_images = []
+
+ for control_image_ in block_state.control_image:
+ control_image = self.prepare_control_image(
+ components,
+ image=control_image_,
+ width=block_state.width,
+ height=block_state.height,
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ device=block_state.device,
+ dtype=controlnet.dtype,
+ crops_coords=block_state.crops_coords,
+ )
+
+ control_images.append(control_image)
+
+ block_state.control_image = control_images
+ else:
+ assert False
+
+ # (1.6)
+ # controlnet_keep
+ block_state.controlnet_keep = []
+ for i in range(len(block_state.timesteps)):
+ keeps = [
+ 1.0 - float(i / len(block_state.timesteps) < s or (i + 1) / len(block_state.timesteps) > e)
+ for s, e in zip(block_state.control_guidance_start, block_state.control_guidance_end)
+ ]
+ block_state.controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)
+
+ block_state.controlnet_cond = block_state.control_image
+ block_state.conditioning_scale = block_state.controlnet_conditioning_scale
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLControlNetUnionInputStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("controlnet", ControlNetUnionModel),
+ ComponentSpec(
+ "control_image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"do_convert_rgb": True, "do_normalize": False}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "step that prepares inputs for the ControlNetUnion model"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("control_image", required=True),
+ InputParam("control_mode", required=True),
+ InputParam("control_guidance_start", default=0.0),
+ InputParam("control_guidance_end", default=1.0),
+ InputParam("controlnet_conditioning_scale", default=1.0),
+ InputParam("guess_mode", default=False),
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Used to determine the shape of the control images. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can be generated in input step.",
+ ),
+ InputParam(
+ "dtype",
+ required=True,
+ type_hint=torch.dtype,
+ description="The dtype of model tensor inputs. Can be generated in input step.",
+ ),
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Needed to determine `controlnet_keep`. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "crops_coords",
+ type_hint=tuple[int] | None,
+ description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("controlnet_cond", type_hint=list[torch.Tensor], description="The processed control images"),
+ OutputParam(
+ "control_type_idx",
+ type_hint=list[int],
+ description="The control mode indices",
+ kwargs_type="controlnet_kwargs",
+ ),
+ OutputParam(
+ "control_type",
+ type_hint=torch.Tensor,
+ description="The control type tensor that specifies which control type is active",
+ kwargs_type="controlnet_kwargs",
+ ),
+ OutputParam("control_guidance_start", type_hint=float, description="The controlnet guidance start value"),
+ OutputParam("control_guidance_end", type_hint=float, description="The controlnet guidance end value"),
+ OutputParam(
+ "conditioning_scale", type_hint=list[float], description="The controlnet conditioning scale values"
+ ),
+ OutputParam("guess_mode", type_hint=bool, description="Whether guess mode is used"),
+ OutputParam("controlnet_keep", type_hint=list[float], description="The controlnet keep values"),
+ ]
+
+ # Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
+ # 1. return image without apply any guidance
+ # 2. add crops_coords and resize_mode to preprocess()
+ @staticmethod
+ def prepare_control_image(
+ components,
+ image,
+ width,
+ height,
+ batch_size,
+ num_images_per_prompt,
+ device,
+ dtype,
+ crops_coords=None,
+ ):
+ if crops_coords is not None:
+ image = components.control_image_processor.preprocess(
+ image, height=height, width=width, crops_coords=crops_coords, resize_mode="fill"
+ ).to(dtype=torch.float32)
+ else:
+ image = components.control_image_processor.preprocess(image, height=height, width=width).to(
+ dtype=torch.float32
+ )
+
+ image_batch_size = image.shape[0]
+ if image_batch_size == 1:
+ repeat_by = batch_size
+ else:
+ # image batch size is the same as prompt batch size
+ repeat_by = num_images_per_prompt
+
+ image = image.repeat_interleave(repeat_by, dim=0)
+ image = image.to(device=device, dtype=dtype)
+ return image
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ controlnet = unwrap_module(components.controlnet)
+
+ device = components._execution_device
+ dtype = block_state.dtype or components.controlnet.dtype
+
+ block_state.height, block_state.width = block_state.latents.shape[-2:]
+ block_state.height = block_state.height * components.vae_scale_factor
+ block_state.width = block_state.width * components.vae_scale_factor
+
+ # control_guidance_start/control_guidance_end (align format)
+ if not isinstance(block_state.control_guidance_start, list) and isinstance(
+ block_state.control_guidance_end, list
+ ):
+ block_state.control_guidance_start = len(block_state.control_guidance_end) * [
+ block_state.control_guidance_start
+ ]
+ elif not isinstance(block_state.control_guidance_end, list) and isinstance(
+ block_state.control_guidance_start, list
+ ):
+ block_state.control_guidance_end = len(block_state.control_guidance_start) * [
+ block_state.control_guidance_end
+ ]
+
+ # guess_mode
+ block_state.global_pool_conditions = controlnet.config.global_pool_conditions
+ block_state.guess_mode = block_state.guess_mode or block_state.global_pool_conditions
+
+ # control_image
+ if not isinstance(block_state.control_image, list):
+ block_state.control_image = [block_state.control_image]
+ # control_mode
+ if not isinstance(block_state.control_mode, list):
+ block_state.control_mode = [block_state.control_mode]
+
+ if len(block_state.control_image) != len(block_state.control_mode):
+ raise ValueError("Expected len(control_image) == len(control_type)")
+
+ # control_type
+ block_state.num_control_type = controlnet.config.num_control_type
+ block_state.control_type = [0 for _ in range(block_state.num_control_type)]
+ for control_idx in block_state.control_mode:
+ block_state.control_type[control_idx] = 1
+ block_state.control_type = torch.Tensor(block_state.control_type)
+
+ block_state.control_type = block_state.control_type.reshape(1, -1).to(device, dtype=block_state.dtype)
+ repeat_by = block_state.batch_size * block_state.num_images_per_prompt // block_state.control_type.shape[0]
+ block_state.control_type = block_state.control_type.repeat_interleave(repeat_by, dim=0)
+
+ # prepare control_image
+ for idx, _ in enumerate(block_state.control_image):
+ block_state.control_image[idx] = self.prepare_control_image(
+ components,
+ image=block_state.control_image[idx],
+ width=block_state.width,
+ height=block_state.height,
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ device=device,
+ dtype=dtype,
+ crops_coords=block_state.crops_coords,
+ )
+ block_state.height, block_state.width = block_state.control_image[idx].shape[-2:]
+
+ # controlnet_keep
+ block_state.controlnet_keep = []
+ for i in range(len(block_state.timesteps)):
+ block_state.controlnet_keep.append(
+ 1.0
+ - float(
+ i / len(block_state.timesteps) < block_state.control_guidance_start
+ or (i + 1) / len(block_state.timesteps) > block_state.control_guidance_end
+ )
+ )
+ block_state.control_type_idx = block_state.control_mode
+ block_state.controlnet_cond = block_state.control_image
+ block_state.conditioning_scale = block_state.controlnet_conditioning_scale
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd8592b51e42a3a10659e82e1d809e8cd579d444
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/decoders.py
@@ -0,0 +1,198 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKL
+from ...utils import deprecate, logging
+from ..modular_pipeline import (
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class StableDiffusionXLDecodeStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("output_type", default="pil"),
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents from the denoising step",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
+ description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array",
+ )
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae with self->components
+ def upcast_vae(components):
+ deprecate(
+ "upcast_vae",
+ "1.0.0",
+ "`upcast_vae` is deprecated. Please use `pipe.vae.to(torch.float32)`. For more details, please refer to: https://github.com/huggingface/diffusers/pull/12619#issue-3606633695.",
+ )
+ components.vae.to(dtype=torch.float32)
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ if not block_state.output_type == "latent":
+ latents = block_state.latents
+ # make sure the VAE is in float32 mode, as it overflows in float16
+ block_state.needs_upcasting = components.vae.dtype == torch.float16 and components.vae.config.force_upcast
+
+ if block_state.needs_upcasting:
+ self.upcast_vae(components)
+ latents = latents.to(next(iter(components.vae.post_quant_conv.parameters())).dtype)
+ elif latents.dtype != components.vae.dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ components.vae = components.vae.to(latents.dtype)
+
+ # unscale/denormalize the latents
+ # denormalize with the mean and std if available and not None
+ block_state.has_latents_mean = (
+ hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None
+ )
+ block_state.has_latents_std = (
+ hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None
+ )
+ if block_state.has_latents_mean and block_state.has_latents_std:
+ block_state.latents_mean = (
+ torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+ )
+ block_state.latents_std = (
+ torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+ )
+ latents = (
+ latents * block_state.latents_std / components.vae.config.scaling_factor + block_state.latents_mean
+ )
+ else:
+ latents = latents / components.vae.config.scaling_factor
+
+ block_state.images = components.vae.decode(latents, return_dict=False)[0]
+
+ # cast back to fp16 if needed
+ if block_state.needs_upcasting:
+ components.vae.to(dtype=torch.float16)
+ else:
+ block_state.images = block_state.latents
+
+ # apply watermark if available
+ if hasattr(components, "watermark") and components.watermark is not None:
+ block_state.images = components.watermark.apply_watermark(block_state.images)
+
+ block_state.images = components.image_processor.postprocess(
+ block_state.images, output_type=block_state.output_type
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLInpaintOverlayMaskStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return (
+ "A post-processing step that overlays the mask on the image (inpainting task only).\n"
+ + "only needed when you are using the `padding_mask_crop` option when pre-processing the image and mask"
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("image"),
+ InputParam("mask_image"),
+ InputParam("padding_mask_crop"),
+ InputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
+ description="The generated images from the decode step",
+ ),
+ InputParam(
+ "crops_coords",
+ type_hint=tuple[int, int],
+ description="The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can be generated in vae_encode step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ if block_state.padding_mask_crop is not None and block_state.crops_coords is not None:
+ block_state.images = [
+ components.image_processor.apply_overlay(
+ block_state.mask_image, block_state.image, i, block_state.crops_coords
+ )
+ for i in block_state.images
+ ]
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..1509ca91348ebc2f6108edd60bfe28c32e74fd1a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/denoise.py
@@ -0,0 +1,798 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import ControlNetModel, UNet2DConditionModel
+from ...schedulers import EulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusionXLModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# YiYi experimenting composible denoise loop
+# loop step (1): prepare latent input for denoiser
+class StableDiffusionXLLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepare the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `StableDiffusionXLDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[str]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int):
+ block_state.scaled_latents = components.scheduler.scale_model_input(block_state.latents, t)
+
+ return components, block_state
+
+
+# loop step (1): prepare latent input for denoiser (with inpainting)
+class StableDiffusionXLInpaintLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepare the latent input for the denoiser (for inpainting workflow only). "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` object"
+ )
+
+ @property
+ def inputs(self) -> list[str]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "mask",
+ type_hint=torch.Tensor | None,
+ description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
+ ),
+ InputParam(
+ "masked_image_latents",
+ type_hint=torch.Tensor | None,
+ description="The masked image latents to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ num_channels_unet = components.num_channels_unet
+ if num_channels_unet == 9:
+ # default case for stable-diffusion-v1-5/stable-diffusion-inpainting
+ if block_state.mask is None or block_state.masked_image_latents is None:
+ raise ValueError("mask and masked_image_latents must be provided for inpainting-specific Unet")
+ num_channels_latents = block_state.latents.shape[1]
+ num_channels_mask = block_state.mask.shape[1]
+ num_channels_masked_image = block_state.masked_image_latents.shape[1]
+ if num_channels_latents + num_channels_mask + num_channels_masked_image != num_channels_unet:
+ raise ValueError(
+ f"Incorrect configuration settings! The config of `components.unet`: {components.unet.config} expects"
+ f" {components.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
+ f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
+ f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of"
+ " `components.unet` or your `mask_image` or `image` input."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int):
+ self.check_inputs(components, block_state)
+
+ block_state.scaled_latents = components.scheduler.scale_model_input(block_state.latents, t)
+ if components.num_channels_unet == 9:
+ block_state.scaled_latents = torch.cat(
+ [block_state.scaled_latents, block_state.mask, block_state.masked_image_latents], dim=1
+ )
+
+ return components, block_state
+
+
+# loop step (2): denoise the latents with guidance
+class StableDiffusionXLLoopDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents with guidance. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `StableDiffusionXLDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("cross_attention_kwargs"),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "timestep_cond",
+ type_hint=torch.Tensor | None,
+ description="The guidance scale embedding to use for Latent Consistency Models(LCMs). Can be generated in prepare_additional_conditioning step.",
+ ),
+ InputParam(
+ kwargs_type="denoiser_input_fields",
+ description=(
+ "All conditional model inputs that need to be prepared with guider. "
+ "It should contain prompt_embeds/negative_prompt_embeds, "
+ "add_time_ids/negative_add_time_ids, "
+ "pooled_prompt_embeds/negative_pooled_prompt_embeds, "
+ "and ip_adapter_embeds/negative_ip_adapter_embeds (optional)."
+ "please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state"
+ ),
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(
+ self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int
+ ) -> PipelineState:
+ # Map the keys we'll see on each `guider_state_batch` (e.g. guider_state_batch.prompt_embeds)
+ # to the corresponding (cond, uncond) fields on block_state. (e.g. block_state.prompt_embeds, block_state.negative_prompt_embeds)
+ guider_inputs = {
+ "prompt_embeds": (
+ getattr(block_state, "prompt_embeds", None),
+ getattr(block_state, "negative_prompt_embeds", None),
+ ),
+ "time_ids": (
+ getattr(block_state, "add_time_ids", None),
+ getattr(block_state, "negative_add_time_ids", None),
+ ),
+ "text_embeds": (
+ getattr(block_state, "pooled_prompt_embeds", None),
+ getattr(block_state, "negative_pooled_prompt_embeds", None),
+ ),
+ "image_embeds": (
+ getattr(block_state, "ip_adapter_embeds", None),
+ getattr(block_state, "negative_ip_adapter_embeds", None),
+ ),
+ }
+
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+ # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+ # you will get a guider_state with two batches:
+ # guider_state = [
+ # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch
+ # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch
+ # ]
+ # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ # run the denoiser for each guidance batch
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.unet)
+ cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+ prompt_embeds = cond_kwargs.pop("prompt_embeds")
+
+ # Predict the noise residual
+ # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+ guider_state_batch.noise_pred = components.unet(
+ block_state.scaled_latents,
+ t,
+ encoder_hidden_states=prompt_embeds,
+ timestep_cond=block_state.timestep_cond,
+ cross_attention_kwargs=block_state.cross_attention_kwargs,
+ added_cond_kwargs=cond_kwargs,
+ return_dict=False,
+ )[0]
+ components.guider.cleanup_models(components.unet)
+
+ # Perform guidance
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+# loop step (2): denoise the latents with guidance (with controlnet)
+class StableDiffusionXLControlNetLoopDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ComponentSpec("controlnet", ControlNetModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that denoise the latents with guidance (with controlnet). "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `StableDiffusionXLDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("cross_attention_kwargs"),
+ InputParam(
+ "controlnet_cond",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ InputParam(
+ "conditioning_scale",
+ type_hint=float,
+ description="The controlnet conditioning scale value to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ InputParam(
+ "guess_mode",
+ required=True,
+ type_hint=bool,
+ description="The guess mode value to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ InputParam(
+ "controlnet_keep",
+ required=True,
+ type_hint=list[float],
+ description="The controlnet keep values to use for the denoising process. Can be generated in prepare_controlnet_inputs step.",
+ ),
+ InputParam(
+ "timestep_cond",
+ type_hint=torch.Tensor | None,
+ description="The guidance scale embedding to use for Latent Consistency Models(LCMs), can be generated by prepare_additional_conditioning step",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ kwargs_type="denoiser_input_fields",
+ description=(
+ "All conditional model inputs that need to be prepared with guider. "
+ "It should contain prompt_embeds/negative_prompt_embeds, "
+ "add_time_ids/negative_add_time_ids, "
+ "pooled_prompt_embeds/negative_pooled_prompt_embeds, "
+ "and ip_adapter_embeds/negative_ip_adapter_embeds (optional)."
+ "please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state"
+ ),
+ ),
+ InputParam(
+ kwargs_type="controlnet_kwargs",
+ description=(
+ "additional kwargs for controlnet (e.g. control_type_idx and control_type from the controlnet union input step )"
+ "please add `kwargs_type=controlnet_kwargs` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state"
+ ),
+ ),
+ ]
+
+ @staticmethod
+ def prepare_extra_kwargs(func, exclude_kwargs=[], **kwargs):
+ accepted_kwargs = set(inspect.signature(func).parameters.keys())
+ extra_kwargs = {}
+ for key, value in kwargs.items():
+ if key in accepted_kwargs and key not in exclude_kwargs:
+ extra_kwargs[key] = value
+
+ return extra_kwargs
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int):
+ extra_controlnet_kwargs = self.prepare_extra_kwargs(
+ components.controlnet.forward, **block_state.controlnet_kwargs
+ )
+
+ # Map the keys we'll see on each `guider_state_batch` (e.g. guider_state_batch.prompt_embeds)
+ # to the corresponding (cond, uncond) fields on block_state. (e.g. block_state.prompt_embeds, block_state.negative_prompt_embeds)
+ guider_inputs = {
+ "prompt_embeds": (
+ getattr(block_state, "prompt_embeds", None),
+ getattr(block_state, "negative_prompt_embeds", None),
+ ),
+ "time_ids": (
+ getattr(block_state, "add_time_ids", None),
+ getattr(block_state, "negative_add_time_ids", None),
+ ),
+ "text_embeds": (
+ getattr(block_state, "pooled_prompt_embeds", None),
+ getattr(block_state, "negative_pooled_prompt_embeds", None),
+ ),
+ "image_embeds": (
+ getattr(block_state, "ip_adapter_embeds", None),
+ getattr(block_state, "negative_ip_adapter_embeds", None),
+ ),
+ }
+
+ # cond_scale for the timestep (controlnet input)
+ if isinstance(block_state.controlnet_keep[i], list):
+ block_state.cond_scale = [
+ c * s for c, s in zip(block_state.conditioning_scale, block_state.controlnet_keep[i])
+ ]
+ else:
+ controlnet_cond_scale = block_state.conditioning_scale
+ if isinstance(controlnet_cond_scale, list):
+ controlnet_cond_scale = controlnet_cond_scale[0]
+ block_state.cond_scale = controlnet_cond_scale * block_state.controlnet_keep[i]
+
+ # default controlnet output/unet input for guess mode + conditional path
+ block_state.down_block_res_samples_zeros = None
+ block_state.mid_block_res_sample_zeros = None
+
+ # guided denoiser step
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+ # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+ # you will get a guider_state with two batches:
+ # guider_state = [
+ # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch
+ # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch
+ # ]
+ # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+ guider_state = components.guider.prepare_inputs(guider_inputs)
+
+ # run the denoiser for each guidance batch
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.unet)
+
+ # Prepare additional conditionings
+ added_cond_kwargs = {
+ "text_embeds": guider_state_batch.text_embeds,
+ "time_ids": guider_state_batch.time_ids,
+ }
+ if hasattr(guider_state_batch, "image_embeds") and guider_state_batch.image_embeds is not None:
+ added_cond_kwargs["image_embeds"] = guider_state_batch.image_embeds
+
+ # Prepare controlnet additional conditionings
+ controlnet_added_cond_kwargs = {
+ "text_embeds": guider_state_batch.text_embeds,
+ "time_ids": guider_state_batch.time_ids,
+ }
+ # run controlnet for the guidance batch
+ if block_state.guess_mode and not components.guider.is_conditional:
+ # guider always run uncond batch first, so these tensors should be set already
+ down_block_res_samples = block_state.down_block_res_samples_zeros
+ mid_block_res_sample = block_state.mid_block_res_sample_zeros
+ else:
+ down_block_res_samples, mid_block_res_sample = components.controlnet(
+ block_state.scaled_latents,
+ t,
+ encoder_hidden_states=guider_state_batch.prompt_embeds,
+ controlnet_cond=block_state.controlnet_cond,
+ conditioning_scale=block_state.cond_scale,
+ guess_mode=block_state.guess_mode,
+ added_cond_kwargs=controlnet_added_cond_kwargs,
+ return_dict=False,
+ **extra_controlnet_kwargs,
+ )
+
+ # assign it to block_state so it will be available for the uncond guidance batch
+ if block_state.down_block_res_samples_zeros is None:
+ block_state.down_block_res_samples_zeros = [torch.zeros_like(d) for d in down_block_res_samples]
+ if block_state.mid_block_res_sample_zeros is None:
+ block_state.mid_block_res_sample_zeros = torch.zeros_like(mid_block_res_sample)
+
+ # Predict the noise
+ # store the noise_pred in guider_state_batch so we can apply guidance across all batches
+ guider_state_batch.noise_pred = components.unet(
+ block_state.scaled_latents,
+ t,
+ encoder_hidden_states=guider_state_batch.prompt_embeds,
+ timestep_cond=block_state.timestep_cond,
+ cross_attention_kwargs=block_state.cross_attention_kwargs,
+ added_cond_kwargs=added_cond_kwargs,
+ down_block_additional_residuals=down_block_res_samples,
+ mid_block_additional_residual=mid_block_res_sample,
+ return_dict=False,
+ )[0]
+ components.guider.cleanup_models(components.unet)
+
+ # Perform guidance
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+# loop step (3): scheduler step to update latents
+class StableDiffusionXLLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that update the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `StableDiffusionXLDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("eta", default=0.0),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
+
+ # YiYi TODO: move this out of here
+ @staticmethod
+ def prepare_extra_kwargs(func, exclude_kwargs=[], **kwargs):
+ accepted_kwargs = set(inspect.signature(func).parameters.keys())
+ extra_kwargs = {}
+ for key, value in kwargs.items():
+ if key in accepted_kwargs and key not in exclude_kwargs:
+ extra_kwargs[key] = value
+
+ return extra_kwargs
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int):
+ # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+ block_state.extra_step_kwargs = self.prepare_extra_kwargs(
+ components.scheduler.step, generator=block_state.generator, eta=block_state.eta
+ )
+
+ # Perform scheduler step using the predicted output
+ block_state.latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ **block_state.extra_step_kwargs,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != block_state.latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ block_state.latents = block_state.latents.to(block_state.latents_dtype)
+
+ return components, block_state
+
+
+# loop step (3): scheduler step to update latents (with inpainting)
+class StableDiffusionXLInpaintLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that update the latents (for inpainting workflow only). "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `StableDiffusionXLDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam("eta", default=0.0),
+ InputParam("generator"),
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "mask",
+ type_hint=torch.Tensor | None,
+ description="The mask to use for the denoising process, for inpainting task only. Can be generated in vae_encode or prepare_latent step.",
+ ),
+ InputParam(
+ "noise",
+ type_hint=torch.Tensor | None,
+ description="The noise added to the image latents, for inpainting task only. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "image_latents",
+ type_hint=torch.Tensor | None,
+ description="The image latents to use for the denoising process, for inpainting/image-to-image task only. Can be generated in vae_encode or prepare_latent step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]
+
+ @staticmethod
+ def prepare_extra_kwargs(func, exclude_kwargs=[], **kwargs):
+ accepted_kwargs = set(inspect.signature(func).parameters.keys())
+ extra_kwargs = {}
+ for key, value in kwargs.items():
+ if key in accepted_kwargs and key not in exclude_kwargs:
+ extra_kwargs[key] = value
+
+ return extra_kwargs
+
+ def check_inputs(self, components, block_state):
+ if components.num_channels_unet == 4:
+ if block_state.image_latents is None:
+ raise ValueError(f"image_latents is required for this step {self.__class__.__name__}")
+ if block_state.mask is None:
+ raise ValueError(f"mask is required for this step {self.__class__.__name__}")
+ if block_state.noise is None:
+ raise ValueError(f"noise is required for this step {self.__class__.__name__}")
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, block_state: BlockState, i: int, t: int):
+ self.check_inputs(components, block_state)
+
+ # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+ block_state.extra_step_kwargs = self.prepare_extra_kwargs(
+ components.scheduler.step, generator=block_state.generator, eta=block_state.eta
+ )
+
+ # Perform scheduler step using the predicted output
+ block_state.latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred,
+ t,
+ block_state.latents,
+ **block_state.extra_step_kwargs,
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != block_state.latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ block_state.latents = block_state.latents.to(block_state.latents_dtype)
+
+ # adjust latent for inpainting
+ if components.num_channels_unet == 4:
+ block_state.init_latents_proper = block_state.image_latents
+ if i < len(block_state.timesteps) - 1:
+ block_state.noise_timestep = block_state.timesteps[i + 1]
+ block_state.init_latents_proper = components.scheduler.add_noise(
+ block_state.init_latents_proper, block_state.noise, torch.tensor([block_state.noise_timestep])
+ )
+
+ block_state.latents = (
+ 1 - block_state.mask
+ ) * block_state.init_latents_proper + block_state.mask * block_state.latents
+
+ return components, block_state
+
+
+# the loop wrapper that iterates over the timesteps
+class StableDiffusionXLDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoise the latents over `timesteps`. "
+ "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("scheduler", EulerDiscreteScheduler),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.disable_guidance = True if components.unet.config.time_cond_proj_dim is not None else False
+ if block_state.disable_guidance:
+ components.guider.disable()
+ else:
+ components.guider.enable()
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+# composing the denoising loops
+class StableDiffusionXLDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
+ block_classes = [
+ StableDiffusionXLLoopBeforeDenoiser,
+ StableDiffusionXLLoopDenoiser,
+ StableDiffusionXLLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `StableDiffusionXLDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `StableDiffusionXLLoopBeforeDenoiser`\n"
+ " - `StableDiffusionXLLoopDenoiser`\n"
+ " - `StableDiffusionXLLoopAfterDenoiser`\n"
+ "This block supports both text2img and img2img tasks."
+ )
+
+
+# control_cond
+class StableDiffusionXLControlNetDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
+ block_classes = [
+ StableDiffusionXLLoopBeforeDenoiser,
+ StableDiffusionXLControlNetLoopDenoiser,
+ StableDiffusionXLLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents with controlnet. \n"
+ "Its loop logic is defined in `StableDiffusionXLDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `StableDiffusionXLLoopBeforeDenoiser`\n"
+ " - `StableDiffusionXLControlNetLoopDenoiser`\n"
+ " - `StableDiffusionXLLoopAfterDenoiser`\n"
+ "This block supports using controlnet for both text2img and img2img tasks."
+ )
+
+
+# mask
+class StableDiffusionXLInpaintDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
+ block_classes = [
+ StableDiffusionXLInpaintLoopBeforeDenoiser,
+ StableDiffusionXLLoopDenoiser,
+ StableDiffusionXLInpaintLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents(for inpainting task only). \n"
+ "Its loop logic is defined in `StableDiffusionXLDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `StableDiffusionXLInpaintLoopBeforeDenoiser`\n"
+ " - `StableDiffusionXLLoopDenoiser`\n"
+ " - `StableDiffusionXLInpaintLoopAfterDenoiser`\n"
+ "This block onlysupports inpainting tasks."
+ )
+
+
+# control_cond + mask
+class StableDiffusionXLInpaintControlNetDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
+ block_classes = [
+ StableDiffusionXLInpaintLoopBeforeDenoiser,
+ StableDiffusionXLControlNetLoopDenoiser,
+ StableDiffusionXLInpaintLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents(for inpainting task only) with controlnet. \n"
+ "Its loop logic is defined in `StableDiffusionXLDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `StableDiffusionXLInpaintLoopBeforeDenoiser`\n"
+ " - `StableDiffusionXLControlNetLoopDenoiser`\n"
+ " - `StableDiffusionXLInpaintLoopAfterDenoiser`\n"
+ "This block only supports using controlnet for inpainting tasks."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..19b7788492d2fe4716b5ae73a4a5208bf5dadd16
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/encoders.py
@@ -0,0 +1,885 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from transformers import (
+ CLIPImageProcessor,
+ CLIPTextModel,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ CLIPVisionModelWithProjection,
+)
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import PipelineImageInput, VaeImageProcessor
+from ...loaders import StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin
+from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel
+from ...models.lora import adjust_lora_scale_text_encoder
+from ...utils import (
+ USE_PEFT_BACKEND,
+ logging,
+ scale_lora_layers,
+ unscale_lora_layers,
+)
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
+from .modular_pipeline import StableDiffusionXLModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class StableDiffusionXLIPAdapterStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return (
+ "IP Adapter step that prepares ip adapter image embeddings.\n"
+ "Note that this step only prepares the embeddings - in order for it to work correctly, "
+ "you need to load ip adapter weights into unet via ModularPipeline.load_ip_adapter() and pipeline.set_ip_adapter_scale().\n"
+ "See [ModularIPAdapterMixin](https://huggingface.co/docs/diffusers/api/loaders/ip_adapter#diffusers.loaders.ModularIPAdapterMixin)"
+ " for more details"
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("image_encoder", CLIPVisionModelWithProjection),
+ ComponentSpec(
+ "feature_extractor",
+ CLIPImageProcessor,
+ config=FrozenDict({"size": 224, "crop_size": 224}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("unet", UNet2DConditionModel),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "ip_adapter_image",
+ PipelineImageInput,
+ required=True,
+ description="The image(s) to be used as ip adapter",
+ )
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("ip_adapter_embeds", type_hint=torch.Tensor, description="IP adapter image embeddings"),
+ OutputParam(
+ "negative_ip_adapter_embeds",
+ type_hint=torch.Tensor,
+ description="Negative IP adapter image embeddings",
+ ),
+ ]
+
+ @staticmethod
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image with self->components
+ def encode_image(components, image, device, num_images_per_prompt, output_hidden_states=None):
+ dtype = next(components.image_encoder.parameters()).dtype
+
+ if not isinstance(image, torch.Tensor):
+ image = components.feature_extractor(image, return_tensors="pt").pixel_values
+
+ image = image.to(device=device, dtype=dtype)
+ if output_hidden_states:
+ image_enc_hidden_states = components.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
+ uncond_image_enc_hidden_states = components.image_encoder(
+ torch.zeros_like(image), output_hidden_states=True
+ ).hidden_states[-2]
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
+ num_images_per_prompt, dim=0
+ )
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
+ else:
+ image_embeds = components.image_encoder(image).image_embeds
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+ uncond_image_embeds = torch.zeros_like(image_embeds)
+
+ return image_embeds, uncond_image_embeds
+
+ # modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
+ def prepare_ip_adapter_image_embeds(
+ self,
+ components,
+ ip_adapter_image,
+ ip_adapter_image_embeds,
+ device,
+ num_images_per_prompt,
+ prepare_unconditional_embeds,
+ ):
+ image_embeds = []
+ if prepare_unconditional_embeds:
+ negative_image_embeds = []
+ if ip_adapter_image_embeds is None:
+ if not isinstance(ip_adapter_image, list):
+ ip_adapter_image = [ip_adapter_image]
+
+ if len(ip_adapter_image) != len(components.unet.encoder_hid_proj.image_projection_layers):
+ raise ValueError(
+ f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(components.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
+ )
+
+ for single_ip_adapter_image, image_proj_layer in zip(
+ ip_adapter_image, components.unet.encoder_hid_proj.image_projection_layers
+ ):
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
+ components, single_ip_adapter_image, device, 1, output_hidden_state
+ )
+
+ image_embeds.append(single_image_embeds[None, :])
+ if prepare_unconditional_embeds:
+ negative_image_embeds.append(single_negative_image_embeds[None, :])
+ else:
+ for single_image_embeds in ip_adapter_image_embeds:
+ if prepare_unconditional_embeds:
+ single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
+ negative_image_embeds.append(single_negative_image_embeds)
+ image_embeds.append(single_image_embeds)
+
+ ip_adapter_image_embeds = []
+ for i, single_image_embeds in enumerate(image_embeds):
+ single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+ if prepare_unconditional_embeds:
+ single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
+
+ single_image_embeds = single_image_embeds.to(device=device)
+ ip_adapter_image_embeds.append(single_image_embeds)
+
+ return ip_adapter_image_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.prepare_unconditional_embeds = components.guider.num_conditions > 1
+ block_state.device = components._execution_device
+
+ block_state.ip_adapter_embeds = self.prepare_ip_adapter_image_embeds(
+ components,
+ ip_adapter_image=block_state.ip_adapter_image,
+ ip_adapter_image_embeds=None,
+ device=block_state.device,
+ num_images_per_prompt=1,
+ prepare_unconditional_embeds=block_state.prepare_unconditional_embeds,
+ )
+ if block_state.prepare_unconditional_embeds:
+ block_state.negative_ip_adapter_embeds = []
+ for i, image_embeds in enumerate(block_state.ip_adapter_embeds):
+ negative_image_embeds, image_embeds = image_embeds.chunk(2)
+ block_state.negative_ip_adapter_embeds.append(negative_image_embeds)
+ block_state.ip_adapter_embeds[i] = image_embeds
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusionXLTextEncoderStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generate text_embeddings to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", CLIPTextModel),
+ ComponentSpec("text_encoder_2", CLIPTextModelWithProjection),
+ ComponentSpec("tokenizer", CLIPTokenizer),
+ ComponentSpec("tokenizer_2", CLIPTokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 7.5}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [ConfigSpec("force_zeros_for_empty_prompt", True)]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("prompt_2"),
+ InputParam("negative_prompt"),
+ InputParam("negative_prompt_2"),
+ InputParam("cross_attention_kwargs"),
+ InputParam("clip_skip"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="negative text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="pooled text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_pooled_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="negative pooled text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ if block_state.prompt is not None and (
+ not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list)
+ ):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+ elif block_state.prompt_2 is not None and (
+ not isinstance(block_state.prompt_2, str) and not isinstance(block_state.prompt_2, list)
+ ):
+ raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(block_state.prompt_2)}")
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str,
+ prompt_2: str | None = None,
+ device: torch.device | None = None,
+ num_images_per_prompt: int = 1,
+ prepare_unconditional_embeds: bool = True,
+ negative_prompt: str | None = None,
+ negative_prompt_2: str | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ pooled_prompt_embeds: torch.Tensor | None = None,
+ negative_pooled_prompt_embeds: torch.Tensor | None = None,
+ lora_scale: float | None = None,
+ clip_skip: int | None = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ prompt to be encoded
+ prompt_2 (`str` or `list[str]`, *optional*):
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+ used in both text-encoders
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ prepare_unconditional_embeds (`bool`):
+ whether to use prepare unconditional embeddings or not
+ negative_prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ negative_prompt_2 (`str` or `list[str]`, *optional*):
+ The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+ `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
+ negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+ input argument.
+ lora_scale (`float`, *optional*):
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+ clip_skip (`int`, *optional*):
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+ the output of the pre-final layer will be used for computing the prompt embeddings.
+ """
+ device = device or components._execution_device
+
+ # set lora scale so that monkey patched LoRA
+ # function of text encoder can correctly access it
+ if lora_scale is not None and isinstance(components, StableDiffusionXLLoraLoaderMixin):
+ components._lora_scale = lora_scale
+
+ # dynamically adjust the LoRA scale
+ if components.text_encoder is not None:
+ if not USE_PEFT_BACKEND:
+ adjust_lora_scale_text_encoder(components.text_encoder, lora_scale)
+ else:
+ scale_lora_layers(components.text_encoder, lora_scale)
+
+ if components.text_encoder_2 is not None:
+ if not USE_PEFT_BACKEND:
+ adjust_lora_scale_text_encoder(components.text_encoder_2, lora_scale)
+ else:
+ scale_lora_layers(components.text_encoder_2, lora_scale)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ # Define tokenizers and text encoders
+ tokenizers = (
+ [components.tokenizer, components.tokenizer_2]
+ if components.tokenizer is not None
+ else [components.tokenizer_2]
+ )
+ text_encoders = (
+ [components.text_encoder, components.text_encoder_2]
+ if components.text_encoder is not None
+ else [components.text_encoder_2]
+ )
+
+ if prompt_embeds is None:
+ prompt_2 = prompt_2 or prompt
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+ # textual inversion: process multi-vector tokens if necessary
+ prompt_embeds_list = []
+ prompts = [prompt, prompt_2]
+ for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+ if isinstance(components, TextualInversionLoaderMixin):
+ prompt = components.maybe_convert_prompt(prompt, tokenizer)
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+ text_input_ids, untruncated_ids
+ ):
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {tokenizer.model_max_length} tokens: {removed_text}"
+ )
+
+ prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+
+ # We are only ALWAYS interested in the pooled output of the final text encoder
+ pooled_prompt_embeds = prompt_embeds[0]
+ if clip_skip is None:
+ prompt_embeds = prompt_embeds.hidden_states[-2]
+ else:
+ # "2" because SDXL always indexes from the penultimate layer.
+ prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+ prompt_embeds_list.append(prompt_embeds)
+
+ prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+
+ # get unconditional embeddings for classifier free guidance
+ zero_out_negative_prompt = negative_prompt is None and components.config.force_zeros_for_empty_prompt
+ if prepare_unconditional_embeds and negative_prompt_embeds is None and zero_out_negative_prompt:
+ negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+ negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
+ elif prepare_unconditional_embeds and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+ # normalize str to list
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+ negative_prompt_2 = (
+ batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+ )
+
+ uncond_tokens: list[str]
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+ else:
+ uncond_tokens = [negative_prompt, negative_prompt_2]
+
+ negative_prompt_embeds_list = []
+ for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
+ if isinstance(components, TextualInversionLoaderMixin):
+ negative_prompt = components.maybe_convert_prompt(negative_prompt, tokenizer)
+
+ max_length = prompt_embeds.shape[1]
+ uncond_input = tokenizer(
+ negative_prompt,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ negative_prompt_embeds = text_encoder(
+ uncond_input.input_ids.to(device),
+ output_hidden_states=True,
+ )
+ # We are only ALWAYS interested in the pooled output of the final text encoder
+ negative_pooled_prompt_embeds = negative_prompt_embeds[0]
+ negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
+
+ negative_prompt_embeds_list.append(negative_prompt_embeds)
+
+ negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
+
+ if components.text_encoder_2 is not None:
+ prompt_embeds = prompt_embeds.to(dtype=components.text_encoder_2.dtype, device=device)
+ else:
+ prompt_embeds = prompt_embeds.to(dtype=components.unet.dtype, device=device)
+
+ bs_embed, seq_len, _ = prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+ if prepare_unconditional_embeds:
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+ seq_len = negative_prompt_embeds.shape[1]
+
+ if components.text_encoder_2 is not None:
+ negative_prompt_embeds = negative_prompt_embeds.to(
+ dtype=components.text_encoder_2.dtype, device=device
+ )
+ else:
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=components.unet.dtype, device=device)
+
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+ pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+ bs_embed * num_images_per_prompt, -1
+ )
+ if prepare_unconditional_embeds:
+ negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+ bs_embed * num_images_per_prompt, -1
+ )
+
+ if components.text_encoder is not None:
+ if isinstance(components, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(components.text_encoder, lora_scale)
+
+ if components.text_encoder_2 is not None:
+ if isinstance(components, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(components.text_encoder_2, lora_scale)
+
+ return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ # Get inputs and intermediates
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.prepare_unconditional_embeds = components.guider.num_conditions > 1
+ block_state.device = components._execution_device
+
+ # Encode input prompt
+ block_state.text_encoder_lora_scale = (
+ block_state.cross_attention_kwargs.get("scale", None)
+ if block_state.cross_attention_kwargs is not None
+ else None
+ )
+ (
+ block_state.prompt_embeds,
+ block_state.negative_prompt_embeds,
+ block_state.pooled_prompt_embeds,
+ block_state.negative_pooled_prompt_embeds,
+ ) = self.encode_prompt(
+ components,
+ block_state.prompt,
+ block_state.prompt_2,
+ block_state.device,
+ 1,
+ block_state.prepare_unconditional_embeds,
+ block_state.negative_prompt,
+ block_state.negative_prompt_2,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ pooled_prompt_embeds=None,
+ negative_pooled_prompt_embeds=None,
+ lora_scale=block_state.text_encoder_lora_scale,
+ clip_skip=block_state.clip_skip,
+ )
+ # Add outputs
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class StableDiffusionXLVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def description(self) -> str:
+ return "Vae Encoder step that encode the input image into a latent representation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image", required=True),
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("generator"),
+ InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"),
+ InputParam(
+ "preprocess_kwargs",
+ type_hint=dict | None,
+ description="A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="The latents representing the reference image for image-to-image/inpainting generation",
+ )
+ ]
+
+ # Modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline._encode_vae_image with self -> components
+ # YiYi TODO: update the _encode_vae_image so that we can use #Coped from
+ def _encode_vae_image(self, components, image: torch.Tensor, generator: torch.Generator):
+ latents_mean = latents_std = None
+ if hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None:
+ latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1)
+ if hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None:
+ latents_std = torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1)
+
+ dtype = image.dtype
+ if components.vae.config.force_upcast:
+ image = image.float()
+ components.vae.to(dtype=torch.float32)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(components.vae.encode(image[i : i + 1]), generator=generator[i])
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(components.vae.encode(image), generator=generator)
+
+ if components.vae.config.force_upcast:
+ components.vae.to(dtype)
+
+ image_latents = image_latents.to(dtype)
+ if latents_mean is not None and latents_std is not None:
+ latents_mean = latents_mean.to(device=image_latents.device, dtype=dtype)
+ latents_std = latents_std.to(device=image_latents.device, dtype=dtype)
+ image_latents = (image_latents - latents_mean) * components.vae.config.scaling_factor / latents_std
+ else:
+ image_latents = components.vae.config.scaling_factor * image_latents
+
+ return image_latents
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ block_state.preprocess_kwargs = block_state.preprocess_kwargs or {}
+ block_state.device = components._execution_device
+ block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype
+
+ image = components.image_processor.preprocess(
+ block_state.image, height=block_state.height, width=block_state.width, **block_state.preprocess_kwargs
+ )
+ image = image.to(device=block_state.device, dtype=block_state.dtype)
+ block_state.batch_size = image.shape[0]
+
+ # if generator is a list, make sure the length of it matches the length of images (both should be batch_size)
+ if isinstance(block_state.generator, list) and len(block_state.generator) != block_state.batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch"
+ f" size of {block_state.batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ block_state.image_latents = self._encode_vae_image(components, image=image, generator=block_state.generator)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class StableDiffusionXLInpaintVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "stable-diffusion-xl"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec(
+ "mask_processor",
+ VaeImageProcessor,
+ config=FrozenDict(
+ {"do_normalize": False, "vae_scale_factor": 8, "do_binarize": True, "do_convert_grayscale": True}
+ ),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Vae encoder step that prepares the image and mask for the inpainting process"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("image", required=True),
+ InputParam("mask_image", required=True),
+ InputParam("padding_mask_crop"),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents", type_hint=torch.Tensor, description="The latents representation of the input image"
+ ),
+ OutputParam("mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process"),
+ OutputParam(
+ "masked_image_latents",
+ type_hint=torch.Tensor,
+ description="The masked image latents to use for the inpainting process (only for inpainting-specifid unet)",
+ ),
+ OutputParam(
+ "crops_coords",
+ type_hint=tuple[int, int] | None,
+ description="The crop coordinates to use for the preprocess/postprocess of the image and mask",
+ ),
+ ]
+
+ # Modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline._encode_vae_image with self -> components
+ # YiYi TODO: update the _encode_vae_image so that we can use #Coped from
+ def _encode_vae_image(self, components, image: torch.Tensor, generator: torch.Generator):
+ latents_mean = latents_std = None
+ if hasattr(components.vae.config, "latents_mean") and components.vae.config.latents_mean is not None:
+ latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, 4, 1, 1)
+ if hasattr(components.vae.config, "latents_std") and components.vae.config.latents_std is not None:
+ latents_std = torch.tensor(components.vae.config.latents_std).view(1, 4, 1, 1)
+
+ dtype = image.dtype
+ if components.vae.config.force_upcast:
+ image = image.float()
+ components.vae.to(dtype=torch.float32)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(components.vae.encode(image[i : i + 1]), generator=generator[i])
+ for i in range(image.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(components.vae.encode(image), generator=generator)
+
+ if components.vae.config.force_upcast:
+ components.vae.to(dtype)
+
+ image_latents = image_latents.to(dtype)
+ if latents_mean is not None and latents_std is not None:
+ latents_mean = latents_mean.to(device=image_latents.device, dtype=dtype)
+ latents_std = latents_std.to(device=image_latents.device, dtype=dtype)
+ image_latents = (image_latents - latents_mean) * self.vae.config.scaling_factor / latents_std
+ else:
+ image_latents = components.vae.config.scaling_factor * image_latents
+
+ return image_latents
+
+ # modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_mask_latents
+ # do not accept do_classifier_free_guidance
+ def prepare_mask_latents(
+ self, components, mask, masked_image, batch_size, height, width, dtype, device, generator
+ ):
+ # resize the mask to latents shape as we concatenate the mask to the latents
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+ # and half precision
+ mask = torch.nn.functional.interpolate(
+ mask, size=(height // components.vae_scale_factor, width // components.vae_scale_factor)
+ )
+ mask = mask.to(device=device, dtype=dtype)
+
+ # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
+ if mask.shape[0] < batch_size:
+ if not batch_size % mask.shape[0] == 0:
+ raise ValueError(
+ "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
+ f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
+ " of masks that you pass is divisible by the total requested batch size."
+ )
+ mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
+
+ if masked_image is not None and masked_image.shape[1] == 4:
+ masked_image_latents = masked_image
+ else:
+ masked_image_latents = None
+
+ if masked_image is not None:
+ if masked_image_latents is None:
+ masked_image = masked_image.to(device=device, dtype=dtype)
+ masked_image_latents = self._encode_vae_image(components, masked_image, generator=generator)
+
+ if masked_image_latents.shape[0] < batch_size:
+ if not batch_size % masked_image_latents.shape[0] == 0:
+ raise ValueError(
+ "The passed images and the required batch size don't match. Images are supposed to be duplicated"
+ f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
+ " Make sure the number of images that you pass is divisible by the total requested batch size."
+ )
+ masked_image_latents = masked_image_latents.repeat(
+ batch_size // masked_image_latents.shape[0], 1, 1, 1
+ )
+
+ # aligning device to prevent device errors when concating it with the latent model input
+ masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+
+ return mask, masked_image_latents
+
+ @torch.no_grad()
+ def __call__(self, components: StableDiffusionXLModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.dtype = block_state.dtype if block_state.dtype is not None else components.vae.dtype
+ block_state.device = components._execution_device
+
+ if block_state.height is None:
+ block_state.height = components.default_height
+ if block_state.width is None:
+ block_state.width = components.default_width
+
+ if block_state.padding_mask_crop is not None:
+ block_state.crops_coords = components.mask_processor.get_crop_region(
+ block_state.mask_image, block_state.width, block_state.height, pad=block_state.padding_mask_crop
+ )
+ block_state.resize_mode = "fill"
+ else:
+ block_state.crops_coords = None
+ block_state.resize_mode = "default"
+
+ image = components.image_processor.preprocess(
+ block_state.image,
+ height=block_state.height,
+ width=block_state.width,
+ crops_coords=block_state.crops_coords,
+ resize_mode=block_state.resize_mode,
+ )
+ image = image.to(dtype=torch.float32)
+
+ mask = components.mask_processor.preprocess(
+ block_state.mask_image,
+ height=block_state.height,
+ width=block_state.width,
+ resize_mode=block_state.resize_mode,
+ crops_coords=block_state.crops_coords,
+ )
+ block_state.masked_image = image * (mask < 0.5)
+
+ block_state.batch_size = image.shape[0]
+ image = image.to(device=block_state.device, dtype=block_state.dtype)
+ block_state.image_latents = self._encode_vae_image(components, image=image, generator=block_state.generator)
+
+ # 7. Prepare mask latent variables
+ block_state.mask, block_state.masked_image_latents = self.prepare_mask_latents(
+ components,
+ mask,
+ block_state.masked_image,
+ block_state.batch_size,
+ block_state.height,
+ block_state.width,
+ block_state.dtype,
+ block_state.device,
+ block_state.generator,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py
new file mode 100644
index 0000000000000000000000000000000000000000..d3fc7f3bfbe169d8046da547c6e15917036bb7e9
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_blocks_stable_diffusion_xl.py
@@ -0,0 +1,512 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ StableDiffusionXLControlNetInputStep,
+ StableDiffusionXLControlNetUnionInputStep,
+ StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep,
+ StableDiffusionXLImg2ImgPrepareLatentsStep,
+ StableDiffusionXLImg2ImgSetTimestepsStep,
+ StableDiffusionXLInpaintPrepareLatentsStep,
+ StableDiffusionXLInputStep,
+ StableDiffusionXLPrepareAdditionalConditioningStep,
+ StableDiffusionXLPrepareLatentsStep,
+ StableDiffusionXLSetTimestepsStep,
+)
+from .decoders import (
+ StableDiffusionXLDecodeStep,
+ StableDiffusionXLInpaintOverlayMaskStep,
+)
+from .denoise import (
+ StableDiffusionXLControlNetDenoiseStep,
+ StableDiffusionXLDenoiseStep,
+ StableDiffusionXLInpaintControlNetDenoiseStep,
+ StableDiffusionXLInpaintDenoiseStep,
+)
+from .encoders import (
+ StableDiffusionXLInpaintVaeEncoderStep,
+ StableDiffusionXLIPAdapterStep,
+ StableDiffusionXLTextEncoderStep,
+ StableDiffusionXLVaeEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# auto blocks & sequential blocks & mappings
+
+
+# vae encoder (run before before_denoise)
+class StableDiffusionXLAutoVaeEncoderStep(AutoPipelineBlocks):
+ block_classes = [StableDiffusionXLInpaintVaeEncoderStep, StableDiffusionXLVaeEncoderStep]
+ block_names = ["inpaint", "img2img"]
+ block_trigger_inputs = ["mask_image", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae encoder step that encode the image inputs into their latent representations.\n"
+ + "This is an auto pipeline block that works for both inpainting and img2img tasks.\n"
+ + " - `StableDiffusionXLInpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n"
+ + " - `StableDiffusionXLVaeEncoderStep` (img2img) is used when only `image` is provided."
+ + " - if neither `mask_image` nor `image` is provided, step will be skipped."
+ )
+
+
+# optional ip-adapter (run before input step)
+class StableDiffusionXLAutoIPAdapterStep(AutoPipelineBlocks):
+ block_classes = [StableDiffusionXLIPAdapterStep]
+ block_names = ["ip_adapter"]
+ block_trigger_inputs = ["ip_adapter_image"]
+
+ @property
+ def description(self):
+ return "Run IP Adapter step if `ip_adapter_image` is provided. This step should be placed before the 'input' step.\n"
+
+
+# before_denoise: text2img
+class StableDiffusionXLBeforeDenoiseStep(SequentialPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLSetTimestepsStep,
+ StableDiffusionXLPrepareLatentsStep,
+ StableDiffusionXLPrepareAdditionalConditioningStep,
+ ]
+ block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is a sequential pipeline blocks:\n"
+ + " - `StableDiffusionXLSetTimestepsStep` is used to set the timesteps\n"
+ + " - `StableDiffusionXLPrepareLatentsStep` is used to prepare the latents\n"
+ + " - `StableDiffusionXLPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n"
+ )
+
+
+# before_denoise: img2img
+class StableDiffusionXLImg2ImgBeforeDenoiseStep(SequentialPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLImg2ImgSetTimestepsStep,
+ StableDiffusionXLImg2ImgPrepareLatentsStep,
+ StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep,
+ ]
+ block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step for img2img task.\n"
+ + "This is a sequential pipeline blocks:\n"
+ + " - `StableDiffusionXLImg2ImgSetTimestepsStep` is used to set the timesteps\n"
+ + " - `StableDiffusionXLImg2ImgPrepareLatentsStep` is used to prepare the latents\n"
+ + " - `StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n"
+ )
+
+
+# before_denoise: inpainting
+class StableDiffusionXLInpaintBeforeDenoiseStep(SequentialPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLImg2ImgSetTimestepsStep,
+ StableDiffusionXLInpaintPrepareLatentsStep,
+ StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep,
+ ]
+ block_names = ["set_timesteps", "prepare_latents", "prepare_add_cond"]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step for inpainting task.\n"
+ + "This is a sequential pipeline blocks:\n"
+ + " - `StableDiffusionXLImg2ImgSetTimestepsStep` is used to set the timesteps\n"
+ + " - `StableDiffusionXLInpaintPrepareLatentsStep` is used to prepare the latents\n"
+ + " - `StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep` is used to prepare the additional conditioning\n"
+ )
+
+
+# before_denoise: all task (text2img, img2img, inpainting)
+class StableDiffusionXLAutoBeforeDenoiseStep(AutoPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLInpaintBeforeDenoiseStep,
+ StableDiffusionXLImg2ImgBeforeDenoiseStep,
+ StableDiffusionXLBeforeDenoiseStep,
+ ]
+ block_names = ["inpaint", "img2img", "text2img"]
+ block_trigger_inputs = ["mask", "image_latents", None]
+
+ @property
+ def description(self):
+ return (
+ "Before denoise step that prepare the inputs for the denoise step.\n"
+ + "This is an auto pipeline block that works for text2img, img2img and inpainting tasks as well as controlnet, controlnet_union.\n"
+ + " - `StableDiffusionXLInpaintBeforeDenoiseStep` (inpaint) is used when both `mask` and `image_latents` are provided.\n"
+ + " - `StableDiffusionXLImg2ImgBeforeDenoiseStep` (img2img) is used when only `image_latents` is provided.\n"
+ + " - `StableDiffusionXLBeforeDenoiseStep` (text2img) is used when both `image_latents` and `mask` are not provided.\n"
+ )
+
+
+# optional controlnet input step (after before_denoise, before denoise)
+# works for both controlnet and controlnet_union
+class StableDiffusionXLAutoControlNetInputStep(AutoPipelineBlocks):
+ block_classes = [StableDiffusionXLControlNetUnionInputStep, StableDiffusionXLControlNetInputStep]
+ block_names = ["controlnet_union", "controlnet"]
+ block_trigger_inputs = ["control_mode", "control_image"]
+
+ @property
+ def description(self):
+ return (
+ "Controlnet Input step that prepare the controlnet input.\n"
+ + "This is an auto pipeline block that works for both controlnet and controlnet_union.\n"
+ + " (it should be called right before the denoise step)"
+ + " - `StableDiffusionXLControlNetUnionInputStep` is called to prepare the controlnet input when `control_mode` and `control_image` are provided.\n"
+ + " - `StableDiffusionXLControlNetInputStep` is called to prepare the controlnet input when `control_image` is provided."
+ + " - if neither `control_mode` nor `control_image` is provided, step will be skipped."
+ )
+
+
+# denoise: controlnet (text2img, img2img, inpainting)
+class StableDiffusionXLAutoControlNetDenoiseStep(AutoPipelineBlocks):
+ block_classes = [StableDiffusionXLInpaintControlNetDenoiseStep, StableDiffusionXLControlNetDenoiseStep]
+ block_names = ["inpaint_controlnet_denoise", "controlnet_denoise"]
+ block_trigger_inputs = ["mask", "controlnet_cond"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents with controlnet. "
+ "This is a auto pipeline block that using controlnet for text2img, img2img and inpainting tasks."
+ "This block should not be used without a controlnet_cond input"
+ " - `StableDiffusionXLInpaintControlNetDenoiseStep` (inpaint_controlnet_denoise) is used when mask is provided."
+ " - `StableDiffusionXLControlNetDenoiseStep` (controlnet_denoise) is used when mask is not provided but controlnet_cond is provided."
+ " - If neither mask nor controlnet_cond are provided, step will be skipped."
+ )
+
+
+# denoise: all task with or without controlnet (text2img, img2img, inpainting)
+class StableDiffusionXLAutoDenoiseStep(AutoPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLAutoControlNetDenoiseStep,
+ StableDiffusionXLInpaintDenoiseStep,
+ StableDiffusionXLDenoiseStep,
+ ]
+ block_names = ["controlnet_denoise", "inpaint_denoise", "denoise"]
+ block_trigger_inputs = ["controlnet_cond", "mask", None]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. "
+ "This is a auto pipeline block that works for text2img, img2img and inpainting tasks. And can be used with or without controlnet."
+ " - `StableDiffusionXLAutoControlNetDenoiseStep` (controlnet_denoise) is used when controlnet_cond is provided (support controlnet withtext2img, img2img and inpainting tasks)."
+ " - `StableDiffusionXLInpaintDenoiseStep` (inpaint_denoise) is used when mask is provided (support inpainting tasks)."
+ " - `StableDiffusionXLDenoiseStep` (denoise) is used when neither mask nor controlnet_cond are provided (support text2img and img2img tasks)."
+ )
+
+
+# decode: inpaint
+class StableDiffusionXLInpaintDecodeStep(SequentialPipelineBlocks):
+ block_classes = [StableDiffusionXLDecodeStep, StableDiffusionXLInpaintOverlayMaskStep]
+ block_names = ["decode", "mask_overlay"]
+
+ @property
+ def description(self):
+ return (
+ "Inpaint decode step that decode the denoised latents into images outputs.\n"
+ + "This is a sequential pipeline blocks:\n"
+ + " - `StableDiffusionXLDecodeStep` is used to decode the denoised latents into images\n"
+ + " - `StableDiffusionXLInpaintOverlayMaskStep` is used to overlay the mask on the image"
+ )
+
+
+# decode: all task (text2img, img2img, inpainting)
+class StableDiffusionXLAutoDecodeStep(AutoPipelineBlocks):
+ block_classes = [StableDiffusionXLInpaintDecodeStep, StableDiffusionXLDecodeStep]
+ block_names = ["inpaint", "non-inpaint"]
+ block_trigger_inputs = ["padding_mask_crop", None]
+
+ @property
+ def description(self):
+ return (
+ "Decode step that decode the denoised latents into images outputs.\n"
+ + "This is an auto pipeline block that works for inpainting and non-inpainting tasks.\n"
+ + " - `StableDiffusionXLInpaintDecodeStep` (inpaint) is used when `padding_mask_crop` is provided.\n"
+ + " - `StableDiffusionXLDecodeStep` (non-inpaint) is used when `padding_mask_crop` is not provided."
+ )
+
+
+class StableDiffusionXLCoreDenoiseStep(SequentialPipelineBlocks):
+ block_classes = [
+ StableDiffusionXLInputStep,
+ StableDiffusionXLAutoBeforeDenoiseStep,
+ StableDiffusionXLAutoControlNetInputStep,
+ StableDiffusionXLAutoDenoiseStep,
+ ]
+ block_names = ["input", "before_denoise", "controlnet_input", "denoise"]
+
+ @property
+ def description(self):
+ return (
+ "Core step that performs the denoising process. \n"
+ + " - `StableDiffusionXLInputStep` (input) standardizes the inputs for the denoising step.\n"
+ + " - `StableDiffusionXLAutoBeforeDenoiseStep` (before_denoise) prepares the inputs for the denoising step.\n"
+ + " - `StableDiffusionXLAutoControlNetInputStep` (controlnet_input) prepares the controlnet input.\n"
+ + " - `StableDiffusionXLAutoDenoiseStep` (denoise) iteratively denoises the latents.\n\n"
+ + "This step support text-to-image, image-to-image, inpainting, with or without controlnet/controlnet_union/ip_adapter for Stable Diffusion XL:\n"
+ + "- for image-to-image generation, you need to provide `image_latents`\n"
+ + "- for inpainting, you need to provide `mask_image` and `image_latents`\n"
+ + "- to run the controlnet workflow, you need to provide `control_image`\n"
+ + "- to run the controlnet_union workflow, you need to provide `control_image` and `control_mode`\n"
+ + "- to run the ip_adapter workflow, you need to load ip_adapter into your unet and provide `ip_adapter_embeds`\n"
+ + "- for text-to-image generation, all you need to provide is prompt embeddings\n"
+ )
+
+
+# ip-adapter, controlnet, text2img, img2img, inpainting
+# auto_docstring
+class StableDiffusionXLAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion
+ XL.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+ - `inpainting`: requires `mask_image`, `image`, `prompt`
+ - `controlnet_text2image`: requires `control_image`, `prompt`
+ - `controlnet_image2image`: requires `control_image`, `image`, `prompt`
+ - `controlnet_inpainting`: requires `control_image`, `mask_image`, `image`, `prompt`
+ - `controlnet_union_text2image`: requires `control_image`, `control_mode`, `prompt`
+ - `controlnet_union_image2image`: requires `control_image`, `control_mode`, `image`, `prompt`
+ - `controlnet_union_inpainting`: requires `control_image`, `control_mode`, `mask_image`, `image`, `prompt`
+ - `ip_adapter_text2image`: requires `ip_adapter_image`, `prompt`
+ - `ip_adapter_image2image`: requires `ip_adapter_image`, `image`, `prompt`
+ - `ip_adapter_inpainting`: requires `ip_adapter_image`, `mask_image`, `image`, `prompt`
+ - `ip_adapter_controlnet_text2image`: requires `ip_adapter_image`, `control_image`, `prompt`
+ - `ip_adapter_controlnet_image2image`: requires `ip_adapter_image`, `control_image`, `image`, `prompt`
+ - `ip_adapter_controlnet_inpainting`: requires `ip_adapter_image`, `control_image`, `mask_image`, `image`,
+ `prompt`
+ - `ip_adapter_controlnet_union_text2image`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `prompt`
+ - `ip_adapter_controlnet_union_image2image`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `image`, `prompt`
+ - `ip_adapter_controlnet_union_inpainting`: requires `ip_adapter_image`, `control_image`, `control_mode`,
+ `mask_image`, `image`, `prompt`
+
+ Components:
+ text_encoder (`CLIPTextModel`) text_encoder_2 (`CLIPTextModelWithProjection`) tokenizer (`CLIPTokenizer`)
+ tokenizer_2 (`CLIPTokenizer`) guider (`ClassifierFreeGuidance`) image_encoder
+ (`CLIPVisionModelWithProjection`) feature_extractor (`CLIPImageProcessor`) unet (`UNet2DConditionModel`) vae
+ (`AutoencoderKL`) image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) scheduler
+ (`EulerDiscreteScheduler`) controlnet (`ControlNetUnionModel`) control_image_processor (`VaeImageProcessor`)
+
+ Configs:
+ force_zeros_for_empty_prompt (default: True) requires_aesthetics_score (default: False)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt_2 (`None`, *optional*):
+ TODO: Add description.
+ cross_attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ clip_skip (`None`, *optional*):
+ TODO: Add description.
+ ip_adapter_image (`Image | ndarray | Tensor | list | list | list`, *optional*):
+ The image(s) to be used as ip adapter
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image (`None`, *optional*):
+ TODO: Add description.
+ mask_image (`None`, *optional*):
+ TODO: Add description.
+ padding_mask_crop (`None`, *optional*):
+ TODO: Add description.
+ dtype (`dtype`, *optional*):
+ The dtype of the model inputs
+ generator (`None`, *optional*):
+ TODO: Add description.
+ preprocess_kwargs (`dict | NoneType`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `ImageProcessor` as defined under
+ `self.image_processor` in [diffusers.image_processor.VaeImageProcessor]
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ ip_adapter_embeds (`list`, *optional*):
+ Pre-generated image embeddings for IP-Adapter. Can be generated from ip_adapter step.
+ negative_ip_adapter_embeds (`list`, *optional*):
+ Pre-generated negative image embeddings for IP-Adapter. Can be generated from ip_adapter step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ denoising_end (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.3):
+ TODO: Add description.
+ denoising_start (`None`, *optional*):
+ TODO: Add description.
+ latents (`None`):
+ TODO: Add description.
+ image_latents (`Tensor`, *optional*):
+ The latents representing the reference image for image-to-image/inpainting generation. Can be generated
+ in vae_encode step.
+ mask (`Tensor`, *optional*):
+ The mask for the inpainting generation. Can be generated in vae_encode step.
+ masked_image_latents (`Tensor`, *optional*):
+ The masked image latents for the inpainting generation (only for inpainting-specific unet). Can be
+ generated in vae_encode step.
+ original_size (`None`, *optional*):
+ TODO: Add description.
+ target_size (`None`, *optional*):
+ TODO: Add description.
+ negative_original_size (`None`, *optional*):
+ TODO: Add description.
+ negative_target_size (`None`, *optional*):
+ TODO: Add description.
+ crops_coords_top_left (`None`, *optional*, defaults to (0, 0)):
+ TODO: Add description.
+ negative_crops_coords_top_left (`None`, *optional*, defaults to (0, 0)):
+ TODO: Add description.
+ aesthetic_score (`None`, *optional*, defaults to 6.0):
+ TODO: Add description.
+ negative_aesthetic_score (`None`, *optional*, defaults to 2.0):
+ TODO: Add description.
+ control_image (`None`, *optional*):
+ TODO: Add description.
+ control_mode (`None`, *optional*):
+ TODO: Add description.
+ control_guidance_start (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ control_guidance_end (`None`, *optional*, defaults to 1.0):
+ TODO: Add description.
+ controlnet_conditioning_scale (`None`, *optional*, defaults to 1.0):
+ TODO: Add description.
+ guess_mode (`None`, *optional*, defaults to False):
+ TODO: Add description.
+ crops_coords (`tuple | NoneType`, *optional*):
+ The crop coordinates to use for preprocess/postprocess the image and mask, for inpainting task only. Can
+ be generated in vae_encode step.
+ controlnet_cond (`Tensor`, *optional*):
+ The control image to use for the denoising process. Can be generated in prepare_controlnet_inputs step.
+ conditioning_scale (`float`, *optional*):
+ The controlnet conditioning scale value to use for the denoising process. Can be generated in
+ prepare_controlnet_inputs step.
+ controlnet_keep (`list`, *optional*):
+ The controlnet keep values to use for the denoising process. Can be generated in
+ prepare_controlnet_inputs step.
+ **denoiser_input_fields (`None`, *optional*):
+ All conditional model inputs that need to be prepared with guider. It should contain
+ prompt_embeds/negative_prompt_embeds, add_time_ids/negative_add_time_ids,
+ pooled_prompt_embeds/negative_pooled_prompt_embeds, and ip_adapter_embeds/negative_ip_adapter_embeds
+ (optional).please add `kwargs_type=denoiser_input_fields` to their parameter spec (`OutputParam`) when
+ they are created and added to the pipeline state
+ eta (`None`, *optional*, defaults to 0.0):
+ TODO: Add description.
+ output_type (`None`, *optional*, defaults to pil):
+ TODO: Add description.
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ block_classes = [
+ StableDiffusionXLTextEncoderStep,
+ StableDiffusionXLAutoIPAdapterStep,
+ StableDiffusionXLAutoVaeEncoderStep,
+ StableDiffusionXLCoreDenoiseStep,
+ StableDiffusionXLAutoDecodeStep,
+ ]
+ block_names = [
+ "text_encoder",
+ "ip_adapter",
+ "vae_encoder",
+ "denoise",
+ "decode",
+ ]
+
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ "inpainting": {"mask_image": True, "image": True, "prompt": True},
+ "controlnet_text2image": {"control_image": True, "prompt": True},
+ "controlnet_image2image": {"control_image": True, "image": True, "prompt": True},
+ "controlnet_inpainting": {"control_image": True, "mask_image": True, "image": True, "prompt": True},
+ "controlnet_union_text2image": {"control_image": True, "control_mode": True, "prompt": True},
+ "controlnet_union_image2image": {"control_image": True, "control_mode": True, "image": True, "prompt": True},
+ "controlnet_union_inpainting": {
+ "control_image": True,
+ "control_mode": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_text2image": {"ip_adapter_image": True, "prompt": True},
+ "ip_adapter_image2image": {"ip_adapter_image": True, "image": True, "prompt": True},
+ "ip_adapter_inpainting": {"ip_adapter_image": True, "mask_image": True, "image": True, "prompt": True},
+ "ip_adapter_controlnet_text2image": {"ip_adapter_image": True, "control_image": True, "prompt": True},
+ "ip_adapter_controlnet_image2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_inpainting": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_text2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_image2image": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "image": True,
+ "prompt": True,
+ },
+ "ip_adapter_controlnet_union_inpainting": {
+ "ip_adapter_image": True,
+ "control_image": True,
+ "control_mode": True,
+ "mask_image": True,
+ "image": True,
+ "prompt": True,
+ },
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and controlnet tasks using Stable Diffusion XL."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a955edacb5517f0d9435b481249f4a2c9d769ee
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/stable_diffusion_xl/modular_pipeline.py
@@ -0,0 +1,360 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import PIL
+import torch
+
+from ...image_processor import PipelineImageInput
+from ...loaders import ModularIPAdapterMixin, StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin
+from ...pipelines.pipeline_utils import StableDiffusionMixin
+from ...pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+from ..modular_pipeline_utils import InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# YiYi TODO: move to a different file? stable_diffusion_xl_module should have its own folder?
+# YiYi Notes: model specific components:
+## (1) it should inherit from ModularPipeline
+## (2) acts like a container that holds components and configs
+## (3) define default config (related to components), e.g. default_sample_size, vae_scale_factor, num_channels_unet, num_channels_latents
+## (4) inherit from model-specic loader class (e.g. StableDiffusionXLLoraLoaderMixin)
+## (5) how to use together with Components_manager?
+class StableDiffusionXLModularPipeline(
+ ModularPipeline,
+ StableDiffusionMixin,
+ TextualInversionLoaderMixin,
+ StableDiffusionXLLoraLoaderMixin,
+ ModularIPAdapterMixin,
+):
+ """
+ A ModularPipeline for Stable Diffusion XL.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "StableDiffusionXLAutoBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_width(self):
+ return self.default_sample_size * self.vae_scale_factor
+
+ @property
+ def default_sample_size(self):
+ default_sample_size = 128
+ if hasattr(self, "unet") and self.unet is not None:
+ default_sample_size = self.unet.config.sample_size
+ return default_sample_size
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
+
+ # YiYi TODO: change to num_channels_latents
+ @property
+ def num_channels_unet(self):
+ num_channels_unet = 4
+ if hasattr(self, "unet") and self.unet is not None:
+ num_channels_unet = self.unet.config.in_channels
+ return num_channels_unet
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 4
+ if hasattr(self, "vae") and self.vae is not None:
+ num_channels_latents = self.vae.config.latent_channels
+ return num_channels_latents
+
+
+# YiYi/Sayak TODO: not used yet, maintain a list of schema that can be used across all pipeline blocks
+# auto_docstring
+SDXL_INPUTS_SCHEMA = {
+ "prompt": InputParam(
+ "prompt", type_hint=str | list[str], description="The prompt or prompts to guide the image generation"
+ ),
+ "prompt_2": InputParam(
+ "prompt_2",
+ type_hint=str | list[str],
+ description="The prompt or prompts to be sent to the tokenizer_2 and text_encoder_2",
+ ),
+ "negative_prompt": InputParam(
+ "negative_prompt",
+ type_hint=str | list[str],
+ description="The prompt or prompts not to guide the image generation",
+ ),
+ "negative_prompt_2": InputParam(
+ "negative_prompt_2",
+ type_hint=str | list[str],
+ description="The negative prompt or prompts for text_encoder_2",
+ ),
+ "cross_attention_kwargs": InputParam(
+ "cross_attention_kwargs",
+ type_hint=dict | None,
+ description="Kwargs dictionary passed to the AttentionProcessor",
+ ),
+ "clip_skip": InputParam(
+ "clip_skip", type_hint=int | None, description="Number of layers to skip in CLIP text encoder"
+ ),
+ "image": InputParam(
+ "image",
+ type_hint=PipelineImageInput,
+ required=True,
+ description="The image(s) to modify for img2img or inpainting",
+ ),
+ "mask_image": InputParam(
+ "mask_image",
+ type_hint=PipelineImageInput,
+ required=True,
+ description="Mask image for inpainting, white pixels will be repainted",
+ ),
+ "generator": InputParam(
+ "generator",
+ type_hint=torch.Generator | list[torch.Generator] | None,
+ description="Generator(s) for deterministic generation",
+ ),
+ "height": InputParam("height", type_hint=int | None, description="Height in pixels of the generated image"),
+ "width": InputParam("width", type_hint=int | None, description="Width in pixels of the generated image"),
+ "num_images_per_prompt": InputParam(
+ "num_images_per_prompt", type_hint=int, default=1, description="Number of images to generate per prompt"
+ ),
+ "num_inference_steps": InputParam(
+ "num_inference_steps", type_hint=int, default=50, description="Number of denoising steps"
+ ),
+ "timesteps": InputParam(
+ "timesteps", type_hint=torch.Tensor | None, description="Custom timesteps for the denoising process"
+ ),
+ "sigmas": InputParam(
+ "sigmas", type_hint=torch.Tensor | None, description="Custom sigmas for the denoising process"
+ ),
+ "denoising_end": InputParam(
+ "denoising_end",
+ type_hint=float | None,
+ description="Fraction of denoising process to complete before termination",
+ ),
+ # YiYi Notes: img2img defaults to 0.3, inpainting defaults to 0.9999
+ "strength": InputParam(
+ "strength", type_hint=float, default=0.3, description="How much to transform the reference image"
+ ),
+ "denoising_start": InputParam(
+ "denoising_start", type_hint=float | None, description="Starting point of the denoising process"
+ ),
+ "latents": InputParam(
+ "latents", type_hint=torch.Tensor | None, description="Pre-generated noisy latents for image generation"
+ ),
+ "padding_mask_crop": InputParam(
+ "padding_mask_crop",
+ type_hint=tuple[int, int] | None,
+ description="Size of margin in crop for image and mask",
+ ),
+ "original_size": InputParam(
+ "original_size",
+ type_hint=tuple[int, int] | None,
+ description="Original size of the image for SDXL's micro-conditioning",
+ ),
+ "target_size": InputParam(
+ "target_size", type_hint=tuple[int, int] | None, description="Target size for SDXL's micro-conditioning"
+ ),
+ "negative_original_size": InputParam(
+ "negative_original_size",
+ type_hint=tuple[int, int] | None,
+ description="Negative conditioning based on image resolution",
+ ),
+ "negative_target_size": InputParam(
+ "negative_target_size",
+ type_hint=tuple[int, int] | None,
+ description="Negative conditioning based on target resolution",
+ ),
+ "crops_coords_top_left": InputParam(
+ "crops_coords_top_left",
+ type_hint=tuple[int, int],
+ default=(0, 0),
+ description="Top-left coordinates for SDXL's micro-conditioning",
+ ),
+ "negative_crops_coords_top_left": InputParam(
+ "negative_crops_coords_top_left",
+ type_hint=tuple[int, int],
+ default=(0, 0),
+ description="Negative conditioning crop coordinates",
+ ),
+ "aesthetic_score": InputParam(
+ "aesthetic_score", type_hint=float, default=6.0, description="Simulates aesthetic score of generated image"
+ ),
+ "negative_aesthetic_score": InputParam(
+ "negative_aesthetic_score", type_hint=float, default=2.0, description="Simulates negative aesthetic score"
+ ),
+ "eta": InputParam("eta", type_hint=float, default=0.0, description="Parameter η in the DDIM paper"),
+ "output_type": InputParam(
+ "output_type", type_hint=str, default="pil", description="Output format (pil/tensor/np.array)"
+ ),
+ "ip_adapter_image": InputParam(
+ "ip_adapter_image",
+ type_hint=PipelineImageInput,
+ required=True,
+ description="Image(s) to be used as IP adapter",
+ ),
+ "control_image": InputParam(
+ "control_image", type_hint=PipelineImageInput, required=True, description="ControlNet input condition"
+ ),
+ "control_guidance_start": InputParam(
+ "control_guidance_start",
+ type_hint=float | list[float],
+ default=0.0,
+ description="When ControlNet starts applying",
+ ),
+ "control_guidance_end": InputParam(
+ "control_guidance_end",
+ type_hint=float | list[float],
+ default=1.0,
+ description="When ControlNet stops applying",
+ ),
+ "controlnet_conditioning_scale": InputParam(
+ "controlnet_conditioning_scale",
+ type_hint=float | list[float],
+ default=1.0,
+ description="Scale factor for ControlNet outputs",
+ ),
+ "guess_mode": InputParam(
+ "guess_mode",
+ type_hint=bool,
+ default=False,
+ description="Enables ControlNet encoder to recognize input without prompts",
+ ),
+ "control_mode": InputParam(
+ "control_mode", type_hint=list[int], required=True, description="Control mode for union controlnet"
+ ),
+ "prompt_embeds": InputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ required=True,
+ description="Text embeddings used to guide image generation",
+ ),
+ "negative_prompt_embeds": InputParam(
+ "negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text embeddings"
+ ),
+ "pooled_prompt_embeds": InputParam(
+ "pooled_prompt_embeds", type_hint=torch.Tensor, required=True, description="Pooled text embeddings"
+ ),
+ "negative_pooled_prompt_embeds": InputParam(
+ "negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="Negative pooled text embeddings"
+ ),
+ "batch_size": InputParam("batch_size", type_hint=int, required=True, description="Number of prompts"),
+ "dtype": InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"),
+ "preprocess_kwargs": InputParam(
+ "preprocess_kwargs", type_hint=dict | None, description="Kwargs for ImageProcessor"
+ ),
+ "latent_timestep": InputParam(
+ "latent_timestep", type_hint=torch.Tensor, required=True, description="Initial noise level timestep"
+ ),
+ "image_latents": InputParam(
+ "image_latents", type_hint=torch.Tensor, required=True, description="Latents representing reference image"
+ ),
+ "mask": InputParam("mask", type_hint=torch.Tensor, required=True, description="Mask for inpainting"),
+ "masked_image_latents": InputParam(
+ "masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting"
+ ),
+ "add_time_ids": InputParam(
+ "add_time_ids", type_hint=torch.Tensor, required=True, description="Time ids for conditioning"
+ ),
+ "negative_add_time_ids": InputParam(
+ "negative_add_time_ids", type_hint=torch.Tensor, description="Negative time ids"
+ ),
+ "timestep_cond": InputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"),
+ "noise": InputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"),
+ "crops_coords": InputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"),
+ "ip_adapter_embeds": InputParam(
+ "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter"
+ ),
+ "negative_ip_adapter_embeds": InputParam(
+ "negative_ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ description="Negative image embeddings for IP-Adapter",
+ ),
+ "images": InputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
+ required=True,
+ description="Generated images",
+ ),
+}
+
+
+SDXL_INTERMEDIATE_OUTPUTS_SCHEMA = {
+ "prompt_embeds": OutputParam(
+ "prompt_embeds", type_hint=torch.Tensor, description="Text embeddings used to guide image generation"
+ ),
+ "negative_prompt_embeds": OutputParam(
+ "negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text embeddings"
+ ),
+ "pooled_prompt_embeds": OutputParam(
+ "pooled_prompt_embeds", type_hint=torch.Tensor, description="Pooled text embeddings"
+ ),
+ "negative_pooled_prompt_embeds": OutputParam(
+ "negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="Negative pooled text embeddings"
+ ),
+ "batch_size": OutputParam("batch_size", type_hint=int, description="Number of prompts"),
+ "dtype": OutputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"),
+ "image_latents": OutputParam(
+ "image_latents", type_hint=torch.Tensor, description="Latents representing reference image"
+ ),
+ "mask": OutputParam("mask", type_hint=torch.Tensor, description="Mask for inpainting"),
+ "masked_image_latents": OutputParam(
+ "masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting"
+ ),
+ "crops_coords": OutputParam("crops_coords", type_hint=tuple[int] | None, description="Crop coordinates"),
+ "timesteps": OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for inference"),
+ "num_inference_steps": OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps"),
+ "latent_timestep": OutputParam(
+ "latent_timestep", type_hint=torch.Tensor, description="Initial noise level timestep"
+ ),
+ "add_time_ids": OutputParam("add_time_ids", type_hint=torch.Tensor, description="Time ids for conditioning"),
+ "negative_add_time_ids": OutputParam(
+ "negative_add_time_ids", type_hint=torch.Tensor, description="Negative time ids"
+ ),
+ "timestep_cond": OutputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"),
+ "latents": OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents"),
+ "noise": OutputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"),
+ "ip_adapter_embeds": OutputParam(
+ "ip_adapter_embeds", type_hint=list[torch.Tensor], description="Image embeddings for IP-Adapter"
+ ),
+ "negative_ip_adapter_embeds": OutputParam(
+ "negative_ip_adapter_embeds",
+ type_hint=list[torch.Tensor],
+ description="Negative image embeddings for IP-Adapter",
+ ),
+ "images": OutputParam(
+ "images",
+ type_hint=list[PIL.Image.Image] | list[torch.Tensor] | list[np.array],
+ description="Generated images",
+ ),
+}
+
+
+SDXL_OUTPUTS_SCHEMA = {
+ "images": OutputParam(
+ "images",
+ type_hint=tuple[list[PIL.Image.Image] | list[torch.Tensor] | list[np.array]] | StableDiffusionXLPipelineOutput,
+ description="The final generated images",
+ )
+}
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e65bcd3723df36cf5285c75978bf17eaf8eb770
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/__init__.py
@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_wan"] = ["WanBlocks"]
+ _import_structure["modular_blocks_wan22"] = ["Wan22Blocks"]
+ _import_structure["modular_blocks_wan22_i2v"] = ["Wan22Image2VideoBlocks"]
+ _import_structure["modular_blocks_wan_i2v"] = ["WanImage2VideoAutoBlocks"]
+ _import_structure["modular_pipeline"] = [
+ "Wan22Image2VideoModularPipeline",
+ "Wan22ModularPipeline",
+ "WanImage2VideoModularPipeline",
+ "WanModularPipeline",
+ ]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_wan import WanBlocks
+ from .modular_blocks_wan22 import Wan22Blocks
+ from .modular_blocks_wan22_i2v import Wan22Image2VideoBlocks
+ from .modular_blocks_wan_i2v import WanImage2VideoAutoBlocks
+ from .modular_pipeline import (
+ Wan22Image2VideoModularPipeline,
+ Wan22ModularPipeline,
+ WanImage2VideoModularPipeline,
+ WanModularPipeline,
+ )
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..92d4e88fc8049136aa45843622e26671844054d0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/before_denoise.py
@@ -0,0 +1,554 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from ...models import WanTransformer3DModel
+from ...schedulers import UniPCMultistepScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import WanModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# TODO(yiyi, aryan): We need another step before text encoder to set the `num_inference_steps` attribute for guider so that
+# things like when to do guidance and how many conditions to be prepared can be determined. Currently, this is done by
+# always assuming you want to do guidance in the Guiders. So, negative embeddings are prepared regardless of what the
+# configuration of guider is.
+
+
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_videos_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_videos_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_videos_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_videos_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_videos_per_prompt (int, optional): Number of videos to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_videos_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_videos_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_videos_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_videos_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_videos_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_videos_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+def calculate_dimension_from_latents(
+ latents: torch.Tensor, vae_scale_factor_temporal: int, vae_scale_factor_spatial: int
+) -> tuple[int, int]:
+ """Calculate image dimensions from latent tensor dimensions.
+
+ This function converts latent temporal and spatial dimensions to image temporal and spatial dimensions by
+ multiplying the latent num_frames/height/width by the VAE scale factor.
+
+ Args:
+ latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions.
+ Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width]
+ vae_scale_factor_temporal (int): The scale factor used by the VAE to compress temporal dimension.
+ Typically 4 for most VAEs (video is 4x larger than latents in temporal dimension)
+ vae_scale_factor_spatial (int): The scale factor used by the VAE to compress spatial dimension.
+ Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
+
+ Returns:
+ tuple[int, int]: The calculated image dimensions as (height, width)
+
+ Raises:
+ ValueError: If latents tensor doesn't have 4 or 5 dimensions
+
+ """
+ if latents.ndim != 5:
+ raise ValueError(f"latents must have 5 dimensions, but got {latents.ndim}")
+
+ _, _, num_latent_frames, latent_height, latent_width = latents.shape
+
+ num_frames = (num_latent_frames - 1) * vae_scale_factor_temporal + 1
+ height = latent_height * vae_scale_factor_spatial
+ width = latent_width * vae_scale_factor_spatial
+
+ return num_frames, height, width
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class WanTextInputStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_videos_per_prompt`\n\n"
+ "All input tensors are expected to have either batch_size=1 or match the batch_size\n"
+ "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n"
+ "have a final batch_size of batch_size * num_videos_per_prompt."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", WanTransformer3DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_videos_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=torch.Tensor,
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `transformer.dtype`)",
+ ),
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.prompt_embeds is not None and block_state.negative_prompt_embeds is not None:
+ if block_state.prompt_embeds.shape != block_state.negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {block_state.negative_prompt_embeds.shape}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
+ block_state.dtype = block_state.prompt_embeds.dtype
+
+ _, seq_len, _ = block_state.prompt_embeds.shape
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_videos_per_prompt, 1)
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
+ block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1
+ )
+
+ if block_state.negative_prompt_embeds is not None:
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+ 1, block_state.num_videos_per_prompt, 1
+ )
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+ block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class WanAdditionalInputsStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[str] = ["image_condition_latents"],
+ additional_batch_inputs: list[str] = [],
+ ):
+ """Initialize a configurable step that standardizes the inputs for the denoising step. It:\n"
+
+ This step handles multiple common tasks to prepare inputs for the denoising step:
+ 1. For encoded image latents, use it update height/width if None, and expands batch size
+ 2. For additional_batch_inputs: Only expands batch dimensions to match final batch size
+
+ This is a dynamic block that allows you to configure which inputs to process.
+
+ Args:
+ image_latent_inputs (list[str], optional): Names of image latent tensors to process.
+ In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be
+ a single string or list of strings. Defaults to ["image_condition_latents"].
+ additional_batch_inputs (List[str], optional):
+ Names of additional conditional input tensors to expand batch size. These tensors will only have their
+ batch dimensions adjusted to match the final batch size. Can be a single string or list of strings.
+ Defaults to [].
+
+ Examples:
+ # Configure to process image_condition_latents (default behavior) WanAdditionalInputsStep() # Configure to
+ process image latents and additional batch inputs WanAdditionalInputsStep(
+ image_latent_inputs=["image_condition_latents"], additional_batch_inputs=["image_embeds"]
+ )
+ """
+ if not isinstance(image_latent_inputs, list):
+ image_latent_inputs = [image_latent_inputs]
+ if not isinstance(additional_batch_inputs, list):
+ additional_batch_inputs = [additional_batch_inputs]
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ # Functionality section
+ summary_section = (
+ "Input processing step that:\n"
+ " 1. For image latent inputs: Updates height/width if None, and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ # Inputs info
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {self._image_latent_inputs}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {self._additional_batch_inputs}"
+
+ # Placement guidance
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(name="num_videos_per_prompt", default=1),
+ InputParam(name="batch_size", required=True),
+ InputParam(name="height"),
+ InputParam(name="width"),
+ InputParam(name="num_frames"),
+ ]
+
+ # Add image latent inputs
+ for image_latent_input_name in self._image_latent_inputs:
+ inputs.append(InputParam(name=image_latent_input_name))
+
+ # Add additional batch inputs
+ for input_name in self._additional_batch_inputs:
+ inputs.append(InputParam(name=input_name))
+
+ return inputs
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs (height/width calculation, patchify, and batch expansion)
+ for image_latent_input_name in self._image_latent_inputs:
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate num_frames, height/width from latents
+ num_frames, height, width = calculate_dimension_from_latents(
+ image_latent_tensor, components.vae_scale_factor_temporal, components.vae_scale_factor_spatial
+ )
+ block_state.num_frames = block_state.num_frames or num_frames
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ # 3. Expand batch size
+ image_latent_tensor = repeat_tensor_to_batch_size(
+ input_name=image_latent_input_name,
+ input_tensor=image_latent_tensor,
+ num_videos_per_prompt=block_state.num_videos_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+ # Process additional batch inputs (only batch expansion)
+ for input_name in self._additional_batch_inputs:
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ # Only expand batch size
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_videos_per_prompt=block_state.num_videos_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_inference_steps", default=50),
+ InputParam("timesteps"),
+ InputParam("sigmas"),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ device,
+ block_state.timesteps,
+ block_state.sigmas,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanPrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the latents for the text-to-video generation process"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam("num_frames", type_hint=int),
+ InputParam("latents", type_hint=torch.Tensor | None),
+ InputParam("num_videos_per_prompt", type_hint=int, default=1),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_videos_per_prompt`. Can be generated in input step.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ )
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+ block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+ )
+ if block_state.num_frames is not None and (
+ block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_scale_factor_temporal != 0
+ ):
+ raise ValueError(
+ f"`num_frames` has to be greater than 0, and (num_frames - 1) must be divisible by {components.vae_scale_factor_temporal}, but got {block_state.num_frames}."
+ )
+
+ @staticmethod
+ # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.prepare_latents with self->comp
+ def prepare_latents(
+ comp,
+ batch_size: int,
+ num_channels_latents: int = 16,
+ height: int = 480,
+ width: int = 832,
+ num_frames: int = 81,
+ dtype: torch.dtype | None = None,
+ device: torch.device | None = None,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ if latents is not None:
+ return latents.to(device=device, dtype=dtype)
+
+ num_latent_frames = (num_frames - 1) // comp.vae_scale_factor_temporal + 1
+ shape = (
+ batch_size,
+ num_channels_latents,
+ num_latent_frames,
+ int(height) // comp.vae_scale_factor_spatial,
+ int(width) // comp.vae_scale_factor_spatial,
+ )
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ device = components._execution_device
+ dtype = torch.float32 # Wan latents should be torch.float32 for best quality
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+ block_state.num_frames = block_state.num_frames or components.default_num_frames
+
+ block_state.latents = self.prepare_latents(
+ components,
+ batch_size=block_state.batch_size * block_state.num_videos_per_prompt,
+ num_channels_latents=components.num_channels_latents,
+ height=block_state.height,
+ width=block_state.width,
+ num_frames=block_state.num_frames,
+ dtype=dtype,
+ device=device,
+ generator=block_state.generator,
+ latents=block_state.latents,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..e265cbb8a5ca4a5d813c4dd6bb36745cad8a1cb3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/decoders.py
@@ -0,0 +1,98 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...models import AutoencoderKLWan
+from ...utils import logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class WanVaeDecoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The denoised latents from the denoising step",
+ ),
+ InputParam(
+ "output_type", default="np", type_hint=str, description="The output type of the decoded videos"
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "videos",
+ type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray],
+ description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae_dtype = components.vae.dtype
+
+ latents = block_state.latents
+ latents_mean = (
+ torch.tensor(components.vae.config.latents_mean)
+ .view(1, components.vae.config.z_dim, 1, 1, 1)
+ .to(latents.device, latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+ 1, components.vae.config.z_dim, 1, 1, 1
+ ).to(latents.device, latents.dtype)
+ latents = latents / latents_std + latents_mean
+ latents = latents.to(vae_dtype)
+ block_state.videos = components.vae.decode(latents, return_dict=False)[0]
+
+ output_type = getattr(block_state, "output_type", "np")
+ block_state.videos = components.video_processor.postprocess_video(block_state.videos, output_type=output_type)
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..a38b9ee7ddc8fa7d6d5c2da75d4208e8584b1ded
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/denoise.py
@@ -0,0 +1,544 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import WanTransformer3DModel
+from ...schedulers import UniPCMultistepScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam
+from .modular_pipeline import WanModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class WanLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `WanDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "dtype",
+ required=True,
+ type_hint=torch.dtype,
+ description="The dtype of the model inputs. Can be generated in input step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = block_state.latents.to(block_state.dtype)
+ return components, block_state
+
+
+class WanImage2VideoLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `WanDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "image_condition_latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The image condition latents to use for the denoising process. Can be generated in prepare_first_frame_latents/prepare_first_last_frame_latents step.",
+ ),
+ InputParam(
+ "dtype",
+ required=True,
+ type_hint=torch.dtype,
+ description="The dtype of the model inputs. Can be generated in input step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ block_state.latent_model_input = torch.cat(
+ [block_state.latents, block_state.image_condition_latents], dim=1
+ ).to(block_state.dtype)
+ return components, block_state
+
+
+class WanLoopDenoiser(ModularPipelineBlocks):
+ model_name = "wan"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
+ ):
+ """Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.1.
+
+ Args:
+ guider_input_fields: A dictionary that maps each argument expected by the denoiser model
+ (for example, "encoder_hidden_states") to data stored on 'block_state'. The value can be either:
+
+ - A tuple of strings. For instance, {"encoder_hidden_states": ("prompt_embeds",
+ "negative_prompt_embeds")} tells the guider to read `block_state.prompt_embeds` and
+ `block_state.negative_prompt_embeds` and pass them as the conditional and unconditional batches of
+ 'encoder_hidden_states'.
+ - A string. For example, {"encoder_hidden_image": "image_embeds"} makes the guider forward
+ `block_state.image_embeds` for both conditional and unconditional batches.
+ """
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", WanTransformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents with guidance. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `WanDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ inputs = [
+ InputParam("attention_kwargs"),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+ guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.extend(value)
+ else:
+ guider_input_names.append(value)
+
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+ # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+ # you will get a guider_state with two batches:
+ # guider_state = [
+ # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch
+ # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch
+ # ]
+ # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ # run the denoiser for each guidance batch
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = guider_state_batch.as_dict()
+ cond_kwargs = {
+ k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v
+ for k, v in cond_kwargs.items()
+ if k in self._guider_input_fields.keys()
+ }
+
+ # Predict the noise residual
+ # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+ guider_state_batch.noise_pred = components.transformer(
+ hidden_states=block_state.latent_model_input.to(block_state.dtype),
+ timestep=t.expand(block_state.latent_model_input.shape[0]).to(block_state.dtype),
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ components.guider.cleanup_models(components.transformer)
+
+ # Perform guidance
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class Wan22LoopDenoiser(ModularPipelineBlocks):
+ model_name = "wan"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
+ ):
+ """Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.2.
+
+ Args:
+ guider_input_fields: A dictionary that maps each argument expected by the denoiser model
+ (for example, "encoder_hidden_states") to data stored on `block_state`. The value can be either:
+
+ - A tuple of strings. For instance, `{"encoder_hidden_states": ("prompt_embeds",
+ "negative_prompt_embeds")}` tells the guider to read `block_state.prompt_embeds` and
+ `block_state.negative_prompt_embeds` and pass them as the conditional and unconditional batches of
+ `encoder_hidden_states`.
+ - A string. For example, `{"encoder_hidden_image": "image_embeds"}` makes the guider forward
+ `block_state.image_embeds` for both conditional and unconditional batches.
+ """
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 4.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec(
+ "guider_2",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 3.0}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", WanTransformer3DModel),
+ ComponentSpec("transformer_2", WanTransformer3DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents with guidance. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `WanDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def expected_configs(self) -> list[ConfigSpec]:
+ return [
+ ConfigSpec(
+ name="boundary_ratio",
+ default=0.875,
+ description="The boundary ratio to divide the denoising loop into high noise and low noise stages.",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ inputs = [
+ InputParam("attention_kwargs"),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+ guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.extend(value)
+ else:
+ guider_input_names.append(value)
+
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ boundary_timestep = components.config.boundary_ratio * components.num_train_timesteps
+ if t >= boundary_timestep:
+ block_state.current_model = components.transformer
+ block_state.guider = components.guider
+ else:
+ block_state.current_model = components.transformer_2
+ block_state.guider = components.guider_2
+
+ block_state.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+ # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+ # you will get a guider_state with two batches:
+ # guider_state = [
+ # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch
+ # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch
+ # ]
+ # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+ guider_state = block_state.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ # run the denoiser for each guidance batch
+ for guider_state_batch in guider_state:
+ block_state.guider.prepare_models(block_state.current_model)
+ cond_kwargs = guider_state_batch.as_dict()
+ cond_kwargs = {
+ k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v
+ for k, v in cond_kwargs.items()
+ if k in self._guider_input_fields.keys()
+ }
+
+ # Predict the noise residual
+ # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+ guider_state_batch.noise_pred = block_state.current_model(
+ hidden_states=block_state.latent_model_input.to(block_state.dtype),
+ timestep=t.expand(block_state.latent_model_input.shape[0]).to(block_state.dtype),
+ attention_kwargs=block_state.attention_kwargs,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ block_state.guider.cleanup_models(block_state.current_model)
+
+ # Perform guidance
+ block_state.noise_pred = block_state.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class WanLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that update the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `WanDenoiseLoopWrapper`)"
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # Perform scheduler step using the predicted output
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred.float(),
+ t,
+ block_state.latents.float(),
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class WanDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoise the latents over `timesteps`. "
+ "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class WanDenoiseStep(WanDenoiseLoopWrapper):
+ block_classes = [
+ WanLoopBeforeDenoiser,
+ WanLoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ }
+ ),
+ WanLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `WanLoopBeforeDenoiser`\n"
+ " - `WanLoopDenoiser`\n"
+ " - `WanLoopAfterDenoiser`\n"
+ "This block supports text-to-video tasks for wan2.1."
+ )
+
+
+class Wan22DenoiseStep(WanDenoiseLoopWrapper):
+ block_classes = [
+ WanLoopBeforeDenoiser,
+ Wan22LoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ }
+ ),
+ WanLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `WanLoopBeforeDenoiser`\n"
+ " - `Wan22LoopDenoiser`\n"
+ " - `WanLoopAfterDenoiser`\n"
+ "This block supports text-to-video tasks for Wan2.2."
+ )
+
+
+class WanImage2VideoDenoiseStep(WanDenoiseLoopWrapper):
+ block_classes = [
+ WanImage2VideoLoopBeforeDenoiser,
+ WanLoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ "encoder_hidden_states_image": "image_embeds",
+ }
+ ),
+ WanLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `WanImage2VideoLoopBeforeDenoiser`\n"
+ " - `WanLoopDenoiser`\n"
+ " - `WanLoopAfterDenoiser`\n"
+ "This block supports image-to-video tasks for wan2.1."
+ )
+
+
+class Wan22Image2VideoDenoiseStep(WanDenoiseLoopWrapper):
+ block_classes = [
+ WanImage2VideoLoopBeforeDenoiser,
+ Wan22LoopDenoiser(
+ guider_input_fields={
+ "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+ }
+ ),
+ WanLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `WanImage2VideoLoopBeforeDenoiser`\n"
+ " - `WanLoopDenoiser`\n"
+ " - `WanLoopAfterDenoiser`\n"
+ "This block supports image-to-video tasks for Wan2.2."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..6df444e56b435e337876de750c454bf5ac83f1ed
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/encoders.py
@@ -0,0 +1,758 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import html
+
+import numpy as np
+import PIL
+import regex as re
+import torch
+from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import PipelineImageInput
+from ...models import AutoencoderKLWan
+from ...utils import is_ftfy_available, is_torchvision_available, logging
+from ...video_processor import VideoProcessor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import WanModularPipeline
+
+
+if is_ftfy_available():
+ import ftfy
+
+if is_torchvision_available():
+ from torchvision import transforms
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def basic_clean(text):
+ text = ftfy.fix_text(text)
+ text = html.unescape(html.unescape(text))
+ return text.strip()
+
+
+def whitespace_clean(text):
+ text = re.sub(r"\s+", " ", text)
+ text = text.strip()
+ return text
+
+
+def prompt_clean(text):
+ text = whitespace_clean(basic_clean(text))
+ return text
+
+
+def get_t5_prompt_embeds(
+ text_encoder: UMT5EncoderModel,
+ tokenizer: AutoTokenizer,
+ prompt: str | list[str],
+ max_sequence_length: int,
+ device: torch.device,
+):
+ dtype = text_encoder.dtype
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ prompt = [prompt_clean(u) for u in prompt]
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+ seq_lens = mask.gt(0).sum(dim=1).long()
+ prompt_embeds = text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+ prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
+ )
+
+ return prompt_embeds
+
+
+def encode_image(
+ image: PipelineImageInput,
+ image_processor: CLIPImageProcessor,
+ image_encoder: CLIPVisionModel,
+ device: torch.device | None = None,
+):
+ image = image_processor(images=image, return_tensors="pt").to(device)
+ image_embeds = image_encoder(**image, output_hidden_states=True)
+ return image_embeds.hidden_states[-2]
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def encode_vae_image(
+ video_tensor: torch.Tensor,
+ vae: AutoencoderKLWan,
+ generator: torch.Generator,
+ device: torch.device,
+ dtype: torch.dtype,
+ latent_channels: int = 16,
+):
+ if not isinstance(video_tensor, torch.Tensor):
+ raise ValueError(f"Expected video_tensor to be a tensor, got {type(video_tensor)}.")
+
+ if isinstance(generator, list) and len(generator) != video_tensor.shape[0]:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but it is not same as number of images {video_tensor.shape[0]}."
+ )
+
+ video_tensor = video_tensor.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list):
+ video_latents = [
+ retrieve_latents(vae.encode(video_tensor[i : i + 1]), generator=generator[i], sample_mode="argmax")
+ for i in range(video_tensor.shape[0])
+ ]
+ video_latents = torch.cat(video_latents, dim=0)
+ else:
+ video_latents = retrieve_latents(vae.encode(video_tensor), sample_mode="argmax")
+
+ latents_mean = (
+ torch.tensor(vae.config.latents_mean)
+ .view(1, latent_channels, 1, 1, 1)
+ .to(video_latents.device, video_latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, latent_channels, 1, 1, 1).to(
+ video_latents.device, video_latents.dtype
+ )
+ video_latents = (video_latents - latents_mean) * latents_std
+
+ return video_latents
+
+
+class WanTextEncoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generate text_embeddings to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", UMT5EncoderModel),
+ ComponentSpec("tokenizer", AutoTokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("negative_prompt"),
+ InputParam("max_sequence_length", default=512),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=torch.Tensor,
+ kwargs_type="denoiser_input_fields",
+ description="negative text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ if block_state.prompt is not None and (
+ not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list)
+ ):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str,
+ device: torch.device | None = None,
+ prepare_unconditional_embeds: bool = True,
+ negative_prompt: str | None = None,
+ max_sequence_length: int = 512,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ prepare_unconditional_embeds (`bool`):
+ whether to use prepare unconditional embeddings or not
+ negative_prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ max_sequence_length (`int`, defaults to `512`):
+ The maximum number of text tokens to be used for the generation process.
+ """
+ device = device or components._execution_device
+ if not isinstance(prompt, list):
+ prompt = [prompt]
+ batch_size = len(prompt)
+
+ prompt_embeds = get_t5_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ if prepare_unconditional_embeds:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = get_t5_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ return prompt_embeds, negative_prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ # Get inputs and intermediates
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ # Encode input prompt
+ (
+ block_state.prompt_embeds,
+ block_state.negative_prompt_embeds,
+ ) = self.encode_prompt(
+ components=components,
+ prompt=block_state.prompt,
+ device=block_state.device,
+ prepare_unconditional_embeds=components.requires_unconditional_embeds,
+ negative_prompt=block_state.negative_prompt,
+ max_sequence_length=block_state.max_sequence_length,
+ )
+
+ # Add outputs
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanImageResizeStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Image Resize step that resize the image to the target area (height * width) while maintaining the aspect ratio."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image", type_hint=PIL.Image.Image, required=True),
+ InputParam("height", type_hint=int, default=480),
+ InputParam("width", type_hint=int, default=832),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("resized_image", type_hint=PIL.Image.Image),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ max_area = block_state.height * block_state.width
+
+ image = block_state.image
+ aspect_ratio = image.height / image.width
+ mod_value = components.vae_scale_factor_spatial * components.patch_size_spatial
+ block_state.height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+ block_state.width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+ block_state.resized_image = image.resize((block_state.width, block_state.height))
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanImageCropResizeStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Image Resize step that resize the last_image to the same size of first frame image with center crop."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "resized_image", type_hint=PIL.Image.Image, required=True, description="The resized first frame image"
+ ),
+ InputParam("last_image", type_hint=PIL.Image.Image, required=True, description="The last frameimage"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("resized_last_image", type_hint=PIL.Image.Image),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ height = block_state.resized_image.height
+ width = block_state.resized_image.width
+ image = block_state.last_image
+
+ # Calculate resize ratio to match first frame dimensions
+ resize_ratio = max(width / image.width, height / image.height)
+
+ # Resize the image
+ width = round(image.width * resize_ratio)
+ height = round(image.height * resize_ratio)
+ size = [width, height]
+ resized_image = transforms.functional.center_crop(image, size)
+ block_state.resized_last_image = resized_image
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanImageEncoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Image Encoder step that generate image_embeds based on first frame image to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("image_processor", CLIPImageProcessor),
+ ComponentSpec("image_encoder", CLIPVisionModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ image = block_state.resized_image
+
+ image_embeds = encode_image(
+ image_processor=components.image_processor,
+ image_encoder=components.image_encoder,
+ image=image,
+ device=device,
+ )
+ block_state.image_embeds = image_embeds
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanFirstLastFrameImageEncoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Image Encoder step that generate image_embeds based on first and last frame images to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("image_processor", CLIPImageProcessor),
+ ComponentSpec("image_encoder", CLIPVisionModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
+ InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("image_embeds", type_hint=torch.Tensor, description="The image embeddings"),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ device = components._execution_device
+
+ first_frame_image = block_state.resized_image
+ last_frame_image = block_state.resized_last_image
+
+ image_embeds = encode_image(
+ image_processor=components.image_processor,
+ image_encoder=components.image_encoder,
+ image=[first_frame_image, last_frame_image],
+ device=device,
+ )
+ block_state.image_embeds = image_embeds
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Vae Image Encoder step that generate condition_latents based on first frame image to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("num_frames", type_hint=int, default=81),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "first_frame_latents",
+ type_hint=torch.Tensor,
+ description="video latent representation with the first frame image condition",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+ block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+ )
+ if block_state.num_frames is not None and (
+ block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_scale_factor_temporal != 0
+ ):
+ raise ValueError(
+ f"`num_frames` has to be greater than 0, and (num_frames - 1) must be divisible by {components.vae_scale_factor_temporal}, but got {block_state.num_frames}."
+ )
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ image = block_state.resized_image
+
+ device = components._execution_device
+ dtype = torch.float32
+ vae_dtype = components.vae.dtype
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ num_frames = block_state.num_frames or components.default_num_frames
+
+ image_tensor = components.video_processor.preprocess(image, height=height, width=width).to(
+ device=device, dtype=dtype
+ )
+
+ if image_tensor.dim() == 4:
+ image_tensor = image_tensor.unsqueeze(2)
+
+ video_tensor = torch.cat(
+ [
+ image_tensor,
+ image_tensor.new_zeros(image_tensor.shape[0], image_tensor.shape[1], num_frames - 1, height, width),
+ ],
+ dim=2,
+ ).to(device=device, dtype=dtype)
+
+ block_state.first_frame_latents = encode_vae_image(
+ video_tensor=video_tensor,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=vae_dtype,
+ latent_channels=components.num_channels_latents,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanPrepareFirstFrameLatentsStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "step that prepares the masked first frame latents and add it to the latent condition"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("first_frame_latents", type_hint=torch.Tensor | None),
+ InputParam("num_frames", required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("image_condition_latents", type_hint=torch.Tensor | None),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ batch_size, _, _, latent_height, latent_width = block_state.first_frame_latents.shape
+
+ mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width)
+ mask_lat_size[:, :, list(range(1, block_state.num_frames))] = 0
+
+ first_frame_mask = mask_lat_size[:, :, 0:1]
+ first_frame_mask = torch.repeat_interleave(
+ first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal
+ )
+ mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+ mask_lat_size = mask_lat_size.view(
+ batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width
+ )
+ mask_lat_size = mask_lat_size.transpose(1, 2)
+ mask_lat_size = mask_lat_size.to(block_state.first_frame_latents.device)
+ block_state.image_condition_latents = torch.concat([mask_lat_size, block_state.first_frame_latents], dim=1)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanFirstLastFrameVaeEncoderStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "Vae Image Encoder step that generate condition_latents based on first and last frame images to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKLWan),
+ ComponentSpec(
+ "video_processor",
+ VideoProcessor,
+ config=FrozenDict({"vae_scale_factor": 8}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("resized_image", type_hint=PIL.Image.Image, required=True),
+ InputParam("resized_last_image", type_hint=PIL.Image.Image, required=True),
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("num_frames", type_hint=int, default=81),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "first_last_frame_latents",
+ type_hint=torch.Tensor,
+ description="video latent representation with the first and last frame images condition",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+ block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+ )
+ if block_state.num_frames is not None and (
+ block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_scale_factor_temporal != 0
+ ):
+ raise ValueError(
+ f"`num_frames` has to be greater than 0, and (num_frames - 1) must be divisible by {components.vae_scale_factor_temporal}, but got {block_state.num_frames}."
+ )
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ first_frame_image = block_state.resized_image
+ last_frame_image = block_state.resized_last_image
+
+ device = components._execution_device
+ dtype = torch.float32
+ vae_dtype = components.vae.dtype
+
+ height = block_state.height or components.default_height
+ width = block_state.width or components.default_width
+ num_frames = block_state.num_frames or components.default_num_frames
+
+ first_image_tensor = components.video_processor.preprocess(first_frame_image, height=height, width=width).to(
+ device=device, dtype=dtype
+ )
+ first_image_tensor = first_image_tensor.unsqueeze(2)
+
+ last_image_tensor = components.video_processor.preprocess(last_frame_image, height=height, width=width).to(
+ device=device, dtype=dtype
+ )
+
+ last_image_tensor = last_image_tensor.unsqueeze(2)
+
+ video_tensor = torch.cat(
+ [
+ first_image_tensor,
+ first_image_tensor.new_zeros(
+ first_image_tensor.shape[0], first_image_tensor.shape[1], num_frames - 2, height, width
+ ),
+ last_image_tensor,
+ ],
+ dim=2,
+ ).to(device=device, dtype=dtype)
+
+ block_state.first_last_frame_latents = encode_vae_image(
+ video_tensor=video_tensor,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=vae_dtype,
+ latent_channels=components.num_channels_latents,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class WanPrepareFirstLastFrameLatentsStep(ModularPipelineBlocks):
+ model_name = "wan"
+
+ @property
+ def description(self) -> str:
+ return "step that prepares the masked latents with first and last frames and add it to the latent condition"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("first_last_frame_latents", type_hint=torch.Tensor | None),
+ InputParam("num_frames", type_hint=int, required=True),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam("image_condition_latents", type_hint=torch.Tensor | None),
+ ]
+
+ def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ batch_size, _, _, latent_height, latent_width = block_state.first_last_frame_latents.shape
+
+ mask_lat_size = torch.ones(batch_size, 1, block_state.num_frames, latent_height, latent_width)
+ mask_lat_size[:, :, list(range(1, block_state.num_frames - 1))] = 0
+
+ first_frame_mask = mask_lat_size[:, :, 0:1]
+ first_frame_mask = torch.repeat_interleave(
+ first_frame_mask, dim=2, repeats=components.vae_scale_factor_temporal
+ )
+ mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+ mask_lat_size = mask_lat_size.view(
+ batch_size, -1, components.vae_scale_factor_temporal, latent_height, latent_width
+ )
+ mask_lat_size = mask_lat_size.transpose(1, 2)
+ mask_lat_size = mask_lat_size.to(block_state.first_last_frame_latents.device)
+ block_state.image_condition_latents = torch.concat(
+ [mask_lat_size, block_state.first_last_frame_latents], dim=1
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py
new file mode 100644
index 0000000000000000000000000000000000000000..cf34a79f82e18ce18cf0ee08fac424c4b53b6c71
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan.py
@@ -0,0 +1,162 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ WanPrepareLatentsStep,
+ WanSetTimestepsStep,
+ WanTextInputStep,
+)
+from .decoders import WanVaeDecoderStep
+from .denoise import (
+ WanDenoiseStep,
+)
+from .encoders import (
+ WanTextEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. DENOISE
+# ====================
+
+
+# inputs(text) -> set_timesteps -> prepare_latents -> denoise
+# auto_docstring
+class WanCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "wan"
+ block_classes = [
+ WanTextInputStep,
+ WanSetTimestepsStep,
+ WanPrepareLatentsStep,
+ WanDenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# ====================
+# 2. BLOCKS (Wan2.1 text2video)
+# ====================
+
+
+# auto_docstring
+class WanBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline blocks for Wan2.1.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer
+ (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) vae (`AutoencoderKLWan`) video_processor
+ (`VideoProcessor`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "wan"
+ block_classes = [
+ WanTextEncoderStep,
+ WanCoreDenoiseStep,
+ WanVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "denoise", "decode"]
+
+ @property
+ def description(self):
+ return "Modular pipeline blocks for Wan2.1."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py
new file mode 100644
index 0000000000000000000000000000000000000000..239fa19187e0b4927774fd64dd5197222520cf24
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22.py
@@ -0,0 +1,176 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ WanPrepareLatentsStep,
+ WanSetTimestepsStep,
+ WanTextInputStep,
+)
+from .decoders import WanVaeDecoderStep
+from .denoise import (
+ Wan22DenoiseStep,
+)
+from .encoders import (
+ WanTextEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. DENOISE
+# ====================
+
+# inputs(text) -> set_timesteps -> prepare_latents -> denoise
+
+
+# auto_docstring
+class Wan22CoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+ guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "wan"
+ block_classes = [
+ WanTextInputStep,
+ WanSetTimestepsStep,
+ WanPrepareLatentsStep,
+ Wan22DenoiseStep,
+ ]
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# ====================
+# 2. BLOCKS (Wan2.2 text2video)
+# ====================
+
+
+# auto_docstring
+class Wan22Blocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for text-to-video using Wan2.2.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer
+ (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`)
+ transformer_2 (`WanTransformer3DModel`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ num_frames (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "wan"
+ block_classes = [
+ WanTextEncoderStep,
+ Wan22CoreDenoiseStep,
+ WanVaeDecoderStep,
+ ]
+ block_names = [
+ "text_encoder",
+ "denoise",
+ "decode",
+ ]
+
+ @property
+ def description(self):
+ return "Modular pipeline for text-to-video using Wan2.2."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py
new file mode 100644
index 0000000000000000000000000000000000000000..30d15cf256ae055cd0c0e7879fbecb848dac39a5
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan22_i2v.py
@@ -0,0 +1,236 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ WanAdditionalInputsStep,
+ WanPrepareLatentsStep,
+ WanSetTimestepsStep,
+ WanTextInputStep,
+)
+from .decoders import WanVaeDecoderStep
+from .denoise import (
+ Wan22Image2VideoDenoiseStep,
+)
+from .encoders import (
+ WanImageResizeStep,
+ WanPrepareFirstFrameLatentsStep,
+ WanTextEncoderStep,
+ WanVaeEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. VAE ENCODER
+# ====================
+
+
+# auto_docstring
+class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent
+ representation
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep]
+ block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"]
+
+ @property
+ def description(self):
+ return "Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent representation"
+
+
+# ====================
+# 2. DENOISE
+# ====================
+
+
+# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents)
+# auto_docstring
+class Wan22Image2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+ guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ num_frames (`None`, *optional*):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [
+ WanTextInputStep,
+ WanAdditionalInputsStep(image_latent_inputs=["image_condition_latents"]),
+ WanSetTimestepsStep,
+ WanPrepareLatentsStep,
+ Wan22Image2VideoDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "set_timesteps",
+ "prepare_latents",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# ====================
+# 3. BLOCKS (Wan2.2 Image2Video)
+# ====================
+
+
+# auto_docstring
+class Wan22Image2VideoBlocks(SequentialPipelineBlocks):
+ """
+ Modular pipeline for image-to-video using Wan2.2.
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler
+ (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanTransformer3DModel`)
+
+ Configs:
+ boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low
+ noise stages.
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [
+ WanTextEncoderStep,
+ WanImage2VideoVaeEncoderStep,
+ Wan22Image2VideoCoreDenoiseStep,
+ WanVaeDecoderStep,
+ ]
+ block_names = [
+ "text_encoder",
+ "vae_encoder",
+ "denoise",
+ "decode",
+ ]
+
+ @property
+ def description(self):
+ return "Modular pipeline for image-to-video using Wan2.2."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py
new file mode 100644
index 0000000000000000000000000000000000000000..74a31ce6cfcd7e34d555acba18b9c2797e763b42
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_blocks_wan_i2v.py
@@ -0,0 +1,481 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ WanAdditionalInputsStep,
+ WanPrepareLatentsStep,
+ WanSetTimestepsStep,
+ WanTextInputStep,
+)
+from .decoders import WanVaeDecoderStep
+from .denoise import (
+ WanImage2VideoDenoiseStep,
+)
+from .encoders import (
+ WanFirstLastFrameImageEncoderStep,
+ WanFirstLastFrameVaeEncoderStep,
+ WanImageCropResizeStep,
+ WanImageEncoderStep,
+ WanImageResizeStep,
+ WanPrepareFirstFrameLatentsStep,
+ WanPrepareFirstLastFrameLatentsStep,
+ WanTextEncoderStep,
+ WanVaeEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+# ====================
+# 1. IMAGE ENCODER
+# ====================
+
+
+# wan2.1 I2V (first frame only)
+# auto_docstring
+class WanImage2VideoImageEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [WanImageResizeStep, WanImageEncoderStep]
+ block_names = ["image_resize", "image_encoder"]
+
+ @property
+ def description(self):
+ return "Image2Video Image Encoder step that resize the image and encode the image to generate the image embeddings"
+
+
+# wan2.1 FLF2V (first and last frame)
+# auto_docstring
+class WanFLF2VImageEncoderStep(SequentialPipelineBlocks):
+ """
+ FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image
+ embeddings
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`):
+ The last frameimage
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [WanImageResizeStep, WanImageCropResizeStep, WanFirstLastFrameImageEncoderStep]
+ block_names = ["image_resize", "last_image_resize", "image_encoder"]
+
+ @property
+ def description(self):
+ return "FLF2V Image Encoder step that resize and encode and encode the first and last frame images to generate the image embeddings"
+
+
+# wan2.1 Auto Image Encoder
+# auto_docstring
+class WanAutoImageEncoderStep(AutoPipelineBlocks):
+ """
+ Image Encoder step that encode the image to generate the image embeddingsThis is an auto pipeline block that works
+ for image2video tasks. - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided. -
+ `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is
+ not provided, step will be skipped.
+
+ Components:
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ The image embeddings
+ """
+
+ block_classes = [WanFLF2VImageEncoderStep, WanImage2VideoImageEncoderStep]
+ block_names = ["flf2v_image_encoder", "image2video_image_encoder"]
+ block_trigger_inputs = ["last_image", "image"]
+ model_name = "wan-i2v"
+
+ @property
+ def description(self):
+ return (
+ "Image Encoder step that encode the image to generate the image embeddings"
+ + "This is an auto pipeline block that works for image2video tasks."
+ + " - `WanFLF2VImageEncoderStep` (flf2v) is used when `last_image` is provided."
+ + " - `WanImage2VideoImageEncoderStep` (image2video) is used when `image` is provided."
+ + " - if `last_image` or `image` is not provided, step will be skipped."
+ )
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# wan2.1 I2V (first frame only)
+# auto_docstring
+class WanImage2VideoVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent
+ representation
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [WanImageResizeStep, WanVaeEncoderStep, WanPrepareFirstFrameLatentsStep]
+ block_names = ["image_resize", "vae_encoder", "prepare_first_frame_latents"]
+
+ @property
+ def description(self):
+ return "Image2Video Vae Image Encoder step that resize the image and encode the first frame image to its latent representation"
+
+
+# wan2.1 FLF2V (first and last frame)
+# auto_docstring
+class WanFLF2VVaeEncoderStep(SequentialPipelineBlocks):
+ """
+ FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the
+ latent conditions
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ first_last_frame_latents (`Tensor`):
+ video latent representation with the first and last frame images condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [
+ WanImageResizeStep,
+ WanImageCropResizeStep,
+ WanFirstLastFrameVaeEncoderStep,
+ WanPrepareFirstLastFrameLatentsStep,
+ ]
+ block_names = ["image_resize", "last_image_resize", "vae_encoder", "prepare_first_last_frame_latents"]
+
+ @property
+ def description(self):
+ return "FLF2V Vae Image Encoder step that resize and encode and encode the first and last frame images to generate the latent conditions"
+
+
+# wan2.1 Auto Vae Encoder
+# auto_docstring
+class WanAutoVaeEncoderStep(AutoPipelineBlocks):
+ """
+ Vae Image Encoder step that encode the image to generate the image latentsThis is an auto pipeline block that works
+ for image2video tasks. - `WanFLF2VVaeEncoderStep` (flf2v) is used when `last_image` is provided. -
+ `WanImage2VideoVaeEncoderStep` (image2video) is used when `image` is provided. - if `last_image` or `image` is not
+ provided, step will be skipped.
+
+ Components:
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ resized_image (`Image`):
+ TODO: Add description.
+ resized_last_image (`Image`):
+ TODO: Add description.
+ first_last_frame_latents (`Tensor`):
+ video latent representation with the first and last frame images condition
+ image_condition_latents (`Tensor | NoneType`):
+ TODO: Add description.
+ first_frame_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [WanFLF2VVaeEncoderStep, WanImage2VideoVaeEncoderStep]
+ block_names = ["flf2v_vae_encoder", "image2video_vae_encoder"]
+ block_trigger_inputs = ["last_image", "image"]
+
+ @property
+ def description(self):
+ return (
+ "Vae Image Encoder step that encode the image to generate the image latents"
+ + "This is an auto pipeline block that works for image2video tasks."
+ + " - `WanFLF2VVaeEncoderStep` (flf2v) is used when `last_image` is provided."
+ + " - `WanImage2VideoVaeEncoderStep` (image2video) is used when `image` is provided."
+ + " - if `last_image` or `image` is not provided, step will be skipped."
+ )
+
+
+# ====================
+# 3. DENOISE (inputs -> set_timesteps -> prepare_latents -> denoise)
+# ====================
+
+
+# wan2.1 I2V core denoise (support both I2V and FLF2V)
+# inputs (text + image_condition_latents) -> set_timesteps -> prepare_latents -> denoise (latents)
+# auto_docstring
+class WanImage2VideoCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`Tensor`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ num_frames (`None`, *optional*):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ TODO: Add description.
+
+ Outputs:
+ batch_size (`int`):
+ Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt
+ dtype (`dtype`):
+ Data type of model tensor inputs (determined by `transformer.dtype`)
+ latents (`Tensor`):
+ The initial latents to use for the denoising process
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [
+ WanTextInputStep,
+ WanAdditionalInputsStep(image_latent_inputs=["image_condition_latents"]),
+ WanSetTimestepsStep,
+ WanPrepareLatentsStep,
+ WanImage2VideoDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "set_timesteps",
+ "prepare_latents",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
+
+
+# ====================
+# 4. BLOCKS (Wan2.1 Image2Video)
+# ====================
+
+
+# wan2.1 Image2Video Auto Blocks
+# auto_docstring
+class WanImage2VideoAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for image-to-video using Wan.
+
+ Supported workflows:
+ - `image2video`: requires `image`, `prompt`
+ - `flf2v`: requires `last_image`, `image`, `prompt`
+
+ Components:
+ text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`)
+ image_processor (`CLIPImageProcessor`) image_encoder (`CLIPVisionModel`) vae (`AutoencoderKLWan`)
+ video_processor (`VideoProcessor`) transformer (`WanTransformer3DModel`) scheduler
+ (`UniPCMultistepScheduler`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`int`, *optional*, defaults to 480):
+ TODO: Add description.
+ width (`int`, *optional*, defaults to 832):
+ TODO: Add description.
+ last_image (`Image`, *optional*):
+ The last frameimage
+ num_frames (`int`, *optional*, defaults to 81):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_videos_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_condition_latents (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 50):
+ TODO: Add description.
+ timesteps (`None`, *optional*):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ attention_kwargs (`None`, *optional*):
+ TODO: Add description.
+ image_embeds (`Tensor`):
+ TODO: Add description.
+ output_type (`str`, *optional*, defaults to np):
+ The output type of the decoded videos
+
+ Outputs:
+ videos (`list`):
+ The generated videos.
+ """
+
+ model_name = "wan-i2v"
+ block_classes = [
+ WanTextEncoderStep,
+ WanAutoImageEncoderStep,
+ WanAutoVaeEncoderStep,
+ WanImage2VideoCoreDenoiseStep,
+ WanVaeDecoderStep,
+ ]
+ block_names = [
+ "text_encoder",
+ "image_encoder",
+ "vae_encoder",
+ "denoise",
+ "decode",
+ ]
+
+ _workflow_map = {
+ "image2video": {"image": True, "prompt": True},
+ "flf2v": {"last_image": True, "image": True, "prompt": True},
+ }
+
+ @property
+ def description(self):
+ return "Auto Modular pipeline for image-to-video using Wan."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("videos")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..09b6f2f83b5a3cc140f92b71d999043176b78040
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/wan/modular_pipeline.py
@@ -0,0 +1,141 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...loaders import WanLoraLoaderMixin
+from ...pipelines.pipeline_utils import StableDiffusionMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class WanModularPipeline(
+ ModularPipeline,
+ StableDiffusionMixin,
+ WanLoraLoaderMixin,
+):
+ """
+ A ModularPipeline for Wan2.1 text2video.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "WanBlocks"
+
+ @property
+ def default_height(self):
+ return self.default_sample_height * self.vae_scale_factor_spatial
+
+ @property
+ def default_width(self):
+ return self.default_sample_width * self.vae_scale_factor_spatial
+
+ @property
+ def default_num_frames(self):
+ return (self.default_sample_num_frames - 1) * self.vae_scale_factor_temporal + 1
+
+ @property
+ def default_sample_height(self):
+ return 60
+
+ @property
+ def default_sample_width(self):
+ return 104
+
+ @property
+ def default_sample_num_frames(self):
+ return 21
+
+ @property
+ def patch_size_spatial(self):
+ patch_size_spatial = 2
+ if hasattr(self, "transformer") and self.transformer is not None:
+ patch_size_spatial = self.transformer.config.patch_size[1]
+ return patch_size_spatial
+
+ @property
+ def vae_scale_factor_spatial(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def vae_scale_factor_temporal(self):
+ vae_scale_factor = 4
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** sum(self.vae.temperal_downsample)
+ return vae_scale_factor
+
+ @property
+ def num_channels_transformer(self):
+ num_channels_transformer = 16
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_transformer = self.transformer.config.in_channels
+ return num_channels_transformer
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if hasattr(self, "vae") and self.vae is not None:
+ num_channels_latents = self.vae.config.z_dim
+ return num_channels_latents
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
+
+ @property
+ def num_train_timesteps(self):
+ num_train_timesteps = 1000
+ if hasattr(self, "scheduler") and self.scheduler is not None:
+ num_train_timesteps = self.scheduler.config.num_train_timesteps
+ return num_train_timesteps
+
+
+class WanImage2VideoModularPipeline(WanModularPipeline):
+ """
+ A ModularPipeline for Wan2.1 image2video (both I2V and FLF2V).
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "WanImage2VideoAutoBlocks"
+
+
+class Wan22ModularPipeline(WanModularPipeline):
+ """
+ A ModularPipeline for Wan2.2 text2video.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "Wan22Blocks"
+
+
+class Wan22Image2VideoModularPipeline(Wan22ModularPipeline):
+ """
+ A ModularPipeline for Wan2.2 image2video.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "Wan22Image2VideoBlocks"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..085b27a7d9e338880b7627fb07db8d2de55f0c6b
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/__init__.py
@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+ DIFFUSERS_SLOW_IMPORT,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_torch_available,
+ is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
+
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+ _import_structure["modular_blocks_z_image"] = ["ZImageAutoBlocks"]
+ _import_structure["modular_pipeline"] = ["ZImageModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+ try:
+ if not (is_transformers_available() and is_torch_available()):
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
+ else:
+ from .modular_blocks_z_image import ZImageAutoBlocks
+ from .modular_pipeline import ZImageModularPipeline
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ module_spec=__spec__,
+ )
+
+ for name, value in _dummy_objects.items():
+ setattr(sys.modules[__name__], name, value)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/before_denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/before_denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..d37ffee9cb48e5ced908f04f371b1cbdf2c36d10
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/before_denoise.py
@@ -0,0 +1,627 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import torch
+
+from ...models import ZImageTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ...utils.torch_utils import randn_tensor
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ZImageModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# TODO(yiyi, aryan): We need another step before text encoder to set the `num_inference_steps` attribute for guider so that
+# things like when to do guidance and how many conditions to be prepared can be determined. Currently, this is done by
+# always assuming you want to do guidance in the Guiders. So, negative embeddings are prepared regardless of what the
+# configuration of guider is.
+
+
+def repeat_tensor_to_batch_size(
+ input_name: str,
+ input_tensor: torch.Tensor,
+ batch_size: int,
+ num_images_per_prompt: int = 1,
+) -> torch.Tensor:
+ """Repeat tensor elements to match the final batch size.
+
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
+ by repeating each element along dimension 0.
+
+ The input tensor must have batch size 1 or batch_size. The function will:
+ - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
+ - If batch size equals batch_size: repeat each element num_images_per_prompt times
+
+ Args:
+ input_name (str): Name of the input tensor (used for error messages)
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+ batch_size (int): The base batch size (number of prompts)
+ num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
+
+ Returns:
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
+
+ Raises:
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+ Examples:
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+ batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+ [4, 3]
+
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+ tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+ - shape: [4, 3]
+ """
+ # make sure input is a tensor
+ if not isinstance(input_tensor, torch.Tensor):
+ raise ValueError(f"`{input_name}` must be a tensor")
+
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+ if input_tensor.shape[0] == 1:
+ repeat_by = batch_size * num_images_per_prompt
+ elif input_tensor.shape[0] == batch_size:
+ repeat_by = num_images_per_prompt
+ else:
+ raise ValueError(
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+ )
+
+ # expand the tensor to match the batch_size * num_images_per_prompt
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+ return input_tensor
+
+
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor_spatial: int) -> tuple[int, int]:
+ """Calculate image dimensions from latent tensor dimensions.
+
+ This function converts latent spatial dimensions to image spatial dimensions by multiplying the latent height/width
+ by the VAE scale factor.
+
+ Args:
+ latents (torch.Tensor): The latent tensor. Must have 4 dimensions.
+ Expected shapes: [batch, channels, height, width]
+ vae_scale_factor (int): The scale factor used by the VAE to compress image spatial dimension.
+ By default, it is 16
+ Returns:
+ tuple[int, int]: The calculated image dimensions as (height, width)
+ """
+ latent_height, latent_width = latents.shape[2:]
+ height = latent_height * vae_scale_factor_spatial // 2
+ width = latent_width * vae_scale_factor_spatial // 2
+
+ return height, width
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+ return torch.linspace(1.0, 1 / num_inference_steps, num_inference_steps).tolist()
+
+
+class ZImageTextInputStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Input processing step that:\n"
+ " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+ " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_images_per_prompt`\n\n"
+ "All input tensors are expected to have either batch_size=1 or match the batch_size\n"
+ "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n"
+ "have a final batch_size of batch_size * num_images_per_prompt."
+ )
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("transformer", ZImageTransformer2DModel),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("num_images_per_prompt", default=1),
+ InputParam(
+ "prompt_embeds",
+ required=True,
+ type_hint=list[torch.Tensor],
+ description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+ ),
+ InputParam(
+ "negative_prompt_embeds",
+ type_hint=list[torch.Tensor],
+ description="Pre-generated negative text embeddings. Can be generated from text_encoder step.",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "batch_size",
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+ ),
+ OutputParam(
+ "dtype",
+ type_hint=torch.dtype,
+ description="Data type of model tensor inputs (determined by `transformer.dtype`)",
+ ),
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.prompt_embeds is not None and block_state.negative_prompt_embeds is not None:
+ if not isinstance(block_state.prompt_embeds, list):
+ raise ValueError(
+ f"`prompt_embeds` must be a list when passed directly, but got {type(block_state.prompt_embeds)}."
+ )
+ if not isinstance(block_state.negative_prompt_embeds, list):
+ raise ValueError(
+ f"`negative_prompt_embeds` must be a list when passed directly, but got {type(block_state.negative_prompt_embeds)}."
+ )
+ if len(block_state.prompt_embeds) != len(block_state.negative_prompt_embeds):
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same length when passed directly, but"
+ f" got: `prompt_embeds` {len(block_state.prompt_embeds)} != `negative_prompt_embeds`"
+ f" {len(block_state.negative_prompt_embeds)}."
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ block_state.batch_size = len(block_state.prompt_embeds)
+ block_state.dtype = block_state.prompt_embeds[0].dtype
+
+ if block_state.num_images_per_prompt > 1:
+ prompt_embeds = [pe for pe in block_state.prompt_embeds for _ in range(block_state.num_images_per_prompt)]
+ block_state.prompt_embeds = prompt_embeds
+
+ if block_state.negative_prompt_embeds is not None:
+ negative_prompt_embeds = [
+ npe for npe in block_state.negative_prompt_embeds for _ in range(block_state.num_images_per_prompt)
+ ]
+ block_state.negative_prompt_embeds = negative_prompt_embeds
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class ZImageAdditionalInputsStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ def __init__(
+ self,
+ image_latent_inputs: list[str] = ["image_latents"],
+ additional_batch_inputs: list[str] = [],
+ ):
+ """Initialize a configurable step that standardizes the inputs for the denoising step. It:\n"
+
+ This step handles multiple common tasks to prepare inputs for the denoising step:
+ 1. For encoded image latents, use it update height/width if None, and expands batch size
+ 2. For additional_batch_inputs: Only expands batch dimensions to match final batch size
+
+ This is a dynamic block that allows you to configure which inputs to process.
+
+ Args:
+ image_latent_inputs (list[str], optional): Names of image latent tensors to process.
+ In additional to adjust batch size of these inputs, they will be used to determine height/width. Can be
+ a single string or list of strings. Defaults to ["image_latents"].
+ additional_batch_inputs (list[str], optional):
+ Names of additional conditional input tensors to expand batch size. These tensors will only have their
+ batch dimensions adjusted to match the final batch size. Can be a single string or list of strings.
+ Defaults to [].
+
+ Examples:
+ # Configure to process image_latents (default behavior) ZImageAdditionalInputsStep()
+
+ # Configure to process multiple image latent inputs
+ ZImageAdditionalInputsStep(image_latent_inputs=["image_latents", "control_image_latents"])
+
+ # Configure to process image latents and additional batch inputs ZImageAdditionalInputsStep(
+ image_latent_inputs=["image_latents"], additional_batch_inputs=["image_embeds"]
+ )
+ """
+ if not isinstance(image_latent_inputs, list):
+ image_latent_inputs = [image_latent_inputs]
+ if not isinstance(additional_batch_inputs, list):
+ additional_batch_inputs = [additional_batch_inputs]
+
+ self._image_latent_inputs = image_latent_inputs
+ self._additional_batch_inputs = additional_batch_inputs
+ super().__init__()
+
+ @property
+ def description(self) -> str:
+ # Functionality section
+ summary_section = (
+ "Input processing step that:\n"
+ " 1. For image latent inputs: Updates height/width if None, and expands batch size\n"
+ " 2. For additional batch inputs: Expands batch dimensions to match final batch size"
+ )
+
+ # Inputs info
+ inputs_info = ""
+ if self._image_latent_inputs or self._additional_batch_inputs:
+ inputs_info = "\n\nConfigured inputs:"
+ if self._image_latent_inputs:
+ inputs_info += f"\n - Image latent inputs: {self._image_latent_inputs}"
+ if self._additional_batch_inputs:
+ inputs_info += f"\n - Additional batch inputs: {self._additional_batch_inputs}"
+
+ # Placement guidance
+ placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+ return summary_section + inputs_info + placement_section
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ inputs = [
+ InputParam(name="num_images_per_prompt", default=1),
+ InputParam(name="batch_size", required=True),
+ InputParam(name="height"),
+ InputParam(name="width"),
+ ]
+
+ # Add image latent inputs
+ for image_latent_input_name in self._image_latent_inputs:
+ inputs.append(InputParam(name=image_latent_input_name))
+
+ # Add additional batch inputs
+ for input_name in self._additional_batch_inputs:
+ inputs.append(InputParam(name=input_name))
+
+ return inputs
+
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ # Process image latent inputs (height/width calculation, patchify, and batch expansion)
+ for image_latent_input_name in self._image_latent_inputs:
+ image_latent_tensor = getattr(block_state, image_latent_input_name)
+ if image_latent_tensor is None:
+ continue
+
+ # 1. Calculate num_frames, height/width from latents
+ height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor_spatial)
+ block_state.height = block_state.height or height
+ block_state.width = block_state.width or width
+
+ # Process additional batch inputs (only batch expansion)
+ for input_name in self._additional_batch_inputs:
+ input_tensor = getattr(block_state, input_name)
+ if input_tensor is None:
+ continue
+
+ # Only expand batch size
+ input_tensor = repeat_tensor_to_batch_size(
+ input_name=input_name,
+ input_tensor=input_tensor,
+ num_images_per_prompt=block_state.num_images_per_prompt,
+ batch_size=block_state.batch_size,
+ )
+
+ setattr(block_state, input_name, input_tensor)
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ZImagePrepareLatentsStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return "Prepare latents step that prepares the latents for the text-to-video generation process"
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("height", type_hint=int),
+ InputParam("width", type_hint=int),
+ InputParam("latents", type_hint=torch.Tensor | None),
+ InputParam("num_images_per_prompt", type_hint=int, default=1),
+ InputParam("generator"),
+ InputParam(
+ "batch_size",
+ required=True,
+ type_hint=int,
+ description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.",
+ ),
+ InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process"
+ )
+ ]
+
+ def check_inputs(self, components, block_state):
+ if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+ block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+ )
+
+ @staticmethod
+ # Copied from diffusers.pipelines.z_image.pipeline_z_image.ZImagePipeline.prepare_latents with self->comp
+ def prepare_latents(
+ comp,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ height = 2 * (int(height) // (comp.vae_scale_factor * 2))
+ width = 2 * (int(width) // (comp.vae_scale_factor * 2))
+
+ shape = (batch_size, num_channels_latents, height, width)
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ if latents.shape != shape:
+ raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+ latents = latents.to(device)
+ return latents
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ device = components._execution_device
+ dtype = torch.float32
+
+ block_state.height = block_state.height or components.default_height
+ block_state.width = block_state.width or components.default_width
+
+ block_state.latents = self.prepare_latents(
+ components,
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
+ num_channels_latents=components.num_channels_latents,
+ height=block_state.height,
+ width=block_state.width,
+ dtype=dtype,
+ device=device,
+ generator=block_state.generator,
+ latents=block_state.latents,
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class ZImageSetTimestepsStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference. Need to run after prepare latents step."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latents", required=True),
+ InputParam("num_inference_steps", default=8),
+ InputParam("sigmas"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process"
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ device = components._execution_device
+
+ latent_height, latent_width = block_state.latents.shape[2], block_state.latents.shape[3]
+ image_seq_len = (latent_height // 2) * (latent_width // 2) # sequence length after patchify
+
+ mu = calculate_shift(
+ image_seq_len,
+ base_seq_len=components.scheduler.config.get("base_image_seq_len", 256),
+ max_seq_len=components.scheduler.config.get("max_image_seq_len", 4096),
+ base_shift=components.scheduler.config.get("base_shift", 0.5),
+ max_shift=components.scheduler.config.get("max_shift", 1.15),
+ )
+ sigmas = block_state.sigmas
+ if sigmas is None:
+ sigmas = get_default_z_image_sigmas(block_state.num_inference_steps)
+
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+ components.scheduler,
+ block_state.num_inference_steps,
+ device,
+ sigmas=sigmas,
+ mu=mu,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ZImageSetTimestepsWithStrengthStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that sets the scheduler's timesteps for inference with strength. Need to run after set timesteps step."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("timesteps", required=True),
+ InputParam("num_inference_steps", required=True),
+ InputParam("strength", default=0.6),
+ ]
+
+ def check_inputs(self, components, block_state):
+ if block_state.strength < 0.0 or block_state.strength > 1.0:
+ raise ValueError(f"Strength must be between 0.0 and 1.0, but got {block_state.strength}")
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ init_timestep = min(block_state.num_inference_steps * block_state.strength, block_state.num_inference_steps)
+
+ t_start = int(max(block_state.num_inference_steps - init_timestep, 0))
+ timesteps = components.scheduler.timesteps[t_start * components.scheduler.order :]
+ if hasattr(components.scheduler, "set_begin_index"):
+ components.scheduler.set_begin_index(t_start * components.scheduler.order)
+
+ block_state.timesteps = timesteps
+ block_state.num_inference_steps = block_state.num_inference_steps - t_start
+
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ZImagePrepareLatentswithImageStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return "step that prepares the latents with image condition, need to run after set timesteps and prepare latents step."
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("latents", required=True),
+ InputParam("image_latents", required=True),
+ InputParam("timesteps", required=True),
+ ]
+
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
+ block_state.latents = components.scheduler.scale_noise(
+ block_state.image_latents, latent_timestep, block_state.latents
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/decoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/decoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f0ee30c8a7e31f5dcee9302c23b45fc2995f832
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/decoders.py
@@ -0,0 +1,91 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import PIL
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKL
+from ...utils import logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ZImageVaeDecoderStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8 * 2}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def description(self) -> str:
+ return "Step that decodes the denoised latents into images"
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ ),
+ InputParam(
+ name="output_type",
+ default="pil",
+ type_hint=str,
+ description="The type of the output images, can be 'pil', 'np', 'pt'",
+ ),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[str]:
+ return [
+ OutputParam(
+ "images",
+ type_hint=list[PIL.Image.Image, list[torch.Tensor], list[np.ndarray]],
+ description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array",
+ )
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ vae_dtype = components.vae.dtype
+
+ latents = block_state.latents.to(vae_dtype)
+ latents = latents / components.vae.config.scaling_factor + components.vae.config.shift_factor
+
+ block_state.images = components.vae.decode(latents, return_dict=False)[0]
+ block_state.images = components.image_processor.postprocess(
+ block_state.images, output_type=block_state.output_type
+ )
+
+ self.set_block_state(state, block_state)
+
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/denoise.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/denoise.py
new file mode 100644
index 0000000000000000000000000000000000000000..d875e2d9db66f50f35b102feaaa83589d701e06a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/denoise.py
@@ -0,0 +1,314 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...models import ZImageTransformer2DModel
+from ...schedulers import FlowMatchEulerDiscreteScheduler
+from ...utils import logging
+from ..modular_pipeline import (
+ BlockState,
+ LoopSequentialPipelineBlocks,
+ ModularPipelineBlocks,
+ PipelineState,
+)
+from ..modular_pipeline_utils import ComponentSpec, InputParam
+from .modular_pipeline import ZImageModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ZImageLoopBeforeDenoiser(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that prepares the latent input for the denoiser. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `ZImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "latents",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+ ),
+ InputParam(
+ "dtype",
+ required=True,
+ type_hint=torch.dtype,
+ description="The dtype of the model inputs. Can be generated in input step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ latents = block_state.latents.unsqueeze(2).to(
+ block_state.dtype
+ ) # [batch_size, num_channels, 1, height, width]
+ block_state.latent_model_input = list(latents.unbind(dim=0)) # list of [num_channels, 1, height, width]
+
+ timestep = t.expand(latents.shape[0]).to(block_state.dtype)
+ timestep = (1000 - timestep) / 1000
+ block_state.timestep = timestep
+ return components, block_state
+
+
+class ZImageLoopDenoiser(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ def __init__(
+ self,
+ guider_input_fields: dict[str, Any] = {"cap_feats": ("prompt_embeds", "negative_prompt_embeds")},
+ ):
+ """Initialize a denoiser block that calls the denoiser model. This block is used in Z-Image.
+
+ Args:
+ guider_input_fields: A dictionary that maps each argument expected by the denoiser model
+ (for example, "encoder_hidden_states") to data stored on 'block_state'. The value can be either:
+
+ - A tuple of strings. For instance, {"encoder_hidden_states": ("prompt_embeds",
+ "negative_prompt_embeds")} tells the guider to read `block_state.prompt_embeds` and
+ `block_state.negative_prompt_embeds` and pass them as the conditional and unconditional batches of
+ 'encoder_hidden_states'.
+ - A string. For example, {"encoder_hidden_image": "image_embeds"} makes the guider forward
+ `block_state.image_embeds` for both conditional and unconditional batches.
+ """
+ if not isinstance(guider_input_fields, dict):
+ raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+ self._guider_input_fields = guider_input_fields
+ super().__init__()
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0, "enabled": False}),
+ default_creation_method="from_config",
+ ),
+ ComponentSpec("transformer", ZImageTransformer2DModel),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Step within the denoising loop that denoise the latents with guidance. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `ZImageDenoiseLoopWrapper`)"
+ )
+
+ @property
+ def inputs(self) -> list[tuple[str, Any]]:
+ inputs = [
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ kwargs_type="denoiser_input_fields",
+ description="The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.",
+ ),
+ ]
+ guider_input_names = []
+ uncond_guider_input_names = []
+ for value in self._guider_input_fields.values():
+ if isinstance(value, tuple):
+ guider_input_names.append(value[0])
+ uncond_guider_input_names.append(value[1])
+ else:
+ guider_input_names.append(value)
+
+ for name in guider_input_names:
+ inputs.append(InputParam(name=name, required=True))
+ for name in uncond_guider_input_names:
+ inputs.append(InputParam(name=name))
+ return inputs
+
+ @torch.no_grad()
+ def __call__(
+ self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+ ) -> PipelineState:
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+ # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+ # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+ # you will get a guider_state with two batches:
+ # guider_state = [
+ # {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"}, # conditional batch
+ # {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"}, # unconditional batch
+ # ]
+ # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+ # run the denoiser for each guidance batch
+ for guider_state_batch in guider_state:
+ components.guider.prepare_models(components.transformer)
+ cond_kwargs = guider_state_batch.as_dict()
+
+ def _convert_dtype(v, dtype):
+ if isinstance(v, torch.Tensor):
+ return v.to(dtype)
+ elif isinstance(v, list):
+ return [_convert_dtype(t, dtype) for t in v]
+ return v
+
+ cond_kwargs = {
+ k: _convert_dtype(v, block_state.dtype)
+ for k, v in cond_kwargs.items()
+ if k in self._guider_input_fields.keys()
+ }
+
+ # Predict the noise residual
+ # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+ model_out_list = components.transformer(
+ x=block_state.latent_model_input,
+ t=block_state.timestep,
+ return_dict=False,
+ **cond_kwargs,
+ )[0]
+ noise_pred = torch.stack(model_out_list, dim=0).squeeze(2)
+ guider_state_batch.noise_pred = -noise_pred
+ components.guider.cleanup_models(components.transformer)
+
+ # Perform guidance
+ block_state.noise_pred = components.guider(guider_state)[0]
+
+ return components, block_state
+
+
+class ZImageLoopAfterDenoiser(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def description(self) -> str:
+ return (
+ "step within the denoising loop that update the latents. "
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+ "object (e.g. `ZImageDenoiseLoopWrapper`)"
+ )
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+ # Perform scheduler step using the predicted output
+ latents_dtype = block_state.latents.dtype
+ block_state.latents = components.scheduler.step(
+ block_state.noise_pred.float(),
+ t,
+ block_state.latents.float(),
+ return_dict=False,
+ )[0]
+
+ if block_state.latents.dtype != latents_dtype:
+ block_state.latents = block_state.latents.to(latents_dtype)
+
+ return components, block_state
+
+
+class ZImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return (
+ "Pipeline block that iteratively denoise the latents over `timesteps`. "
+ "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+ )
+
+ @property
+ def loop_expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+ ]
+
+ @property
+ def loop_inputs(self) -> list[InputParam]:
+ return [
+ InputParam(
+ "timesteps",
+ required=True,
+ type_hint=torch.Tensor,
+ description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ InputParam(
+ "num_inference_steps",
+ required=True,
+ type_hint=int,
+ description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+ ),
+ ]
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+
+ block_state.num_warmup_steps = max(
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+ )
+
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+ for i, t in enumerate(block_state.timesteps):
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
+ if i == len(block_state.timesteps) - 1 or (
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+ ):
+ progress_bar.update()
+
+ self.set_block_state(state, block_state)
+
+ return components, state
+
+
+class ZImageDenoiseStep(ZImageDenoiseLoopWrapper):
+ block_classes = [
+ ZImageLoopBeforeDenoiser,
+ ZImageLoopDenoiser(
+ guider_input_fields={
+ "cap_feats": ("prompt_embeds", "negative_prompt_embeds"),
+ }
+ ),
+ ZImageLoopAfterDenoiser,
+ ]
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. \n"
+ "Its loop logic is defined in `ZImageDenoiseLoopWrapper.__call__` method \n"
+ "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+ " - `ZImageLoopBeforeDenoiser`\n"
+ " - `ZImageLoopDenoiser`\n"
+ " - `ZImageLoopAfterDenoiser`\n"
+ "This block supports text-to-image and image-to-image tasks for Z-Image."
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/encoders.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/encoders.py
new file mode 100644
index 0000000000000000000000000000000000000000..78829b504f349d330a024e158ea47dd2acc4fb78
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/encoders.py
@@ -0,0 +1,343 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import PIL
+import torch
+from transformers import Qwen2Tokenizer, Qwen3Model
+
+from ...configuration_utils import FrozenDict
+from ...guiders import ClassifierFreeGuidance
+from ...image_processor import VaeImageProcessor
+from ...models import AutoencoderKL
+from ...utils import is_ftfy_available, logging
+from ..modular_pipeline import ModularPipelineBlocks, PipelineState
+from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import ZImageModularPipeline
+
+
+if is_ftfy_available():
+ pass
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def get_qwen_prompt_embeds(
+ text_encoder: Qwen3Model,
+ tokenizer: Qwen2Tokenizer,
+ prompt: str | list[str],
+ device: torch.device,
+ max_sequence_length: int = 512,
+) -> list[torch.Tensor]:
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ for i, prompt_item in enumerate(prompt):
+ messages = [
+ {"role": "user", "content": prompt_item},
+ ]
+ prompt_item = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=True,
+ )
+ prompt[i] = prompt_item
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids.to(device)
+ prompt_masks = text_inputs.attention_mask.to(device).bool()
+
+ prompt_embeds = text_encoder(
+ input_ids=text_input_ids,
+ attention_mask=prompt_masks,
+ output_hidden_states=True,
+ ).hidden_states[-2]
+
+ prompt_embeds_list = []
+
+ for i in range(len(prompt_embeds)):
+ prompt_embeds_list.append(prompt_embeds[i][prompt_masks[i]])
+
+ return prompt_embeds_list
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def encode_vae_image(
+ image_tensor: torch.Tensor,
+ vae: AutoencoderKL,
+ generator: torch.Generator,
+ device: torch.device,
+ dtype: torch.dtype,
+ latent_channels: int = 16,
+):
+ if not isinstance(image_tensor, torch.Tensor):
+ raise ValueError(f"Expected image_tensor to be a tensor, got {type(image_tensor)}.")
+
+ if isinstance(generator, list) and len(generator) != image_tensor.shape[0]:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but it is not same as number of images {image_tensor.shape[0]}."
+ )
+
+ image_tensor = image_tensor.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list):
+ image_latents = [
+ retrieve_latents(vae.encode(image_tensor[i : i + 1]), generator=generator[i])
+ for i in range(image_tensor.shape[0])
+ ]
+ image_latents = torch.cat(image_latents, dim=0)
+ else:
+ image_latents = retrieve_latents(vae.encode(image_tensor), generator=generator)
+
+ image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+ return image_latents
+
+
+class ZImageTextEncoderStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return "Text Encoder step that generate text_embeddings to guide the video generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("text_encoder", Qwen3Model),
+ ComponentSpec("tokenizer", Qwen2Tokenizer),
+ ComponentSpec(
+ "guider",
+ ClassifierFreeGuidance,
+ config=FrozenDict({"guidance_scale": 5.0, "enabled": False}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("prompt"),
+ InputParam("negative_prompt"),
+ InputParam("max_sequence_length", default=512),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "prompt_embeds",
+ type_hint=list[torch.Tensor],
+ kwargs_type="denoiser_input_fields",
+ description="text embeddings used to guide the image generation",
+ ),
+ OutputParam(
+ "negative_prompt_embeds",
+ type_hint=list[torch.Tensor],
+ kwargs_type="denoiser_input_fields",
+ description="negative text embeddings used to guide the image generation",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(block_state):
+ if block_state.prompt is not None and (
+ not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list)
+ ):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
+
+ @staticmethod
+ def encode_prompt(
+ components,
+ prompt: str,
+ device: torch.device | None = None,
+ prepare_unconditional_embeds: bool = True,
+ negative_prompt: str | None = None,
+ max_sequence_length: int = 512,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ prepare_unconditional_embeds (`bool`):
+ whether to use prepare unconditional embeddings or not
+ negative_prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ max_sequence_length (`int`, defaults to `512`):
+ The maximum number of text tokens to be used for the generation process.
+ """
+ device = device or components._execution_device
+ if not isinstance(prompt, list):
+ prompt = [prompt]
+ batch_size = len(prompt)
+
+ prompt_embeds = get_qwen_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ negative_prompt_embeds = None
+ if prepare_unconditional_embeds:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = get_qwen_prompt_embeds(
+ text_encoder=components.text_encoder,
+ tokenizer=components.tokenizer,
+ prompt=negative_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+
+ return prompt_embeds, negative_prompt_embeds
+
+ @torch.no_grad()
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ # Get inputs and intermediates
+ block_state = self.get_block_state(state)
+ self.check_inputs(block_state)
+
+ block_state.device = components._execution_device
+
+ # Encode input prompt
+ (
+ block_state.prompt_embeds,
+ block_state.negative_prompt_embeds,
+ ) = self.encode_prompt(
+ components=components,
+ prompt=block_state.prompt,
+ device=block_state.device,
+ prepare_unconditional_embeds=components.requires_unconditional_embeds,
+ negative_prompt=block_state.negative_prompt,
+ max_sequence_length=block_state.max_sequence_length,
+ )
+
+ # Add outputs
+ self.set_block_state(state, block_state)
+ return components, state
+
+
+class ZImageVaeImageEncoderStep(ModularPipelineBlocks):
+ model_name = "z-image"
+
+ @property
+ def description(self) -> str:
+ return "Vae Image Encoder step that generate condition_latents based on image to guide the image generation"
+
+ @property
+ def expected_components(self) -> list[ComponentSpec]:
+ return [
+ ComponentSpec("vae", AutoencoderKL),
+ ComponentSpec(
+ "image_processor",
+ VaeImageProcessor,
+ config=FrozenDict({"vae_scale_factor": 8 * 2}),
+ default_creation_method="from_config",
+ ),
+ ]
+
+ @property
+ def inputs(self) -> list[InputParam]:
+ return [
+ InputParam("image", type_hint=PIL.Image.Image, required=True),
+ InputParam("height"),
+ InputParam("width"),
+ InputParam("generator"),
+ ]
+
+ @property
+ def intermediate_outputs(self) -> list[OutputParam]:
+ return [
+ OutputParam(
+ "image_latents",
+ type_hint=torch.Tensor,
+ description="video latent representation with the first frame image condition",
+ ),
+ ]
+
+ @staticmethod
+ def check_inputs(components, block_state):
+ if (block_state.height is not None and block_state.height % components.vae_scale_factor_spatial != 0) or (
+ block_state.width is not None and block_state.width % components.vae_scale_factor_spatial != 0
+ ):
+ raise ValueError(
+ f"`height` and `width` have to be divisible by {components.vae_scale_factor_spatial} but are {block_state.height} and {block_state.width}."
+ )
+
+ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState:
+ block_state = self.get_block_state(state)
+ self.check_inputs(components, block_state)
+
+ image = block_state.image
+
+ device = components._execution_device
+ dtype = torch.float32
+ vae_dtype = components.vae.dtype
+
+ image_tensor = components.image_processor.preprocess(
+ image, height=block_state.height, width=block_state.width
+ ).to(device=device, dtype=dtype)
+
+ block_state.image_latents = encode_vae_image(
+ image_tensor=image_tensor,
+ vae=components.vae,
+ generator=block_state.generator,
+ device=device,
+ dtype=vae_dtype,
+ latent_channels=components.num_channels_latents,
+ )
+
+ self.set_block_state(state, block_state)
+ return components, state
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..7769d15ad725b13709a866e95ade0a7b53081d9d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py
@@ -0,0 +1,334 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...utils import logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import OutputParam
+from .before_denoise import (
+ ZImageAdditionalInputsStep,
+ ZImagePrepareLatentsStep,
+ ZImagePrepareLatentswithImageStep,
+ ZImageSetTimestepsStep,
+ ZImageSetTimestepsWithStrengthStep,
+ ZImageTextInputStep,
+)
+from .decoders import ZImageVaeDecoderStep
+from .denoise import (
+ ZImageDenoiseStep,
+)
+from .encoders import (
+ ZImageTextEncoderStep,
+ ZImageVaeImageEncoderStep,
+)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+# ====================
+# 1. DENOISE
+# ====================
+
+
+# text2image: inputs(text) -> set_timesteps -> prepare_latents -> denoise
+# auto_docstring
+class ZImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded conditions and runs the denoising process.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`int`, *optional*):
+ TODO: Add description.
+ width (`int`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 8):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageTextInputStep,
+ ZImagePrepareLatentsStep,
+ ZImageSetTimestepsStep,
+ ZImageDenoiseStep,
+ ]
+ block_names = ["input", "prepare_latents", "set_timesteps", "denoise"]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# image2image: inputs(text + image_latents) -> prepare_latents -> set_timesteps -> set_timesteps_with_strength -> prepare_latents_with_image -> denoise
+# auto_docstring
+class ZImageImage2ImageCoreDenoiseStep(SequentialPipelineBlocks):
+ """
+ denoise block that takes encoded text and image latent conditions and runs the denoising process.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`, *optional*, defaults to 8):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageTextInputStep,
+ ZImageAdditionalInputsStep(image_latent_inputs=["image_latents"]),
+ ZImagePrepareLatentsStep,
+ ZImageSetTimestepsStep,
+ ZImageSetTimestepsWithStrengthStep,
+ ZImagePrepareLatentswithImageStep,
+ ZImageDenoiseStep,
+ ]
+ block_names = [
+ "input",
+ "additional_inputs",
+ "prepare_latents",
+ "set_timesteps",
+ "set_timesteps_with_strength",
+ "prepare_latents_with_image",
+ "denoise",
+ ]
+
+ @property
+ def description(self):
+ return "denoise block that takes encoded text and image latent conditions and runs the denoising process."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("latents")]
+
+
+# auto_docstring
+class ZImageAutoDenoiseStep(AutoPipelineBlocks):
+ """
+ Denoise step that iteratively denoise the latents. This is a auto pipeline block that works for text2image and
+ image2image tasks. - `ZImageCoreDenoiseStep` (text2image) for text2image tasks. -
+ `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks. - if `image_latents` is provided,
+ `ZImageImage2ImageCoreDenoiseStep` will be used.
+ - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.
+
+ Components:
+ transformer (`ZImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
+ (`ClassifierFreeGuidance`)
+
+ Inputs:
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ prompt_embeds (`list`):
+ Pre-generated text embeddings. Can be generated from text_encoder step.
+ negative_prompt_embeds (`list`, *optional*):
+ Pre-generated negative text embeddings. Can be generated from text_encoder step.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+ Outputs:
+ latents (`Tensor`):
+ Denoised latents.
+ """
+
+ block_classes = [
+ ZImageImage2ImageCoreDenoiseStep,
+ ZImageCoreDenoiseStep,
+ ]
+ block_names = ["image2image", "text2image"]
+ block_trigger_inputs = ["image_latents", None]
+
+ @property
+ def description(self) -> str:
+ return (
+ "Denoise step that iteratively denoise the latents. "
+ "This is a auto pipeline block that works for text2image and image2image tasks."
+ " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks."
+ " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks."
+ + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n"
+ + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n"
+ )
+
+
+# auto_docstring
+class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks):
+ """
+ Vae Image Encoder step that encode the image to generate the image latents
+
+ Components:
+ vae (`AutoencoderKL`) image_processor (`VaeImageProcessor`)
+
+ Inputs:
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+
+ Outputs:
+ image_latents (`Tensor`):
+ video latent representation with the first frame image condition
+ """
+
+ block_classes = [ZImageVaeImageEncoderStep]
+ block_names = ["vae_encoder"]
+ block_trigger_inputs = ["image"]
+
+ @property
+ def description(self) -> str:
+ return "Vae Image Encoder step that encode the image to generate the image latents"
+ +"This is an auto pipeline block that works for image2image tasks."
+ +" - `ZImageVaeImageEncoderStep` is used when `image` is provided."
+ +" - if `image` is not provided, step will be skipped."
+
+
+# auto_docstring
+class ZImageAutoBlocks(SequentialPipelineBlocks):
+ """
+ Auto Modular pipeline for text-to-image and image-to-image using ZImage.
+
+ Supported workflows:
+ - `text2image`: requires `prompt`
+ - `image2image`: requires `image`, `prompt`
+
+ Components:
+ text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`) vae
+ (`AutoencoderKL`) image_processor (`VaeImageProcessor`) transformer (`ZImageTransformer2DModel`) scheduler
+ (`FlowMatchEulerDiscreteScheduler`)
+
+ Inputs:
+ prompt (`None`, *optional*):
+ TODO: Add description.
+ negative_prompt (`None`, *optional*):
+ TODO: Add description.
+ max_sequence_length (`None`, *optional*, defaults to 512):
+ TODO: Add description.
+ image (`Image`, *optional*):
+ TODO: Add description.
+ height (`None`, *optional*):
+ TODO: Add description.
+ width (`None`, *optional*):
+ TODO: Add description.
+ generator (`None`, *optional*):
+ TODO: Add description.
+ num_images_per_prompt (`None`, *optional*, defaults to 1):
+ TODO: Add description.
+ image_latents (`None`, *optional*):
+ TODO: Add description.
+ latents (`Tensor | NoneType`):
+ TODO: Add description.
+ num_inference_steps (`None`):
+ TODO: Add description.
+ sigmas (`None`, *optional*):
+ TODO: Add description.
+ strength (`None`, *optional*, defaults to 0.6):
+ TODO: Add description.
+ **denoiser_input_fields (`None`, *optional*):
+ The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+ output_type (`str`, *optional*, defaults to pil):
+ The type of the output images, can be 'pil', 'np', 'pt'
+
+ Outputs:
+ images (`list`):
+ Generated images.
+ """
+
+ block_classes = [
+ ZImageTextEncoderStep,
+ ZImageAutoVaeImageEncoderStep,
+ ZImageAutoDenoiseStep,
+ ZImageVaeDecoderStep,
+ ]
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
+ _workflow_map = {
+ "text2image": {"prompt": True},
+ "image2image": {"image": True, "prompt": True},
+ }
+
+ @property
+ def description(self) -> str:
+ return "Auto Modular pipeline for text-to-image and image-to-image using ZImage."
+
+ @property
+ def outputs(self):
+ return [OutputParam.template("images")]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_pipeline.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..7029288b43103dbd7e27044fbc70d283729d74a9
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/z_image/modular_pipeline.py
@@ -0,0 +1,72 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ...loaders import ZImageLoraLoaderMixin
+from ...utils import logging
+from ..modular_pipeline import ModularPipeline
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class ZImageModularPipeline(
+ ModularPipeline,
+ ZImageLoraLoaderMixin,
+):
+ """
+ A ModularPipeline for Z-Image.
+
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
+ """
+
+ default_blocks_name = "ZImageAutoBlocks"
+
+ @property
+ def default_height(self):
+ return 1024
+
+ @property
+ def default_width(self):
+ return 1024
+
+ @property
+ def vae_scale_factor_spatial(self):
+ vae_scale_factor_spatial = 16
+ if hasattr(self, "image_processor") and self.image_processor is not None:
+ vae_scale_factor_spatial = self.image_processor.config.vae_scale_factor
+ return vae_scale_factor_spatial
+
+ @property
+ def vae_scale_factor(self):
+ vae_scale_factor = 8
+ if hasattr(self, "vae") and self.vae is not None:
+ vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ return vae_scale_factor
+
+ @property
+ def num_channels_latents(self):
+ num_channels_latents = 16
+ if hasattr(self, "transformer") and self.transformer is not None:
+ num_channels_latents = self.transformer.config.in_channels
+ return num_channels_latents
+
+ @property
+ def requires_unconditional_embeds(self):
+ requires_unconditional_embeds = False
+
+ if hasattr(self, "guider") and self.guider is not None:
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+ return requires_unconditional_embeds
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/optimization.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/optimization.py
new file mode 100644
index 0000000000000000000000000000000000000000..152bc222b667ec3935c4ae1b1bfe4c5139ce46c3
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/optimization.py
@@ -0,0 +1,368 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch optimization for diffusion models."""
+
+import math
+from enum import Enum
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class SchedulerType(Enum):
+ LINEAR = "linear"
+ COSINE = "cosine"
+ COSINE_WITH_RESTARTS = "cosine_with_restarts"
+ POLYNOMIAL = "polynomial"
+ CONSTANT = "constant"
+ CONSTANT_WITH_WARMUP = "constant_with_warmup"
+ PIECEWISE_CONSTANT = "piecewise_constant"
+
+
+def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1) -> LambdaLR:
+ """
+ Create a schedule with a constant learning rate, using the learning rate set in optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+ return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
+
+
+def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1) -> LambdaLR:
+ """
+ Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
+ increases linearly between 0 and the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ def lr_lambda(current_step: int):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1.0, num_warmup_steps))
+ return 1.0
+
+ return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
+
+
+def get_piecewise_constant_schedule(optimizer: Optimizer, step_rules: str, last_epoch: int = -1) -> LambdaLR:
+ """
+ Create a schedule with a constant learning rate, using the learning rate set in optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ step_rules (`string`):
+ The rules for the learning rate. ex: rule_steps="1:10,0.1:20,0.01:30,0.005" it means that the learning rate
+ if multiple 1 for the first 10 steps, multiple 0.1 for the next 20 steps, multiple 0.01 for the next 30
+ steps and multiple 0.005 for the other steps.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ rules_dict = {}
+ rule_list = step_rules.split(",")
+ for rule_str in rule_list[:-1]:
+ value_str, steps_str = rule_str.split(":")
+ steps = int(steps_str)
+ value = float(value_str)
+ rules_dict[steps] = value
+ last_lr_multiple = float(rule_list[-1])
+
+ def create_rules_function(rules_dict, last_lr_multiple):
+ def rule_func(steps: int) -> float:
+ sorted_steps = sorted(rules_dict.keys())
+ for i, sorted_step in enumerate(sorted_steps):
+ if steps < sorted_step:
+ return rules_dict[sorted_steps[i]]
+ return last_lr_multiple
+
+ return rule_func
+
+ rules_func = create_rules_function(rules_dict, last_lr_multiple)
+
+ return LambdaLR(optimizer, rules_func, last_epoch=last_epoch)
+
+
+def get_linear_schedule_with_warmup(
+ optimizer: Optimizer,
+ num_warmup_steps: int,
+ num_training_steps: int,
+ last_epoch: int = -1,
+ f_min: float = 0.0,
+ f_max: float = 1.0,
+) -> LambdaLR:
+ """
+ Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+ a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+ f_min (`float`, *optional*, defaults to 0.0):
+ Minimum lr multiplier (floor of the linear decay). The lr will not fall below `f_min * initial_lr`.
+ f_max (`float`, *optional*, defaults to 1.0):
+ Maximum lr multiplier (peak reached after warmup). The lr peaks at `f_max * initial_lr`.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ def lr_lambda(current_step: int):
+ if current_step < num_warmup_steps:
+ return f_max * float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
+ return f_min + (f_max - f_min) * max(0.0, progress)
+
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_cosine_schedule_with_warmup(
+ optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+) -> LambdaLR:
+ """
+ Create a schedule with a learning rate that decreases following the values of the cosine function between the
+ initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+ initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ num_periods (`float`, *optional*, defaults to 0.5):
+ The number of periods of the cosine function in a schedule (the default is to just decrease from the max
+ value to 0 following a half-cosine).
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ def lr_lambda(current_step):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_cosine_with_hard_restarts_schedule_with_warmup(
+ optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
+) -> LambdaLR:
+ """
+ Create a schedule with a learning rate that decreases following the values of the cosine function between the
+ initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
+ linearly between 0 and the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ num_cycles (`int`, *optional*, defaults to 1):
+ The number of hard restarts to use.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ def lr_lambda(current_step):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ if progress >= 1.0:
+ return 0.0
+ return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
+
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_polynomial_decay_schedule_with_warmup(
+ optimizer: Optimizer,
+ num_warmup_steps: int,
+ num_training_steps: int,
+ lr_end: float = 1e-7,
+ power: float = 1.0,
+ last_epoch: int = -1,
+) -> LambdaLR:
+ """
+ Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
+ optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the
+ initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ lr_end (`float`, *optional*, defaults to 1e-7):
+ The end LR.
+ power (`float`, *optional*, defaults to 1.0):
+ Power factor.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
+ implementation at
+ https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+
+ """
+
+ lr_init = optimizer.defaults["lr"]
+ if not (lr_init > lr_end):
+ raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})")
+
+ def lr_lambda(current_step: int):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ elif current_step > num_training_steps:
+ return lr_end / lr_init # as LambdaLR multiplies by lr_init
+ else:
+ lr_range = lr_init - lr_end
+ decay_steps = num_training_steps - num_warmup_steps
+ pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
+ decay = lr_range * pct_remaining**power + lr_end
+ return decay / lr_init # as LambdaLR multiplies by lr_init
+
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+TYPE_TO_SCHEDULER_FUNCTION = {
+ SchedulerType.LINEAR: get_linear_schedule_with_warmup,
+ SchedulerType.COSINE: get_cosine_schedule_with_warmup,
+ SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
+ SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
+ SchedulerType.CONSTANT: get_constant_schedule,
+ SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
+ SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
+}
+
+
+def get_scheduler(
+ name: str | SchedulerType,
+ optimizer: Optimizer,
+ step_rules: str | None = None,
+ num_warmup_steps: int | None = None,
+ num_training_steps: int | None = None,
+ num_cycles: int = 1,
+ power: float = 1.0,
+ last_epoch: int = -1,
+) -> LambdaLR:
+ """
+ Unified API to get any scheduler from its name.
+
+ Args:
+ name (`str` or `SchedulerType`):
+ The name of the scheduler to use.
+ optimizer (`torch.optim.Optimizer`):
+ The optimizer that will be used during training.
+ step_rules (`str`, *optional*):
+ A string representing the step rules to use. This is only used by the `PIECEWISE_CONSTANT` scheduler.
+ num_warmup_steps (`int`, *optional*):
+ The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+ optional), the function will raise an error if it's unset and the scheduler type requires it.
+ num_training_steps (`int``, *optional*):
+ The number of training steps to do. This is not required by all schedulers (hence the argument being
+ optional), the function will raise an error if it's unset and the scheduler type requires it.
+ num_cycles (`int`, *optional*):
+ The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
+ power (`float`, *optional*, defaults to 1.0):
+ Power factor. See `POLYNOMIAL` scheduler
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+ """
+ name = SchedulerType(name)
+ schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+ if name == SchedulerType.CONSTANT:
+ return schedule_func(optimizer, last_epoch=last_epoch)
+
+ if name == SchedulerType.PIECEWISE_CONSTANT:
+ return schedule_func(optimizer, step_rules=step_rules, last_epoch=last_epoch)
+
+ # All other schedulers require `num_warmup_steps`
+ if num_warmup_steps is None:
+ raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+ if name == SchedulerType.CONSTANT_WITH_WARMUP:
+ return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, last_epoch=last_epoch)
+
+ # All other schedulers require `num_training_steps`
+ if num_training_steps is None:
+ raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+ if name == SchedulerType.COSINE_WITH_RESTARTS:
+ return schedule_func(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ num_cycles=num_cycles,
+ last_epoch=last_epoch,
+ )
+
+ if name == SchedulerType.POLYNOMIAL:
+ return schedule_func(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ power=power,
+ last_epoch=last_epoch,
+ )
+
+ return schedule_func(
+ optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, last_epoch=last_epoch
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/py.typed b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/py.typed
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/training_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/training_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fc2aca47d3e9551fd2dd6e6a2b421ed5ab3d911
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/training_utils.py
@@ -0,0 +1,963 @@
+import contextlib
+import copy
+import gc
+import math
+import random
+import re
+import warnings
+from contextlib import contextmanager
+from functools import partial
+from typing import Any, Iterable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+if torch.distributed.is_available():
+ from torch.distributed.fsdp import CPUOffload, ShardingStrategy
+ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+ from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+else:
+ CPUOffload = None
+ ShardingStrategy = None
+ FSDP = None
+ transformer_auto_wrap_policy = None
+
+from .models import UNet2DConditionModel
+from .pipelines import DiffusionPipeline
+from .schedulers import SchedulerMixin
+from .utils import (
+ convert_state_dict_to_diffusers,
+ convert_state_dict_to_peft,
+ deprecate,
+ is_accelerate_available,
+ is_peft_available,
+ is_torch_npu_available,
+ is_torchvision_available,
+ is_transformers_available,
+)
+
+
+if is_transformers_available():
+ import transformers
+
+ if transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():
+ import deepspeed
+
+if is_accelerate_available():
+ from accelerate.logging import get_logger
+
+if is_peft_available():
+ from peft import set_peft_model_state_dict
+
+if is_torchvision_available():
+ from torchvision import transforms
+
+if is_torch_npu_available():
+ import torch_npu # noqa: F401
+
+
+def set_seed(seed: int):
+ """
+ Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+ Args:
+ seed (`int`): The seed to set.
+
+ Returns:
+ `None`
+ """
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ if is_torch_npu_available():
+ torch.npu.manual_seed_all(seed)
+ else:
+ torch.cuda.manual_seed_all(seed)
+ # ^^ safe to call this function even if cuda is not available
+
+
+def compute_snr(noise_scheduler, timesteps):
+ """
+ Computes SNR as per
+ https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
+ for the given timesteps using the provided noise scheduler.
+
+ Args:
+ noise_scheduler (`NoiseScheduler`):
+ An object containing the noise schedule parameters, specifically `alphas_cumprod`, which is used to compute
+ the SNR values.
+ timesteps (`torch.Tensor`):
+ A tensor of timesteps for which the SNR is computed.
+
+ Returns:
+ `torch.Tensor`: A tensor containing the computed SNR values for each timestep.
+ """
+ alphas_cumprod = noise_scheduler.alphas_cumprod
+ sqrt_alphas_cumprod = alphas_cumprod**0.5
+ sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
+
+ # Expand the tensors.
+ # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
+ while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
+ alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
+
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
+ while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
+ sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
+
+ # Compute SNR.
+ snr = (alpha / sigma) ** 2
+ return snr
+
+
+def compute_confidence_aware_loss(
+ logits: torch.Tensor,
+ labels: torch.Tensor,
+ *,
+ lambda_conf: float = 0.0,
+ temperature: float = 1.0,
+ per_token_weights: torch.Tensor | None = None,
+ ignore_index: int = -100,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Computes a confidence-aware training loss for token classification-style heads.
+
+ This loss combines:
+ - `loss_sft`: standard supervised cross-entropy on all non-ignored labels.
+ - `loss_conf`: an entropy penalty applied only on tokens that are already predicted correctly.
+
+ Args:
+ logits (`torch.Tensor`): Logits of shape `(..., vocab_size)`.
+ labels (`torch.Tensor`): Labels of shape `(...)`, matching `logits.shape[:-1]`. Values set to `ignore_index`
+ are excluded from both losses.
+ lambda_conf (`float`, *optional*, defaults to `0.0`): Weight for the confidence term.
+ temperature (`float`, *optional*, defaults to `1.0`): Temperature used for the entropy term only. Lower values
+ sharpen the distribution and change the strength of the confidence gradients.
+ per_token_weights (`torch.Tensor`, *optional*): Optional weights of shape `(...)` to reweight both losses per
+ token (e.g. schedule-aware weights). Tokens with weight `0` contribute nothing.
+ ignore_index (`int`, *optional*, defaults to `-100`): Ignore index for labels.
+
+ Returns:
+ `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`: `(loss, loss_sft, loss_conf)`.
+ """
+ if logits.ndim < 2:
+ raise ValueError(f"`logits` must have at least 2 dims, got shape {tuple(logits.shape)}.")
+ if labels.shape != logits.shape[:-1]:
+ raise ValueError(
+ f"`labels` shape must match `logits.shape[:-1]`, got labels={tuple(labels.shape)} logits={tuple(logits.shape)}."
+ )
+ if temperature <= 0:
+ raise ValueError(f"`temperature` must be > 0, got {temperature}.")
+
+ valid = labels.ne(ignore_index)
+ if per_token_weights is None:
+ weights = torch.ones_like(labels, dtype=logits.dtype)
+ else:
+ if per_token_weights.shape != labels.shape:
+ raise ValueError(
+ f"`per_token_weights` shape must match `labels` shape, got {tuple(per_token_weights.shape)} != {tuple(labels.shape)}."
+ )
+ weights = per_token_weights.to(dtype=logits.dtype)
+
+ # Supervised CE (optionally weighted).
+ vocab_size = logits.shape[-1]
+ per_token_nll = F.cross_entropy(
+ logits.reshape(-1, vocab_size),
+ labels.reshape(-1),
+ reduction="none",
+ ignore_index=ignore_index,
+ ).reshape_as(labels)
+
+ denom_sft = (weights * valid.to(weights.dtype)).sum().clamp_min(1)
+ loss_sft = (per_token_nll * weights * valid.to(per_token_nll.dtype)).sum() / denom_sft
+
+ # Confidence loss: penalize entropy only where prediction is already correct.
+ if lambda_conf == 0.0:
+ loss_conf = torch.zeros((), device=logits.device, dtype=loss_sft.dtype)
+ return loss_sft, loss_sft, loss_conf
+
+ with torch.no_grad():
+ pred = logits.argmax(dim=-1)
+ correct = valid & pred.eq(labels)
+
+ scaled_logits = logits.float()
+ if temperature != 1.0:
+ scaled_logits = scaled_logits / float(temperature)
+
+ probs = torch.softmax(scaled_logits, dim=-1)
+ eps = torch.finfo(probs.dtype).tiny
+ log_probs = torch.log(probs.clamp_min(eps))
+ entropy = -(probs * log_probs).sum(dim=-1).to(dtype=logits.dtype)
+
+ denom_conf = (weights * correct.to(weights.dtype)).sum().clamp_min(1)
+ loss_conf = (entropy * weights * correct.to(entropy.dtype)).sum() / denom_conf
+
+ loss = loss_sft + float(lambda_conf) * loss_conf
+ return loss, loss_sft, loss_conf
+
+
+def resolve_interpolation_mode(interpolation_type: str):
+ """
+ Maps a string describing an interpolation function to the corresponding torchvision `InterpolationMode` enum. The
+ full list of supported enums is documented at
+ https://pytorch.org/vision/0.9/transforms.html#torchvision.transforms.functional.InterpolationMode.
+
+ Args:
+ interpolation_type (`str`):
+ A string describing an interpolation method. Currently, `bilinear`, `bicubic`, `box`, `nearest`,
+ `nearest_exact`, `hamming`, and `lanczos` are supported, corresponding to the supported interpolation modes
+ in torchvision.
+
+ Returns:
+ `torchvision.transforms.InterpolationMode`: an `InterpolationMode` enum used by torchvision's `resize`
+ transform.
+ """
+ if not is_torchvision_available():
+ raise ImportError(
+ "Please make sure to install `torchvision` to be able to use the `resolve_interpolation_mode()` function."
+ )
+
+ if interpolation_type == "bilinear":
+ interpolation_mode = transforms.InterpolationMode.BILINEAR
+ elif interpolation_type == "bicubic":
+ interpolation_mode = transforms.InterpolationMode.BICUBIC
+ elif interpolation_type == "box":
+ interpolation_mode = transforms.InterpolationMode.BOX
+ elif interpolation_type == "nearest":
+ interpolation_mode = transforms.InterpolationMode.NEAREST
+ elif interpolation_type == "nearest_exact":
+ interpolation_mode = transforms.InterpolationMode.NEAREST_EXACT
+ elif interpolation_type == "hamming":
+ interpolation_mode = transforms.InterpolationMode.HAMMING
+ elif interpolation_type == "lanczos":
+ interpolation_mode = transforms.InterpolationMode.LANCZOS
+ else:
+ raise ValueError(
+ f"The given interpolation mode {interpolation_type} is not supported. Currently supported interpolation"
+ f" modes are `bilinear`, `bicubic`, `box`, `nearest`, `nearest_exact`, `hamming`, and `lanczos`."
+ )
+
+ return interpolation_mode
+
+
+def compute_dream_and_update_latents(
+ unet: UNet2DConditionModel,
+ noise_scheduler: SchedulerMixin,
+ timesteps: torch.Tensor,
+ noise: torch.Tensor,
+ noisy_latents: torch.Tensor,
+ target: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ dream_detail_preservation: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Implements "DREAM (Diffusion Rectification and Estimation-Adaptive Models)" from
+ https://huggingface.co/papers/2312.00210. DREAM helps align training with sampling to help training be more
+ efficient and accurate at the cost of an extra forward step without gradients.
+
+ Args:
+ `unet`: The state unet to use to make a prediction.
+ `noise_scheduler`: The noise scheduler used to add noise for the given timestep.
+ `timesteps`: The timesteps for the noise_scheduler to user.
+ `noise`: A tensor of noise in the shape of noisy_latents.
+ `noisy_latents`: Previously noise latents from the training loop.
+ `target`: The ground-truth tensor to predict after eps is removed.
+ `encoder_hidden_states`: Text embeddings from the text model.
+ `dream_detail_preservation`: A float value that indicates detail preservation level.
+ See reference.
+
+ Returns:
+ `tuple[torch.Tensor, torch.Tensor]`: Adjusted noisy_latents and target.
+ """
+ alphas_cumprod = noise_scheduler.alphas_cumprod.to(timesteps.device)[timesteps, None, None, None]
+ sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
+
+ # The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments.
+ dream_lambda = sqrt_one_minus_alphas_cumprod**dream_detail_preservation
+
+ pred = None
+ with torch.no_grad():
+ pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+ _noisy_latents, _target = (None, None)
+ if noise_scheduler.config.prediction_type == "epsilon":
+ predicted_noise = pred
+ delta_noise = (noise - predicted_noise).detach()
+ delta_noise.mul_(dream_lambda)
+ _noisy_latents = noisy_latents.add(sqrt_one_minus_alphas_cumprod * delta_noise)
+ _target = target.add(delta_noise)
+ elif noise_scheduler.config.prediction_type == "v_prediction":
+ raise NotImplementedError("DREAM has not been implemented for v-prediction")
+ else:
+ raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+ return _noisy_latents, _target
+
+
+def unet_lora_state_dict(unet: UNet2DConditionModel) -> dict[str, torch.Tensor]:
+ r"""
+ Returns:
+ A state dict containing just the LoRA parameters.
+ """
+ lora_state_dict = {}
+
+ for name, module in unet.named_modules():
+ if hasattr(module, "set_lora_layer"):
+ lora_layer = getattr(module, "lora_layer")
+ if lora_layer is not None:
+ current_lora_layer_sd = lora_layer.state_dict()
+ for lora_layer_matrix_name, lora_param in current_lora_layer_sd.items():
+ # The matrix name can either be "down" or "up".
+ lora_state_dict[f"{name}.lora.{lora_layer_matrix_name}"] = lora_param
+
+ return lora_state_dict
+
+
+def cast_training_params(model: torch.nn.Module | list[torch.nn.Module], dtype=torch.float32):
+ """
+ Casts the training parameters of the model to the specified data type.
+
+ Args:
+ model: The PyTorch model whose parameters will be cast.
+ dtype: The data type to which the model parameters will be cast.
+ """
+ if not isinstance(model, list):
+ model = [model]
+ for m in model:
+ for param in m.parameters():
+ # only upcast trainable parameters into fp32
+ if param.requires_grad:
+ param.data = param.to(dtype)
+
+
+def _set_state_dict_into_text_encoder(
+ lora_state_dict: dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module
+):
+ """
+ Sets the `lora_state_dict` into `text_encoder` coming from `transformers`.
+
+ Args:
+ lora_state_dict: The state dictionary to be set.
+ prefix: String identifier to retrieve the portion of the state dict that belongs to `text_encoder`.
+ text_encoder: Where the `lora_state_dict` is to be set.
+ """
+
+ text_encoder_state_dict = {
+ f"{k.replace(prefix, '')}": v for k, v in lora_state_dict.items() if k.startswith(prefix)
+ }
+ text_encoder_state_dict = convert_state_dict_to_peft(convert_state_dict_to_diffusers(text_encoder_state_dict))
+ set_peft_model_state_dict(text_encoder, text_encoder_state_dict, adapter_name="default")
+
+
+def _collate_lora_metadata(modules_to_save: dict[str, torch.nn.Module]) -> dict[str, Any]:
+ metadatas = {}
+ for module_name, module in modules_to_save.items():
+ if module is not None:
+ metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict()
+ return metadatas
+
+
+def compute_density_for_timestep_sampling(
+ weighting_scheme: str,
+ batch_size: int,
+ logit_mean: float = None,
+ logit_std: float = None,
+ mode_scale: float = None,
+ device: torch.device | str = "cpu",
+ generator: torch.Generator | None = None,
+):
+ """
+ Compute the density for sampling the timesteps when doing SD3 training.
+
+ Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+ SD3 paper reference: https://huggingface.co/papers/2403.03206v1.
+ """
+ if weighting_scheme == "logit_normal":
+ u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device=device, generator=generator)
+ u = torch.nn.functional.sigmoid(u)
+ elif weighting_scheme == "mode":
+ u = torch.rand(size=(batch_size,), device=device, generator=generator)
+ u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
+ else:
+ u = torch.rand(size=(batch_size,), device=device, generator=generator)
+ return u
+
+
+def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None):
+ """
+ Computes loss weighting scheme for SD3 training.
+
+ Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+ SD3 paper reference: https://huggingface.co/papers/2403.03206v1.
+ """
+ if weighting_scheme == "sigma_sqrt":
+ weighting = (sigmas**-2.0).float()
+ elif weighting_scheme == "cosmap":
+ bot = 1 - 2 * sigmas + 2 * sigmas**2
+ weighting = 2 / (math.pi * bot)
+ else:
+ weighting = torch.ones_like(sigmas)
+ return weighting
+
+
+def free_memory():
+ """
+ Runs garbage collection. Then clears the cache of the available accelerator.
+ """
+ gc.collect()
+
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ elif torch.backends.mps.is_available():
+ torch.mps.empty_cache()
+ elif is_torch_npu_available():
+ torch_npu.npu.empty_cache()
+ elif hasattr(torch, "xpu") and torch.xpu.is_available():
+ torch.xpu.empty_cache()
+
+
+@contextmanager
+def offload_models(*modules: torch.nn.Module | DiffusionPipeline, device: str | torch.device, offload: bool = True):
+ """
+ Context manager that, if offload=True, moves each module to `device` on enter, then moves it back to its original
+ device on exit.
+
+ Args:
+ device (`str` or `torch.Device`): Device to move the `modules` to.
+ offload (`bool`): Flag to enable offloading.
+ """
+ if offload:
+ is_model = not any(isinstance(m, DiffusionPipeline) for m in modules)
+ # record where each module was
+ if is_model:
+ original_devices = [next(m.parameters()).device for m in modules]
+ else:
+ assert len(modules) == 1
+ # For DiffusionPipeline, wrap the device in a list to make it iterable
+ original_devices = [modules[0].device]
+ # move to target device
+ for m in modules:
+ m.to(device)
+
+ try:
+ yield
+ finally:
+ if offload:
+ # move back to original devices
+ for m, orig_dev in zip(modules, original_devices):
+ m.to(orig_dev)
+
+
+def parse_buckets_string(buckets_str):
+ """Parses a string defining buckets into a list of (height, width) tuples."""
+ if not buckets_str:
+ raise ValueError("Bucket string cannot be empty.")
+
+ bucket_pairs = buckets_str.strip().split(";")
+ parsed_buckets = []
+ for pair_str in bucket_pairs:
+ match = re.match(r"^\s*(\d+)\s*,\s*(\d+)\s*$", pair_str)
+ if not match:
+ raise ValueError(f"Invalid bucket format: '{pair_str}'. Expected 'height,width'.")
+ try:
+ height = int(match.group(1))
+ width = int(match.group(2))
+ if height <= 0 or width <= 0:
+ raise ValueError("Bucket dimensions must be positive integers.")
+ if height % 8 != 0 or width % 8 != 0:
+ warnings.warn(f"Bucket dimension ({height},{width}) not divisible by 8. This might cause issues.")
+ parsed_buckets.append((height, width))
+ except ValueError as e:
+ raise ValueError(f"Invalid integer in bucket pair '{pair_str}': {e}") from e
+
+ if not parsed_buckets:
+ raise ValueError("No valid buckets found in the provided string.")
+
+ return parsed_buckets
+
+
+def find_nearest_bucket(h, w, bucket_options):
+ """Finds the closes bucket to the given height and width."""
+ min_metric = float("inf")
+ best_bucket_idx = None
+ for bucket_idx, (bucket_h, bucket_w) in enumerate(bucket_options):
+ metric = abs(h * bucket_w - w * bucket_h)
+ if metric <= min_metric:
+ min_metric = metric
+ best_bucket_idx = bucket_idx
+ return best_bucket_idx
+
+
+# Common aspect ratios (width, height) used to generate an aspect-ratio bucket ladder on the fly.
+_DEFAULT_BUCKET_ASPECT_RATIOS = [
+ (1, 1),
+ (4, 5),
+ (5, 4),
+ (2, 3),
+ (3, 2),
+ (3, 4),
+ (4, 3),
+ (9, 16),
+ (16, 9),
+ (1, 2),
+ (2, 1),
+]
+
+
+def generate_aspect_ratio_buckets(resolution, divisibility=16, base_resolutions=None):
+ """Generate a discrete list of ``(height, width)`` buckets for a target pixel budget.
+
+ Buckets are sized so that ``height * width`` is close to ``resolution ** 2`` while spanning a range of aspect
+ ratios, with each dimension rounded to a multiple of ``divisibility``. When ``base_resolutions`` (a list of
+ ``(height, width)`` pairs, e.g. a model's preferred resolutions) is provided, its aspect ratios are used instead of
+ the default set and rescaled to the target pixel budget.
+
+ Args:
+ resolution (`int`): Target resolution; the pixel budget is ``resolution ** 2``.
+ divisibility (`int`, *optional*, defaults to 16): Each bucket dimension is rounded to a
+ multiple of this value (16 satisfies the patch-size constraints of the Flux/Z-Image VAEs).
+ base_resolutions (`list[tuple[int, int]]`, *optional*): ``(height, width)`` pairs whose aspect
+ ratios seed the ladder instead of the default aspect-ratio set.
+
+ Returns:
+ `list[tuple[int, int]]`: The unique ``(height, width)`` buckets.
+ """
+ target_pixels = resolution * resolution
+ if base_resolutions is not None:
+ aspect_ratios = [(w, h) for (h, w) in base_resolutions]
+ else:
+ aspect_ratios = _DEFAULT_BUCKET_ASPECT_RATIOS
+
+ buckets = []
+ seen = set()
+ for rw, rh in aspect_ratios:
+ aspect = rw / rh
+ # h * w == target_pixels and w / h == aspect => h = sqrt(target_pixels / aspect)
+ h = (target_pixels / aspect) ** 0.5
+ w = h * aspect
+ h = max(divisibility, round(h / divisibility) * divisibility)
+ w = max(divisibility, round(w / divisibility) * divisibility)
+ if (h, w) not in seen:
+ seen.add((h, w))
+ buckets.append((h, w))
+ return buckets
+
+
+def _to_cpu_contiguous(state_dicts) -> dict:
+ return {k: v.detach().cpu().contiguous() if isinstance(v, torch.Tensor) else v for k, v in state_dicts.items()}
+
+
+def get_fsdp_kwargs_from_accelerator(accelerator) -> dict:
+ """
+ Extract and convert FSDP config from Accelerator into PyTorch FSDP kwargs.
+ """
+
+ kwargs = {}
+ fsdp_state = getattr(accelerator.state, "fsdp_plugin", None)
+
+ if fsdp_state is None:
+ raise ValueError("Accelerate isn't configured to handle FSDP. Please update your installation.")
+
+ fsdp_plugin = accelerator.state.fsdp_plugin
+
+ if fsdp_plugin is None:
+ # FSDP not enabled in Accelerator
+ kwargs["sharding_strategy"] = ShardingStrategy.FULL_SHARD
+ else:
+ # FSDP is enabled → use plugin's strategy, or default if None
+ kwargs["sharding_strategy"] = fsdp_plugin.sharding_strategy or ShardingStrategy.FULL_SHARD
+
+ return kwargs
+
+
+def wrap_with_fsdp(
+ model: torch.nn.Module,
+ device: str | torch.device,
+ offload: bool = True,
+ use_orig_params: bool = True,
+ limit_all_gathers: bool = True,
+ fsdp_kwargs: dict[str, Any] | None = None,
+ transformer_layer_cls: set[type[torch.nn.Module]] | None = None,
+) -> FSDP:
+ """
+ Wrap a model with FSDP using common defaults and optional transformer auto-wrapping.
+
+ Args:
+ model: Model to wrap
+ device: Target device (e.g., accelerator.device)
+ offload: Whether to enable CPU parameter offloading
+ use_orig_params: Whether to use original parameters
+ limit_all_gathers: Whether to limit all gathers
+ fsdp_kwargs: FSDP arguments (sharding_strategy, etc.) — usually from Accelerate config
+ transformer_layer_cls: Classes for auto-wrapping (if not using policy from fsdp_kwargs)
+
+ Returns:
+ FSDP-wrapped model
+ """
+
+ logger = get_logger(__name__)
+
+ if transformer_layer_cls is None:
+ # Set the default layers if transformer_layer_cls is not provided
+ transformer_layer_cls = type(model.model.language_model.layers[0])
+ logger.info(f"transformer_layer_cls is not provided, auto-inferred as {transformer_layer_cls.__name__}")
+
+ # Add auto-wrap policy if transformer layers specified
+ auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer_cls})
+
+ config = {
+ "device_id": device,
+ "cpu_offload": CPUOffload(offload_params=offload) if offload else None,
+ "use_orig_params": use_orig_params,
+ "limit_all_gathers": limit_all_gathers,
+ "auto_wrap_policy": auto_wrap_policy,
+ }
+
+ if fsdp_kwargs:
+ config.update(fsdp_kwargs)
+
+ fsdp_model = FSDP(model, **config)
+ return fsdp_model
+
+
+# Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14
+class EMAModel:
+ """
+ Exponential Moving Average of models weights
+ """
+
+ def __init__(
+ self,
+ parameters: Iterable[torch.nn.Parameter],
+ decay: float = 0.9999,
+ min_decay: float = 0.0,
+ update_after_step: int = 0,
+ use_ema_warmup: bool = False,
+ inv_gamma: float | int = 1.0,
+ power: float | int = 2 / 3,
+ foreach: bool = False,
+ model_cls: Any | None = None,
+ model_config: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ parameters (Iterable[torch.nn.Parameter]): The parameters to track.
+ decay (float): The decay factor for the exponential moving average.
+ min_decay (float): The minimum decay factor for the exponential moving average.
+ update_after_step (int): The number of steps to wait before starting to update the EMA weights.
+ use_ema_warmup (bool): Whether to use EMA warmup.
+ inv_gamma (float):
+ Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
+ power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
+ foreach (bool): Use torch._foreach functions for updating shadow parameters. Should be faster.
+ device (str | torch.device | None): The device to store the EMA weights on. If None, the EMA
+ weights will be stored on CPU.
+
+ @crowsonkb's notes on EMA Warmup:
+ If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
+ to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
+ gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
+ at 215.4k steps).
+ """
+
+ if isinstance(parameters, torch.nn.Module):
+ deprecation_message = (
+ "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "
+ "Please pass the parameters of the module instead."
+ )
+ deprecate(
+ "passing a `torch.nn.Module` to `ExponentialMovingAverage`",
+ "1.0.0",
+ deprecation_message,
+ standard_warn=False,
+ )
+ parameters = parameters.parameters()
+
+ # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
+ use_ema_warmup = True
+
+ if kwargs.get("max_value", None) is not None:
+ deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."
+ deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)
+ decay = kwargs["max_value"]
+
+ if kwargs.get("min_value", None) is not None:
+ deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."
+ deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)
+ min_decay = kwargs["min_value"]
+
+ parameters = list(parameters)
+ self.shadow_params = [p.clone().detach() for p in parameters]
+
+ if kwargs.get("device", None) is not None:
+ deprecation_message = "The `device` argument is deprecated. Please use `to` instead."
+ deprecate("device", "1.0.0", deprecation_message, standard_warn=False)
+ self.to(device=kwargs["device"])
+
+ self.temp_stored_params = None
+
+ self.decay = decay
+ self.min_decay = min_decay
+ self.update_after_step = update_after_step
+ self.use_ema_warmup = use_ema_warmup
+ self.inv_gamma = inv_gamma
+ self.power = power
+ self.optimization_step = 0
+ self.cur_decay_value = None # set in `step()`
+ self.foreach = foreach
+
+ self.model_cls = model_cls
+ self.model_config = model_config
+
+ @classmethod
+ def from_pretrained(cls, path, model_cls, foreach=False) -> "EMAModel":
+ _, ema_kwargs = model_cls.from_config(path, return_unused_kwargs=True)
+ model = model_cls.from_pretrained(path)
+
+ ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config, foreach=foreach)
+
+ ema_model.load_state_dict(ema_kwargs)
+ return ema_model
+
+ def save_pretrained(self, path):
+ if self.model_cls is None:
+ raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")
+
+ if self.model_config is None:
+ raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")
+
+ model = self.model_cls.from_config(self.model_config)
+ state_dict = self.state_dict()
+ state_dict.pop("shadow_params", None)
+
+ model.register_to_config(**state_dict)
+ self.copy_to(model.parameters())
+ model.save_pretrained(path)
+
+ def get_decay(self, optimization_step: int) -> float:
+ """
+ Compute the decay factor for the exponential moving average.
+ """
+ step = max(0, optimization_step - self.update_after_step - 1)
+
+ if step <= 0:
+ return 0.0
+
+ if self.use_ema_warmup:
+ cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power
+ else:
+ cur_decay_value = (1 + step) / (10 + step)
+
+ cur_decay_value = min(cur_decay_value, self.decay)
+ # make sure decay is not smaller than min_decay
+ cur_decay_value = max(cur_decay_value, self.min_decay)
+ return cur_decay_value
+
+ @torch.no_grad()
+ def step(self, parameters: Iterable[torch.nn.Parameter]):
+ if isinstance(parameters, torch.nn.Module):
+ deprecation_message = (
+ "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. "
+ "Please pass the parameters of the module instead."
+ )
+ deprecate(
+ "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`",
+ "1.0.0",
+ deprecation_message,
+ standard_warn=False,
+ )
+ parameters = parameters.parameters()
+
+ parameters = list(parameters)
+
+ self.optimization_step += 1
+
+ # Compute the decay factor for the exponential moving average.
+ decay = self.get_decay(self.optimization_step)
+ self.cur_decay_value = decay
+ one_minus_decay = 1 - decay
+
+ context_manager = contextlib.nullcontext()
+
+ if self.foreach:
+ if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():
+ context_manager = deepspeed.zero.GatheredParameters(parameters, modifier_rank=None)
+
+ with context_manager:
+ params_grad = [param for param in parameters if param.requires_grad]
+ s_params_grad = [
+ s_param for s_param, param in zip(self.shadow_params, parameters) if param.requires_grad
+ ]
+
+ if len(params_grad) < len(parameters):
+ torch._foreach_copy_(
+ [s_param for s_param, param in zip(self.shadow_params, parameters) if not param.requires_grad],
+ [param for param in parameters if not param.requires_grad],
+ non_blocking=True,
+ )
+
+ torch._foreach_sub_(
+ s_params_grad, torch._foreach_sub(s_params_grad, params_grad), alpha=one_minus_decay
+ )
+
+ else:
+ for s_param, param in zip(self.shadow_params, parameters):
+ if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():
+ context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None)
+
+ with context_manager:
+ if param.requires_grad:
+ s_param.sub_(one_minus_decay * (s_param - param))
+ else:
+ s_param.copy_(param)
+
+ def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+ """
+ Copy current averaged parameters into given collection of parameters.
+
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ updated with the stored moving averages. If `None`, the parameters with which this
+ `ExponentialMovingAverage` was initialized will be used.
+ """
+ parameters = list(parameters)
+ if self.foreach:
+ torch._foreach_copy_(
+ [param.data for param in parameters],
+ [s_param.to(param.device).data for s_param, param in zip(self.shadow_params, parameters)],
+ )
+ else:
+ for s_param, param in zip(self.shadow_params, parameters):
+ param.data.copy_(s_param.to(param.device).data)
+
+ def pin_memory(self) -> None:
+ r"""
+ Move internal buffers of the ExponentialMovingAverage to pinned memory. Useful for non-blocking transfers for
+ offloading EMA params to the host.
+ """
+
+ self.shadow_params = [p.pin_memory() for p in self.shadow_params]
+
+ def to(self, device=None, dtype=None, non_blocking=False) -> None:
+ r"""
+ Move internal buffers of the ExponentialMovingAverage to `device`.
+
+ Args:
+ device: like `device` argument to `torch.Tensor.to`
+ """
+ # .to() on the tensors handles None correctly
+ self.shadow_params = [
+ p.to(device=device, dtype=dtype, non_blocking=non_blocking)
+ if p.is_floating_point()
+ else p.to(device=device, non_blocking=non_blocking)
+ for p in self.shadow_params
+ ]
+
+ def state_dict(self) -> dict:
+ r"""
+ Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during
+ checkpointing to save the ema state dict.
+ """
+ # Following PyTorch conventions, references to tensors are returned:
+ # "returns a reference to the state and not its copy!" -
+ # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict
+ return {
+ "decay": self.decay,
+ "min_decay": self.min_decay,
+ "optimization_step": self.optimization_step,
+ "update_after_step": self.update_after_step,
+ "use_ema_warmup": self.use_ema_warmup,
+ "inv_gamma": self.inv_gamma,
+ "power": self.power,
+ "shadow_params": self.shadow_params,
+ }
+
+ def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+ r"""
+ Saves the current parameters for restoring later.
+
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`. The parameters to be temporarily stored.
+ """
+ self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]
+
+ def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+ r"""
+ Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters
+ without: affecting the original optimization process. Store the parameters before the `copy_to()` method. After
+ validation (or model saving), use this to restore the former parameters.
+
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ updated with the stored parameters. If `None`, the parameters with which this
+ `ExponentialMovingAverage` was initialized will be used.
+ """
+
+ if self.temp_stored_params is None:
+ raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights to `restore()`")
+ if self.foreach:
+ torch._foreach_copy_(
+ [param.data for param in parameters], [c_param.data for c_param in self.temp_stored_params]
+ )
+ else:
+ for c_param, param in zip(self.temp_stored_params, parameters):
+ param.data.copy_(c_param.data)
+
+ # Better memory-wise.
+ self.temp_stored_params = None
+
+ def load_state_dict(self, state_dict: dict) -> None:
+ r"""
+ Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the
+ ema state dict.
+
+ Args:
+ state_dict (dict): EMA state. Should be an object returned
+ from a call to :meth:`state_dict`.
+ """
+ # deepcopy, to be consistent with module API
+ state_dict = copy.deepcopy(state_dict)
+
+ self.decay = state_dict.get("decay", self.decay)
+ if self.decay < 0.0 or self.decay > 1.0:
+ raise ValueError("Decay must be between 0 and 1")
+
+ self.min_decay = state_dict.get("min_decay", self.min_decay)
+ if not isinstance(self.min_decay, float):
+ raise ValueError("Invalid min_decay")
+
+ self.optimization_step = state_dict.get("optimization_step", self.optimization_step)
+ if not isinstance(self.optimization_step, int):
+ raise ValueError("Invalid optimization_step")
+
+ self.update_after_step = state_dict.get("update_after_step", self.update_after_step)
+ if not isinstance(self.update_after_step, int):
+ raise ValueError("Invalid update_after_step")
+
+ self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)
+ if not isinstance(self.use_ema_warmup, bool):
+ raise ValueError("Invalid use_ema_warmup")
+
+ self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)
+ if not isinstance(self.inv_gamma, (float, int)):
+ raise ValueError("Invalid inv_gamma")
+
+ self.power = state_dict.get("power", self.power)
+ if not isinstance(self.power, (float, int)):
+ raise ValueError("Invalid power")
+
+ shadow_params = state_dict.get("shadow_params", None)
+ if shadow_params is not None:
+ self.shadow_params = shadow_params
+ if not isinstance(self.shadow_params, list):
+ raise ValueError("shadow_params must be a list")
+ if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):
+ raise ValueError("shadow_params must all be Tensors")
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/__init__.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0cc5ead5c5c6f31b54c1e0c3d7c7e9317750d716
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/__init__.py
@@ -0,0 +1,176 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+
+from packaging import version
+
+from .. import __version__
+from .constants import (
+ CONFIG_NAME,
+ DEFAULT_HF_PARALLEL_LOADING_WORKERS,
+ DEPRECATED_REVISION_ARGS,
+ DIFFUSERS_DYNAMIC_MODULE_NAME,
+ DIFFUSERS_LOAD_ID_FIELDS,
+ FLASHPACK_FILE_EXTENSION,
+ FLASHPACK_WEIGHTS_NAME,
+ GGUF_FILE_EXTENSION,
+ HF_ENABLE_PARALLEL_LOADING,
+ HF_MODULES_CACHE,
+ HUGGINGFACE_CO_RESOLVE_ENDPOINT,
+ MIN_PEFT_VERSION,
+ ONNX_EXTERNAL_WEIGHTS_NAME,
+ ONNX_WEIGHTS_NAME,
+ SAFE_WEIGHTS_INDEX_NAME,
+ SAFETENSORS_FILE_EXTENSION,
+ SAFETENSORS_WEIGHTS_NAME,
+ USE_PEFT_BACKEND,
+ WEIGHTS_INDEX_NAME,
+ WEIGHTS_NAME,
+)
+from .deprecation_utils import _maybe_remap_transformers_class, _resolve_dtype, deprecate
+from .doc_utils import replace_example_docstring
+from .dynamic_modules_utils import get_class_from_dynamic_module
+from .export_utils import encode_video, export_to_gif, export_to_obj, export_to_ply, export_to_video
+from .hub_utils import (
+ PushToHubMixin,
+ _add_variant,
+ _get_checkpoint_shard_files,
+ _get_model_file,
+ extract_commit_hash,
+ http_user_agent,
+)
+from .import_utils import (
+ BACKENDS_MAPPING,
+ DIFFUSERS_SLOW_IMPORT,
+ ENV_VARS_TRUE_AND_AUTO_VALUES,
+ ENV_VARS_TRUE_VALUES,
+ USE_TF,
+ USE_TORCH,
+ DummyObject,
+ OptionalDependencyNotAvailable,
+ _LazyModule,
+ get_objects_from_module,
+ is_accelerate_available,
+ is_accelerate_version,
+ is_aiter_available,
+ is_aiter_version,
+ is_auto_round_available,
+ is_av_available,
+ is_better_profanity_available,
+ is_bitsandbytes_available,
+ is_bitsandbytes_version,
+ is_bs4_available,
+ is_cosmos_guardrail_available,
+ is_flash_attn_3_available,
+ is_flash_attn_available,
+ is_flash_attn_version,
+ is_flashpack_available,
+ is_ftfy_available,
+ is_gguf_available,
+ is_gguf_version,
+ is_google_colab,
+ is_hf_hub_version,
+ is_hpu_available,
+ is_inflect_available,
+ is_invisible_watermark_available,
+ is_kernels_available,
+ is_kernels_version,
+ is_kornia_available,
+ is_librosa_available,
+ is_matplotlib_available,
+ is_nltk_available,
+ is_note_seq_available,
+ is_nvidia_modelopt_available,
+ is_nvidia_modelopt_version,
+ is_onnx_available,
+ is_opencv_available,
+ is_optimum_quanto_available,
+ is_optimum_quanto_version,
+ is_outlines_available,
+ is_peft_available,
+ is_peft_version,
+ is_pytorch_retinaface_available,
+ is_safetensors_available,
+ is_sageattention_available,
+ is_sageattention_version,
+ is_scipy_available,
+ is_sdnq_available,
+ is_sdnq_version,
+ is_sentencepiece_available,
+ is_tensorboard_available,
+ is_timm_available,
+ is_torch_available,
+ is_torch_mlu_available,
+ is_torch_neuronx_available,
+ is_torch_npu_available,
+ is_torch_version,
+ is_torch_xla_available,
+ is_torch_xla_version,
+ is_torchao_available,
+ is_torchao_version,
+ is_torchsde_available,
+ is_torchvision_available,
+ is_transformers_available,
+ is_transformers_version,
+ is_unidecode_available,
+ is_wandb_available,
+ is_xformers_available,
+ is_xformers_version,
+ requires_backends,
+)
+from .loading_utils import get_module_from_name, get_submodule_by_name, load_image, load_video
+from .logging import get_logger
+from .outputs import BaseOutput
+from .peft_utils import (
+ apply_lora_scale,
+ check_peft_version,
+ delete_adapter_layers,
+ get_adapter_name,
+ get_peft_kwargs,
+ recurse_remove_peft_layers,
+ scale_lora_layers,
+ set_adapter_layers,
+ set_weights_and_activate_adapters,
+ unscale_lora_layers,
+)
+from .pil_utils import PIL_INTERPOLATION, make_image_grid, numpy_to_pil, pt_to_pil
+from .remote_utils import remote_decode
+from .state_dict_utils import (
+ convert_all_state_dict_to_peft,
+ convert_sai_sd_control_lora_state_dict_to_peft,
+ convert_state_dict_to_diffusers,
+ convert_state_dict_to_kohya,
+ convert_state_dict_to_peft,
+ convert_unet_state_dict_to_peft,
+ state_dict_all_zero,
+)
+from .typing_utils import _get_detailed_type, _is_valid_type
+
+
+logger = get_logger(__name__)
+
+
+def check_min_version(min_version):
+ if version.parse(__version__) < version.parse(min_version):
+ if "dev" in min_version:
+ error_message = (
+ "This example requires a source install from HuggingFace diffusers (see "
+ "`https://huggingface.co/docs/diffusers/installation#install-from-source`),"
+ )
+ else:
+ error_message = f"This example requires a minimum version of {min_version},"
+ error_message += f" but the version found is {__version__}.\n"
+ raise ImportError(error_message)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/accelerate_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/accelerate_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..de6d0ff171862c4e38982a43d61db00975b1ce1e
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/accelerate_utils.py
@@ -0,0 +1,48 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Accelerate utilities: Utilities related to accelerate
+"""
+
+from packaging import version
+
+from .import_utils import is_accelerate_available
+
+
+if is_accelerate_available():
+ import accelerate
+
+
+def apply_forward_hook(method):
+ """
+ Decorator that applies a registered CpuOffload hook to an arbitrary function rather than `forward`. This is useful
+ for cases where a PyTorch module provides functions other than `forward` that should trigger a move to the
+ appropriate acceleration device. This is the case for `encode` and `decode` in [`AutoencoderKL`].
+
+ This decorator looks inside the internal `_hf_hook` property to find a registered offload hook.
+
+ :param method: The method to decorate. This method should be a method of a PyTorch module.
+ """
+ if not is_accelerate_available():
+ return method
+ accelerate_version = version.parse(accelerate.__version__).base_version
+ if version.parse(accelerate_version) < version.parse("0.17.0"):
+ return method
+
+ def wrapper(self, *args, **kwargs):
+ if hasattr(self, "_hf_hook") and hasattr(self._hf_hook, "pre_forward"):
+ self._hf_hook.pre_forward(self)
+ return method(self, *args, **kwargs)
+
+ return wrapper
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/constants.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..96e30d41641f2da3cd0b08b25e929ef180b43a04
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/constants.py
@@ -0,0 +1,85 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import importlib
+import os
+
+from huggingface_hub.constants import HF_HOME
+from packaging import version
+
+from ..dependency_versions_check import dep_version_check
+from .import_utils import ENV_VARS_TRUE_VALUES, is_peft_available, is_transformers_available
+
+
+MIN_PEFT_VERSION = "0.6.0"
+MIN_TRANSFORMERS_VERSION = "4.34.0"
+_CHECK_PEFT = os.environ.get("_CHECK_PEFT", "1") in ENV_VARS_TRUE_VALUES
+
+
+CONFIG_NAME = "config.json"
+WEIGHTS_NAME = "diffusion_pytorch_model.bin"
+WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.bin.index.json"
+ONNX_WEIGHTS_NAME = "model.onnx"
+SAFETENSORS_WEIGHTS_NAME = "diffusion_pytorch_model.safetensors"
+SAFE_WEIGHTS_INDEX_NAME = "diffusion_pytorch_model.safetensors.index.json"
+SAFETENSORS_FILE_EXTENSION = "safetensors"
+FLASHPACK_WEIGHTS_NAME = "model.flashpack"
+FLASHPACK_FILE_EXTENSION = "flashpack"
+GGUF_FILE_EXTENSION = "gguf"
+ONNX_EXTERNAL_WEIGHTS_NAME = "weights.pb"
+HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co")
+DIFFUSERS_DYNAMIC_MODULE_NAME = "diffusers_modules"
+HF_MODULES_CACHE = os.getenv("HF_MODULES_CACHE", os.path.join(HF_HOME, "modules"))
+DEPRECATED_REVISION_ARGS = ["fp16", "non-ema"]
+DIFFUSERS_REQUEST_TIMEOUT = 60
+DIFFUSERS_ATTN_BACKEND = os.getenv("DIFFUSERS_ATTN_BACKEND", "native")
+DIFFUSERS_ATTN_CHECKS = os.getenv("DIFFUSERS_ATTN_CHECKS", "0").upper() in ENV_VARS_TRUE_VALUES
+DEFAULT_HF_PARALLEL_LOADING_WORKERS = 8
+HF_ENABLE_PARALLEL_LOADING = os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES
+DIFFUSERS_DISABLE_REMOTE_CODE = os.getenv("DIFFUSERS_DISABLE_REMOTE_CODE", "false").upper() in ENV_VARS_TRUE_VALUES
+DIFFUSERS_SDNQ_TRANSFORMERS = os.getenv("DIFFUSERS_SDNQ_TRANSFORMERS", "false").upper() in ENV_VARS_TRUE_VALUES
+
+# Below should be `True` if the current version of `peft` and `transformers` are compatible with
+# PEFT backend. Will automatically fall back to PEFT backend if the correct versions of the libraries are
+# available.
+# For PEFT it is has to be greater than or equal to 0.6.0 and for transformers it has to be greater than or equal to 4.34.0.
+_required_peft_version = is_peft_available() and version.parse(
+ version.parse(importlib.metadata.version("peft")).base_version
+) >= version.parse(MIN_PEFT_VERSION)
+_required_transformers_version = is_transformers_available() and version.parse(
+ version.parse(importlib.metadata.version("transformers")).base_version
+) >= version.parse(MIN_TRANSFORMERS_VERSION)
+
+USE_PEFT_BACKEND = _required_peft_version and _required_transformers_version
+
+if USE_PEFT_BACKEND and _CHECK_PEFT:
+ dep_version_check("peft")
+
+
+DECODE_ENDPOINT_SD_V1 = "https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/"
+DECODE_ENDPOINT_SD_XL = "https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud/"
+DECODE_ENDPOINT_FLUX = "https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/"
+DECODE_ENDPOINT_HUNYUAN_VIDEO = "https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/"
+
+
+ENCODE_ENDPOINT_SD_V1 = "https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud/"
+ENCODE_ENDPOINT_SD_XL = "https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud/"
+ENCODE_ENDPOINT_FLUX = "https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud/"
+
+
+DIFFUSERS_LOAD_ID_FIELDS = [
+ "pretrained_model_name_or_path",
+ "subfolder",
+ "variant",
+ "revision",
+]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/deprecation_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/deprecation_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d3038d179059726d033494d49b384901d3bb1fe
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/deprecation_utils.py
@@ -0,0 +1,112 @@
+import inspect
+import warnings
+from typing import Any
+
+from packaging import version
+
+from ..utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+_TORCH_DTYPE_DEPRECATION_MESSAGE = "Please use `dtype` instead."
+
+
+def _resolve_dtype(dtype, torch_dtype):
+ """Resolve the deprecated `torch_dtype` argument into `dtype`."""
+ if torch_dtype is None:
+ return dtype
+ if dtype is not None:
+ raise ValueError(
+ "You have passed both `dtype` and `torch_dtype`. Please only pass `dtype`, `torch_dtype` is deprecated."
+ )
+ deprecate("torch_dtype", "1.0.0", _TORCH_DTYPE_DEPRECATION_MESSAGE)
+ return torch_dtype
+
+
+# Mapping for deprecated Transformers classes to their replacements
+# This is used to handle models that reference deprecated class names in their configs
+# Reference: https://github.com/huggingface/transformers/issues/40822
+# Format: {
+# "DeprecatedClassName": {
+# "new_class": "NewClassName",
+# "transformers_version": (">=", "5.0.0"), # (operation, version) tuple
+# }
+# }
+_TRANSFORMERS_CLASS_REMAPPING = {
+ "CLIPFeatureExtractor": {
+ "new_class": "CLIPImageProcessor",
+ "transformers_version": (">", "4.57.0"),
+ },
+}
+
+
+def _maybe_remap_transformers_class(class_name: str) -> str | None:
+ """
+ Check if a Transformers class should be remapped to a newer version.
+
+ Args:
+ class_name: The name of the class to check
+
+ Returns:
+ The new class name if remapping should occur, None otherwise
+ """
+ if class_name not in _TRANSFORMERS_CLASS_REMAPPING:
+ return None
+
+ from .import_utils import is_transformers_version
+
+ mapping = _TRANSFORMERS_CLASS_REMAPPING[class_name]
+ operation, required_version = mapping["transformers_version"]
+
+ # Only remap if the transformers version meets the requirement
+ if is_transformers_version(operation, required_version):
+ new_class = mapping["new_class"]
+ logger.warning(f"{class_name} appears to have been deprecated in transformers. Using {new_class} instead.")
+ return mapping["new_class"]
+
+ return None
+
+
+def deprecate(*args, take_from: dict | Any | None = None, standard_warn=True, stacklevel=2):
+ from .. import __version__
+
+ deprecated_kwargs = take_from
+ values = ()
+ if not isinstance(args[0], tuple):
+ args = (args,)
+
+ for attribute, version_name, message in args:
+ if version.parse(version.parse(__version__).base_version) >= version.parse(version_name):
+ raise ValueError(
+ f"The deprecation tuple {(attribute, version_name, message)} should be removed since diffusers'"
+ f" version {__version__} is >= {version_name}"
+ )
+
+ warning = None
+ if isinstance(deprecated_kwargs, dict) and attribute in deprecated_kwargs:
+ values += (deprecated_kwargs.pop(attribute),)
+ warning = f"The `{attribute}` argument is deprecated and will be removed in version {version_name}."
+ elif hasattr(deprecated_kwargs, attribute):
+ values += (getattr(deprecated_kwargs, attribute),)
+ warning = f"The `{attribute}` attribute is deprecated and will be removed in version {version_name}."
+ elif deprecated_kwargs is None:
+ warning = f"`{attribute}` is deprecated and will be removed in version {version_name}."
+
+ if warning is not None:
+ warning = warning + " " if standard_warn else ""
+ warnings.warn(warning + message, FutureWarning, stacklevel=stacklevel)
+
+ if isinstance(deprecated_kwargs, dict) and len(deprecated_kwargs) > 0:
+ call_frame = inspect.getouterframes(inspect.currentframe())[1]
+ filename = call_frame.filename
+ line_number = call_frame.lineno
+ function = call_frame.function
+ key, value = next(iter(deprecated_kwargs.items()))
+ raise TypeError(f"{function} in {filename} line {line_number - 1} got an unexpected keyword argument `{key}`")
+
+ if len(values) == 0:
+ return
+ elif len(values) == 1:
+ return values[0]
+ return values
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/distributed_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/distributed_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..1446e6eac6a46d6636ae6442c33ffaa70ed053e2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/distributed_utils.py
@@ -0,0 +1,36 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+try:
+ import torch
+except ImportError:
+ torch = None
+
+
+def is_torch_dist_rank_zero() -> bool:
+ if torch is None:
+ return True
+
+ dist_module = getattr(torch, "distributed", None)
+ if dist_module is None or not dist_module.is_available():
+ return True
+
+ if not dist_module.is_initialized():
+ return True
+
+ try:
+ return dist_module.get_rank() == 0
+ except (RuntimeError, ValueError):
+ return True
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/doc_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/doc_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..bbf9b585cd07ae26f2d4c507309dfab571a12807
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/doc_utils.py
@@ -0,0 +1,39 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Doc utilities: Utilities related to documentation
+"""
+
+import re
+
+
+def replace_example_docstring(example_docstring):
+ def docstring_decorator(fn):
+ func_doc = fn.__doc__
+ lines = func_doc.split("\n")
+ i = 0
+ while i < len(lines) and re.search(r"^\s*Examples?:\s*$", lines[i]) is None:
+ i += 1
+ if i < len(lines):
+ lines[i] = example_docstring
+ func_doc = "\n".join(lines)
+ else:
+ raise ValueError(
+ f"The function {fn} should have an empty 'Examples:' in its docstring as placeholder, "
+ f"current docstring is:\n{func_doc}"
+ )
+ fn.__doc__ = func_doc
+ return fn
+
+ return docstring_decorator
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_auto_round_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_auto_round_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..6bf983e9581ef12244c9129ca1154e4b0134a647
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_auto_round_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class AutoRoundConfig(metaclass=DummyObject):
+ _backends = ["auto_round"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["auto_round"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["auto_round"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["auto_round"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_bitsandbytes_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_bitsandbytes_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6feaea86ce30832a225d8779cc4c95f24a1f0f8
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_bitsandbytes_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class BitsAndBytesConfig(metaclass=DummyObject):
+ _backends = ["bitsandbytes"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["bitsandbytes"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["bitsandbytes"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["bitsandbytes"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_gguf_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_gguf_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b10a2826c0e04828df95fe253a5dd8ddd780659
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_gguf_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class GGUFQuantizationConfig(metaclass=DummyObject):
+ _backends = ["gguf"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["gguf"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["gguf"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["gguf"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_note_seq_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_note_seq_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..a853a306b98a7338ae343ec1fb22c66a12bc9438
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_note_seq_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class MidiProcessor(metaclass=DummyObject):
+ _backends = ["note_seq"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["note_seq"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["note_seq"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["note_seq"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nunchaku_lite_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nunchaku_lite_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9d7956aebe775080497a945c158397241c87bd6
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nunchaku_lite_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class NunchakuLiteQuantizationConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nvidia_modelopt_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nvidia_modelopt_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..beddf9be31719069d613da2b2a47cbb76f5a1f30
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_nvidia_modelopt_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class NVIDIAModelOptConfig(metaclass=DummyObject):
+ _backends = ["nvidia_modelopt"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["nvidia_modelopt"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["nvidia_modelopt"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["nvidia_modelopt"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_onnx_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_onnx_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed42bcd1b9c20d811107ac621ea488019e691696
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_onnx_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class OnnxRuntimeModel(metaclass=DummyObject):
+ _backends = ["onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["onnx"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_optimum_quanto_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_optimum_quanto_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c6f06fea338194e64c2c101eeb8f6833cf6f46d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_optimum_quanto_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class QuantoConfig(metaclass=DummyObject):
+ _backends = ["optimum_quanto"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["optimum_quanto"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["optimum_quanto"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["optimum_quanto"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_pt_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_pt_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0109c1a990a2c178117ab24549aaff281676cce
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_pt_objects.py
@@ -0,0 +1,3662 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class AdaptiveProjectedGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AdaptiveProjectedMixGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BaseGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ClassifierFreeGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ClassifierFreeZeroStarGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FrequencyDecoupledGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PerturbedAttentionGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SkipLayerGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SmoothedEnergyGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class TangentialClassifierFreeGuidance(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FasterCacheConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FirstBlockCacheConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HookRegistry(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LayerSkipConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MagCacheConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PyramidAttentionBroadcastConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SmoothedEnergyGuidanceConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class TaylorSeerCacheConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class TextKVCacheConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+def apply_faster_cache(*args, **kwargs):
+ requires_backends(apply_faster_cache, ["torch"])
+
+
+def apply_first_block_cache(*args, **kwargs):
+ requires_backends(apply_first_block_cache, ["torch"])
+
+
+def apply_layer_skip(*args, **kwargs):
+ requires_backends(apply_layer_skip, ["torch"])
+
+
+def apply_mag_cache(*args, **kwargs):
+ requires_backends(apply_mag_cache, ["torch"])
+
+
+def apply_pyramid_attention_broadcast(*args, **kwargs):
+ requires_backends(apply_pyramid_attention_broadcast, ["torch"])
+
+
+def apply_taylorseer_cache(*args, **kwargs):
+ requires_backends(apply_taylorseer_cache, ["torch"])
+
+
+def apply_text_kv_cache(*args, **kwargs):
+ requires_backends(apply_text_kv_cache, ["torch"])
+
+
+class InpaintProcessor(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class IPAdapterMaskProcessor(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PixArtImageProcessor(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class VaeImageProcessor(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class VaeImageProcessorLDM3D(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AceStepTransformer1DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AllegroTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AnimaTextConditioner(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AnyFlowFARTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AnyFlowTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AsymmetricAutoencoderKL(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AttentionBackendName(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AuraFlowTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderDC(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKL(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLAllegro(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLCogVideoX(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLCosmos(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLFlux2(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLHunyuanImage(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLHunyuanImageRefiner(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLHunyuanVideo(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLHunyuanVideo15(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLKVAE(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLKVAEVideo(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLLTX2Audio(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLLTX2Video(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLLTXVideo(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLMagvit(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLMiniMaxH3(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLMiniMaxH3Audio(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLMochi(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLQwenImage(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLTemporalDecoder(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderKLWan(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderOobleck(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderRAE(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderTiny(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoencoderVidTok(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BriaFiboTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BriaTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CacheMixin(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ChromaTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ChronoEditTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CogVideoXTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CogView3PlusTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CogView4Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConsisIDTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConsistencyDecoderVAE(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ContextParallelConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ControlNetUnionModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ControlNetXSAdapter(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Cosmos3AVAEAudioTokenizer(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Cosmos3OmniTransformer(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CosmosControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CosmosTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiTTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DreamLiteTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DreamLiteUNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EasyAnimateTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ErnieImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Flux2Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FluxControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FluxMultiControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FluxTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class GlmImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HeliosTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HiDreamImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanDiT2DControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanDiT2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanDiT2DMultiControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanVideo15Transformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanVideoFramepackTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HunyuanVideoTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class I2VGenXLUNet(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Ideogram4Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class JoyImageEditPlusTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class JoyImageEditTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Kandinsky3UNet(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Kandinsky5Transformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Krea2Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LatteTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LongCatAudioDiTTransformer(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LongCatAudioDiTVae(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LongCatImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LTX2VideoTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LTXVideoTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Lumina2Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LuminaNextDiT2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MiniMaxH3Transformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MochiTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ModelMixin(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MotifVideoTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MotionAdapter(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MultiAdapter(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MultiControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class NucleusMoEImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class OmniGenTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class OvisImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ParallelConfig(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PixArtTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PriorTransformer(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PRXTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class QwenImageControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class QwenImageMultiControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class QwenImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SanaControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SanaTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SanaVideoTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SD3ControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SD3MultiControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SD3Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SkyReelsV2Transformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SparseControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class StableAudioDiTModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class T2IAdapter(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class T5FilmDecoder(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class Transformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class TransformerTemporalModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNet1DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNet2DConditionModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNet2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNet3DConditionModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNetControlNetXSModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNetMotionModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UNetSpatioTemporalConditionModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UVit2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class VQModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class WanAnimateTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class WanTransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class WanVACETransformer3DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ZImageControlNetModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ZImageTransformer2DModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+def attention_backend(*args, **kwargs):
+ requires_backends(attention_backend, ["torch"])
+
+
+class AutoPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ComponentsManager(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ComponentSpec(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConditionalPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConfigSpec(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class InputParam(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LoopSequentialPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ModularPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class OutputParam(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SequentialPipelineBlocks(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+def get_constant_schedule(*args, **kwargs):
+ requires_backends(get_constant_schedule, ["torch"])
+
+
+def get_constant_schedule_with_warmup(*args, **kwargs):
+ requires_backends(get_constant_schedule_with_warmup, ["torch"])
+
+
+def get_cosine_schedule_with_warmup(*args, **kwargs):
+ requires_backends(get_cosine_schedule_with_warmup, ["torch"])
+
+
+def get_cosine_with_hard_restarts_schedule_with_warmup(*args, **kwargs):
+ requires_backends(get_cosine_with_hard_restarts_schedule_with_warmup, ["torch"])
+
+
+def get_linear_schedule_with_warmup(*args, **kwargs):
+ requires_backends(get_linear_schedule_with_warmup, ["torch"])
+
+
+def get_polynomial_decay_schedule_with_warmup(*args, **kwargs):
+ requires_backends(get_polynomial_decay_schedule_with_warmup, ["torch"])
+
+
+def get_scheduler(*args, **kwargs):
+ requires_backends(get_scheduler, ["torch"])
+
+
+class AudioPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoPipelineForImage2Image(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoPipelineForInpainting(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoPipelineForText2Audio(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AutoPipelineForText2Image(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BlipDiffusionControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BlipDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CLIPImageProjection(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ConsistencyModelPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DanceDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDIMPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDPMPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiTPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ImagePipelineOutput(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class KarrasVePipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LDMPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LDMSuperResolutionPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PNDMPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class RePaintPipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ScoreSdeVePipeline(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class StableDiffusionMixin(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiffusersQuantizer(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class AmusedScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BlockRefinementScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class BlockRefinementSchedulerOutput(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CMStochasticIterativeScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CogVideoXDDIMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class CogVideoXDPMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDIMInverseScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDIMParallelScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDIMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDPMParallelScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDPMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DDPMWuerstchenScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DEISMultistepScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiscreteDDIMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DiscreteDDIMSchedulerOutput(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DPMSolverMultistepInverseScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DPMSolverMultistepScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class DPMSolverSinglestepScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EDMDPMSolverMultistepScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EDMEulerScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EntropyBoundScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EntropyBoundSchedulerOutput(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EulerAncestralDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EulerDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FlowMapEulerDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FlowMatchEulerDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FlowMatchHeunDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class FlowMatchLCMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HeliosDMDScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HeliosScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class HeunDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class IPNDMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class KarrasVeScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class KDPM2AncestralDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class KDPM2DiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LCMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class LTXEulerAncestralRFScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class MiniMaxH3Scheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class PNDMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class RePaintScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SASolverScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SchedulerMixin(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class SCMScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class ScoreSdeVeScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class TCDScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UnCLIPScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class UniPCMultistepScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class VQDiffusionScheduler(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class EMAModel(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+
+class VideoProcessor(metaclass=DummyObject):
+ _backends = ["torch"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_sdnq_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_sdnq_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5808d4141f73bca3c0927e037b6f71d760aeb4c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_sdnq_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class SDNQConfig(metaclass=DummyObject):
+ _backends = ["sdnq"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["sdnq"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["sdnq"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["sdnq"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_librosa_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_librosa_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2c66c223a2c607c8c7bc738b174c7f8d156d2f2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_librosa_objects.py
@@ -0,0 +1,32 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class AudioDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "librosa"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "librosa"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "librosa"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "librosa"])
+
+
+class Mel(metaclass=DummyObject):
+ _backends = ["torch", "librosa"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "librosa"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "librosa"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "librosa"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_scipy_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_scipy_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbd96af6170b33428f8502625fbd41dbdf097f91
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_scipy_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class LMSDiscreteScheduler(metaclass=DummyObject):
+ _backends = ["torch", "scipy"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "scipy"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "scipy"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "scipy"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_torchsde_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_torchsde_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d4e9fc044b4727a04800032abea7111b98116e2
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_torchsde_objects.py
@@ -0,0 +1,32 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class CosineDPMSolverMultistepScheduler(metaclass=DummyObject):
+ _backends = ["torch", "torchsde"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "torchsde"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "torchsde"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "torchsde"])
+
+
+class DPMSolverSDEScheduler(metaclass=DummyObject):
+ _backends = ["torch", "torchsde"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "torchsde"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "torchsde"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "torchsde"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..6cb7ef13273229ff68917afbaf50b2e268005de4
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py
@@ -0,0 +1,92 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class OnnxStableDiffusionImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+
+class OnnxStableDiffusionInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+
+class OnnxStableDiffusionInpaintPipelineLegacy(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+
+class OnnxStableDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+
+class OnnxStableDiffusionUpscalePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+
+class StableDiffusionOnnxPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "onnx"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "onnx"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..0582d7d871aaeb7c526837fee5aaa9728c50dc41
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_opencv_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class ConsisIDPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "opencv"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "opencv"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "opencv"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "opencv"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..fce7f61280111113c5c5def5b28c41b22f14fafa
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py
@@ -0,0 +1,47 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class KolorsImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "sentencepiece"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
+
+
+class KolorsPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "sentencepiece"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
+
+
+class KolorsPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers", "sentencepiece"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers", "sentencepiece"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef23667a95b6c2356bc58adcb8b9d308a92d3f24
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torch_and_transformers_objects.py
@@ -0,0 +1,5222 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class AnimaAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimaModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos3DistilledBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos3DistilledModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos3OmniBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos3OmniModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ErnieImageAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ErnieImageModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2AutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinBaseAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinBaseModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxKontextAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxKontextModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPyramidAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPyramidDistilledAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPyramidDistilledModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPyramidModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideo15AutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideo15ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Ideogram4AutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Ideogram4ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Krea2AutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Krea2ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Krea2TurboAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Krea2TurboModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MiniMaxH3Blocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MiniMaxH3ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditPlusAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditPlusModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageLayeredAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageLayeredModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3AutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Wan22Blocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Wan22Image2VideoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Wan22Image2VideoModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Wan22ModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanImage2VideoAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanImage2VideoModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageAutoBlocks(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageModularPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AceStepAudioTokenDetokenizer(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AceStepAudioTokenizer(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AceStepConditionEncoder(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AceStepPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AllegroPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AltDiffusionImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AltDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AmusedImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AmusedInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AmusedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffSDXLPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffSparseControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffVideoToVideoControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnimateDiffVideoToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnyFlowFARPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AnyFlowPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AudioLDM2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AudioLDM2ProjectionModel(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AudioLDM2UNet2DConditionModel(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AudioLDMPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class AuraFlowPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class BriaFiboEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class BriaFiboPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class BriaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ChromaImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ChromaInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ChromaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ChronoEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CLIPImageProjection(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogVideoXFunControlPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogVideoXImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogVideoXPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogVideoXVideoToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogView3PlusPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogView4ControlPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CogView4Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ConsisIDPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos2_5_PredictBasePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos2_5_TransferPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos2TextToImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos2VideoToWorldPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Cosmos3OmniPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CosmosActionCondition(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CosmosTextToWorldPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CosmosVideoToWorldPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class CycleDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class DiffusionGemmaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class DiffusionGemmaPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class DreamLiteMobilePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class DreamLitePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class DreamLitePipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class EasyAnimateControlPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class EasyAnimateInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class EasyAnimatePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ErnieImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinKVPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2KleinPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlNetImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlNetInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxControlPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxFillPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxKontextInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxKontextPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class FluxPriorReduxPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class GlmImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HeliosPyramidPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HiDreamImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanDiTControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanDiTPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanDiTPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanImageRefinerPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanSkyreelsImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideo15ImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideo15Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideoFramepackPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideoImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class HunyuanVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class I2VGenXLPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Ideogram4Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Ideogram4PromptEnhancerHead(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFImg2ImgSuperResolutionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFInpaintingPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFInpaintingSuperResolutionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class IFSuperResolutionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ImageTextPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class JoyImageEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class JoyImageEditPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class JoyImageEditPlusPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class JoyImageEditPlusPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky3Img2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky3Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky5I2IPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky5I2VPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky5T2IPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Kandinsky5T2VPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyImg2ImgCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyInpaintCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyPriorPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22CombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22ControlnetImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22ControlnetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22Img2ImgCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22Img2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22InpaintCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22InpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22PriorEmb2EmbPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class KandinskyV22PriorPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Krea2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LatentConsistencyModelImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LatentConsistencyModelPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LattePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LDMTextToImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LEditsPPPipelineStableDiffusion(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LEditsPPPipelineStableDiffusionXL(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LLaDA2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LLaDA2PipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LongCatAudioDiTPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LongCatImageEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LongCatImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2ConditionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2HDRPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2ImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2InContextPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2LatentUpsamplePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTX2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXConditionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXI2VLongMultiPromptPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXLatentUpsamplePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LTXPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LucyEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Lumina2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class Lumina2Text2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LuminaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class LuminaText2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MarigoldDepthPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MarigoldIntrinsicsPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MarigoldNormalsPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MochiPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MotifVideoImage2VideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MotifVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MotifVideoPipelineOutput(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class MusicLDMPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class NucleusMoEImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class OmniGenPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class OvisImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PaintByExamplePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PIAPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PixArtAlphaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PixArtSigmaPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PixArtSigmaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PRXPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class PRXPixelPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageControlNetInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageEditPlusPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImageLayeredPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class QwenImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ReduxImageEncoder(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaSprintImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaSprintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SanaVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SemanticStableDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ShapEImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ShapEPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SkyReelsV2DiffusionForcingImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SkyReelsV2DiffusionForcingPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SkyReelsV2DiffusionForcingVideoToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SkyReelsV2ImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class SkyReelsV2Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableAudioPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableAudioProjectionModel(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableCascadeCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableCascadeDecoderPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableCascadePriorPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3ControlNetInpaintingPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3ControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3Img2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3InpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3PAGImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3PAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusion3Pipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionAdapterPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionAttendAndExcitePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetPAGInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionControlNetXSPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionDepth2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionDiffEditPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionGLIGENPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionGLIGENTextImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionImageVariationPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionInpaintPipelineLegacy(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionInstructPix2PixPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionLatentUpscalePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionLDM3DPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionModelEditingPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPAGImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPAGInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPanoramaPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionParadigmsPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPipelineSafe(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionPix2PixZeroPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionSAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionUpscalePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLAdapterPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetPAGImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetUnionImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetUnionInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetUnionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLControlNetXSPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLInstructPix2PixPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLPAGImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLPAGInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLPAGPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableDiffusionXLPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableUnCLIPImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableUnCLIPPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class StableVideoDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class TextToVideoSDPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class TextToVideoZeroPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class TextToVideoZeroSDXLPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class UnCLIPImageVariationPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class UnCLIPPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class UniDiffuserModel(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class UniDiffuserPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class UniDiffuserTextDecoder(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VersatileDiffusionDualGuidedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VersatileDiffusionImageVariationPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VersatileDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VersatileDiffusionTextToImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VideoToVideoSDPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VisualClozeGenerationPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VisualClozePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class VQDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanAnimatePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanImageToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanVACEPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WanVideoToVideoPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WuerstchenCombinedPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WuerstchenDecoderPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class WuerstchenPriorPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageControlNetInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageControlNetPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageImg2ImgPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageInpaintPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImageOmniPipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+
+class ZImagePipeline(metaclass=DummyObject):
+ _backends = ["torch", "transformers"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torch", "transformers"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torch", "transformers"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torchao_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torchao_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2cdc1916db548acf0deabe78df7fd696dcee7c0
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_torchao_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class TorchAoConfig(metaclass=DummyObject):
+ _backends = ["torchao"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["torchao"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["torchao"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["torchao"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py
new file mode 100644
index 0000000000000000000000000000000000000000..34ef89f8c06182b0bd845ad77384306fb15641e7
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py
@@ -0,0 +1,17 @@
+# This file is autogenerated by the command `make fix-copies`, do not edit.
+from ..utils import DummyObject, requires_backends
+
+
+class SpectrogramDiffusionPipeline(metaclass=DummyObject):
+ _backends = ["transformers", "torch", "note_seq"]
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, ["transformers", "torch", "note_seq"])
+
+ @classmethod
+ def from_config(cls, *args, **kwargs):
+ requires_backends(cls, ["transformers", "torch", "note_seq"])
+
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ requires_backends(cls, ["transformers", "torch", "note_seq"])
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dynamic_modules_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dynamic_modules_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb98783faa7be2509cfa799dd53435c0c0ec28af
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/dynamic_modules_utils.py
@@ -0,0 +1,544 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities to dynamically load objects from the Hub."""
+
+import importlib
+import inspect
+import json
+import os
+import re
+import shutil
+import sys
+import threading
+from pathlib import Path
+from types import ModuleType
+from urllib import request
+
+from huggingface_hub import hf_hub_download, model_info
+from huggingface_hub.utils import RevisionNotFoundError, validate_hf_hub_args
+from packaging import version
+
+from .. import __version__
+from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging
+from .constants import DIFFUSERS_DISABLE_REMOTE_CODE
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+# See https://huggingface.co/datasets/diffusers/community-pipelines-mirror
+COMMUNITY_PIPELINES_MIRROR_ID = "diffusers/community-pipelines-mirror"
+TIME_OUT_REMOTE_CODE = int(os.getenv("DIFFUSERS_TIMEOUT_REMOTE_CODE", 15))
+_HF_REMOTE_CODE_LOCK = threading.Lock()
+
+
+def get_diffusers_versions():
+ url = "https://pypi.org/pypi/diffusers/json"
+ releases = json.loads(request.urlopen(url).read())["releases"].keys()
+ return sorted(releases, key=lambda x: version.Version(x))
+
+
+def init_hf_modules():
+ """
+ Creates the cache directory for modules with an init, and adds it to the Python path.
+ """
+ # This function has already been executed if HF_MODULES_CACHE already is in the Python path.
+ if HF_MODULES_CACHE in sys.path:
+ return
+
+ sys.path.append(HF_MODULES_CACHE)
+ os.makedirs(HF_MODULES_CACHE, exist_ok=True)
+ init_path = Path(HF_MODULES_CACHE) / "__init__.py"
+ if not init_path.exists():
+ init_path.touch()
+
+
+def create_dynamic_module(name: str | os.PathLike):
+ """
+ Creates a dynamic module in the cache directory for modules.
+ """
+ init_hf_modules()
+ dynamic_module_path = Path(HF_MODULES_CACHE) / name
+ # If the parent module does not exist yet, recursively create it.
+ if not dynamic_module_path.parent.exists():
+ create_dynamic_module(dynamic_module_path.parent)
+ os.makedirs(dynamic_module_path, exist_ok=True)
+ init_path = dynamic_module_path / "__init__.py"
+ if not init_path.exists():
+ init_path.touch()
+
+
+def get_relative_imports(module_file):
+ """
+ Get the list of modules that are relatively imported in a module file.
+
+ Args:
+ module_file (`str` or `os.PathLike`): The module file to inspect.
+ """
+ with open(module_file, "r", encoding="utf-8") as f:
+ content = f.read()
+
+ # Imports of the form `import .xxx`
+ relative_imports = re.findall(r"^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE)
+ # Imports of the form `from .xxx import yyy`
+ relative_imports += re.findall(r"^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE)
+ # Unique-ify
+ return list(set(relative_imports))
+
+
+def get_relative_import_files(module_file):
+ """
+ Get the list of all files that are needed for a given module. Note that this function recurses through the relative
+ imports (if a imports b and b imports c, it will return module files for b and c).
+
+ Args:
+ module_file (`str` or `os.PathLike`): The module file to inspect.
+ """
+ no_change = False
+ files_to_check = [module_file]
+ all_relative_imports = []
+
+ # Let's recurse through all relative imports
+ while not no_change:
+ new_imports = []
+ for f in files_to_check:
+ new_imports.extend(get_relative_imports(f))
+
+ module_path = Path(module_file).parent
+ new_import_files = [str(module_path / m) for m in new_imports]
+ new_import_files = [f for f in new_import_files if f not in all_relative_imports]
+ files_to_check = [f"{f}.py" for f in new_import_files]
+
+ no_change = len(new_import_files) == 0
+ all_relative_imports.extend(files_to_check)
+
+ return all_relative_imports
+
+
+def check_imports(filename):
+ """
+ Check if the current Python environment contains all the libraries that are imported in a file.
+ """
+ with open(filename, "r", encoding="utf-8") as f:
+ content = f.read()
+
+ # Imports of the form `import xxx`
+ imports = re.findall(r"^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE)
+ # Imports of the form `from xxx import yyy`
+ imports += re.findall(r"^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE)
+ # Only keep the top-level module
+ imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")]
+
+ # Unique-ify and test we got them all
+ imports = list(set(imports))
+ missing_packages = []
+ for imp in imports:
+ try:
+ importlib.import_module(imp)
+ except ImportError:
+ missing_packages.append(imp)
+
+ if len(missing_packages) > 0:
+ logger.warning(
+ "This modeling file might require the following packages that were not found in your environment: "
+ f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`"
+ )
+
+ return get_relative_imports(filename)
+
+
+def resolve_trust_remote_code(trust_remote_code, model_name, has_remote_code):
+ trust_remote_code = trust_remote_code and not DIFFUSERS_DISABLE_REMOTE_CODE
+ if DIFFUSERS_DISABLE_REMOTE_CODE:
+ logger.warning(
+ "Downloading remote code is disabled globally via the DIFFUSERS_DISABLE_REMOTE_CODE environment variable. Ignoring `trust_remote_code`."
+ )
+
+ if has_remote_code and not trust_remote_code:
+ error_msg = f"The repository for {model_name} contains custom code. "
+ error_msg += (
+ "Downloading remote code is disabled globally via the DIFFUSERS_DISABLE_REMOTE_CODE environment variable."
+ if DIFFUSERS_DISABLE_REMOTE_CODE
+ else "Pass `trust_remote_code=True` to allow loading remote code modules."
+ )
+ raise ValueError(error_msg)
+
+ elif has_remote_code and trust_remote_code:
+ logger.warning(
+ f"`trust_remote_code` is enabled. Downloading code from {model_name}. Please ensure you trust the contents of this repository"
+ )
+
+ return trust_remote_code
+
+
+def get_class_in_module(class_name, module_path, force_reload=False):
+ """
+ Import a module on the cache directory for modules and extract a class from it.
+ """
+ name = os.path.normpath(module_path)
+ if name.endswith(".py"):
+ name = name[:-3]
+ name = name.replace(os.path.sep, ".")
+ module_file: Path = Path(HF_MODULES_CACHE) / module_path
+
+ with _HF_REMOTE_CODE_LOCK:
+ if force_reload:
+ sys.modules.pop(name, None)
+ importlib.invalidate_caches()
+ cached_module: ModuleType | None = sys.modules.get(name)
+ module_spec = importlib.util.spec_from_file_location(name, location=module_file)
+
+ module: ModuleType
+ if cached_module is None:
+ module = importlib.util.module_from_spec(module_spec)
+ # insert it into sys.modules before any loading begins
+ sys.modules[name] = module
+ else:
+ module = cached_module
+
+ module_spec.loader.exec_module(module)
+
+ if class_name is None:
+ return find_pipeline_class(module)
+
+ return getattr(module, class_name)
+
+
+def find_pipeline_class(loaded_module):
+ """
+ Retrieve pipeline class that inherits from `DiffusionPipeline`. Note that there has to be exactly one class
+ inheriting from `DiffusionPipeline`.
+ """
+ from ..pipelines import DiffusionPipeline
+
+ cls_members = dict(inspect.getmembers(loaded_module, inspect.isclass))
+
+ pipeline_class = None
+ for cls_name, cls in cls_members.items():
+ if (
+ cls_name != DiffusionPipeline.__name__
+ and issubclass(cls, DiffusionPipeline)
+ and cls.__module__.split(".")[0] != "diffusers"
+ ):
+ if pipeline_class is not None:
+ raise ValueError(
+ f"Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:"
+ f" {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in"
+ f" {loaded_module}."
+ )
+ pipeline_class = cls
+
+ return pipeline_class
+
+
+@validate_hf_hub_args
+def get_cached_module_file(
+ pretrained_model_name_or_path: str | os.PathLike,
+ module_file: str,
+ subfolder: str | None = None,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ local_dir: str | None = None,
+ trust_remote_code: bool = False,
+):
+ """
+ Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached
+ Transformers module.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
+ under a user or organization name, like `dbmdz/bert-base-german-cased`.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ module_file (`str`):
+ The name of the module file containing the class to look for.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `transformers-cli login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom pipelines and components defined on the Hub in their own files. This
+ option should only be set to `True` for repositories you trust and in which you have read the code, as it
+ will execute code present on the Hub on your local machine.
+
+ > [!TIP] > You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or
+ [gated > models](https://huggingface.co/docs/hub/models-gated#gated-models).
+
+ Returns:
+ `str`: The path to the module inside the cache.
+ """
+ # Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file.
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+
+ if DIFFUSERS_DISABLE_REMOTE_CODE:
+ raise ValueError(
+ "Downloading remote code is disabled globally via the DIFFUSERS_DISABLE_REMOTE_CODE environment variable."
+ )
+
+ if subfolder is not None:
+ module_file_or_url = os.path.join(pretrained_model_name_or_path, subfolder, module_file)
+ else:
+ module_file_or_url = os.path.join(pretrained_model_name_or_path, module_file)
+
+ is_local_file = os.path.isfile(module_file_or_url)
+ is_community_pipeline = not is_local_file and pretrained_model_name_or_path.count("/") == 0
+
+ if is_local_file:
+ resolved_module_file = module_file_or_url
+ submodule = "local"
+ if not trust_remote_code:
+ raise ValueError(
+ f"The directory {pretrained_model_name_or_path} contains custom code in {module_file} which must be executed to correctly "
+ f"load the model. You can inspect the file content at {module_file_or_url}.\n"
+ f"Pass `trust_remote_code=True` to allow loading remote code modules."
+ )
+ elif is_community_pipeline:
+ available_versions = get_diffusers_versions()
+ # cut ".dev0"
+ latest_version = "v" + ".".join(__version__.split(".")[:3])
+
+ # retrieve github version that matches
+ if revision is None:
+ revision = latest_version if latest_version[1:] in available_versions else "main"
+ logger.info(f"Defaulting to latest_version: {revision}.")
+ elif revision in available_versions:
+ revision = f"v{revision}"
+ elif revision == "main":
+ revision = revision
+ else:
+ raise ValueError(
+ f"`custom_revision`: {revision} does not exist. Please make sure to choose one of"
+ f" {', '.join(available_versions + ['main'])}."
+ )
+
+ try:
+ resolved_module_file = hf_hub_download(
+ repo_id=COMMUNITY_PIPELINES_MIRROR_ID,
+ repo_type="dataset",
+ filename=f"{revision}/{pretrained_model_name_or_path}.py",
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ local_dir=local_dir,
+ )
+ submodule = "git"
+ module_file = pretrained_model_name_or_path + ".py"
+ except RevisionNotFoundError as e:
+ raise EnvironmentError(
+ f"Revision '{revision}' not found in the community pipelines mirror. Check available revisions on"
+ " https://huggingface.co/datasets/diffusers/community-pipelines-mirror/tree/main."
+ " If you don't find the revision you are looking for, please open an issue on https://github.com/huggingface/diffusers/issues."
+ ) from e
+ except EnvironmentError:
+ logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")
+ raise
+ else:
+ if not trust_remote_code:
+ raise ValueError(
+ f"The repository for {pretrained_model_name_or_path} contains custom code in {module_file} which must be executed to correctly "
+ f"load the model. You can inspect the repository content at https://hf.co/{pretrained_model_name_or_path}/blob/main/{module_file}.\n"
+ f"Pass `trust_remote_code=True` to allow loading remote code modules."
+ )
+ try:
+ # Load from URL or cache if already cached
+ resolved_module_file = hf_hub_download(
+ pretrained_model_name_or_path,
+ module_file,
+ subfolder=subfolder,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ local_dir=local_dir,
+ revision=revision,
+ token=token,
+ )
+ submodule = os.path.join("local", "--".join(pretrained_model_name_or_path.split("/")))
+ except EnvironmentError:
+ logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.")
+ raise
+
+ # Check we have all the requirements in our environment
+ modules_needed = check_imports(resolved_module_file)
+
+ # Now we move the module inside our cached dynamic modules.
+ full_submodule = DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
+ create_dynamic_module(full_submodule)
+ submodule_path = Path(HF_MODULES_CACHE) / full_submodule
+ if submodule == "local" or submodule == "git":
+ # We always copy local files (we could hash the file to see if there was a change, and give them the name of
+ # that hash, to only copy when there is a modification but it seems overkill for now).
+ # The only reason we do the copy is to avoid putting too many folders in sys.path.
+ shutil.copyfile(resolved_module_file, submodule_path / module_file)
+ for module_needed in modules_needed:
+ if len(module_needed.split(".")) == 2:
+ module_needed = "/".join(module_needed.split("."))
+ module_folder = module_needed.split("/")[0]
+ if not os.path.exists(submodule_path / module_folder):
+ os.makedirs(submodule_path / module_folder)
+ module_needed = f"{module_needed}.py"
+ if subfolder is not None:
+ source_path = os.path.join(pretrained_model_name_or_path, subfolder, module_needed)
+ else:
+ source_path = os.path.join(pretrained_model_name_or_path, module_needed)
+ shutil.copyfile(source_path, submodule_path / module_needed)
+ else:
+ # Get the commit hash
+ # TODO: we will get this info in the etag soon, so retrieve it from there and not here.
+ commit_hash = model_info(pretrained_model_name_or_path, revision=revision, token=token).sha
+
+ # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the
+ # benefit of versioning.
+ submodule_path = submodule_path / commit_hash
+ full_submodule = full_submodule + os.path.sep + commit_hash
+ create_dynamic_module(full_submodule)
+
+ if not (submodule_path / module_file).exists():
+ if len(module_file.split("/")) == 2:
+ module_folder = module_file.split("/")[0]
+ if not os.path.exists(submodule_path / module_folder):
+ os.makedirs(submodule_path / module_folder)
+ shutil.copyfile(resolved_module_file, submodule_path / module_file)
+
+ # Make sure we also have every file with relative
+ for module_needed in modules_needed:
+ if len(module_needed.split(".")) == 2:
+ module_needed = "/".join(module_needed.split("."))
+ if not (submodule_path / module_needed).exists():
+ get_cached_module_file(
+ pretrained_model_name_or_path,
+ f"{module_needed}.py",
+ subfolder=subfolder,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ local_dir=local_dir,
+ trust_remote_code=trust_remote_code,
+ )
+ return os.path.join(full_submodule, module_file)
+
+
+@validate_hf_hub_args
+def get_class_from_dynamic_module(
+ pretrained_model_name_or_path: str | os.PathLike,
+ module_file: str,
+ subfolder: str | None = None,
+ class_name: str | None = None,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ local_dir: str | None = None,
+ trust_remote_code: bool = False,
+):
+ """
+ Extracts a class from a module file, present in the local folder or repository of a model.
+
+ > [!WARNING] > Calling this function will execute the code in the module file found locally or downloaded from the
+ Hub. It should > therefore only be called on trusted repos.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co. Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced
+ under a user or organization name, like `dbmdz/bert-base-german-cased`.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ module_file (`str`):
+ The name of the module file containing the class to look for.
+ class_name (`str`):
+ The name of the class to import in the module.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or `bool`, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `transformers-cli login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom pipelines and components defined on the Hub in their own files. This
+ option should only be set to `True` for repositories you trust and in which you have read the code, as it
+ will execute code present on the Hub on your local machine.
+
+ > [!TIP] > You may pass a token in `token` if you are not logged in (`hf auth login`) and want to use private or
+ [gated > models](https://huggingface.co/docs/hub/models-gated#gated-models).
+
+ Returns:
+ `type`: The class, dynamically imported from the module.
+
+ Examples:
+
+ ```python
+ # Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this
+ # module.
+ cls = get_class_from_dynamic_module("sgugger/my-bert-model", "modeling.py", "MyBertModel")
+ ```"""
+ # And lastly we get the class inside our newly created module
+ final_module = get_cached_module_file(
+ pretrained_model_name_or_path,
+ module_file,
+ subfolder=subfolder,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ local_dir=local_dir,
+ trust_remote_code=trust_remote_code,
+ )
+ return get_class_in_module(class_name, final_module)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/export_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/export_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..005ffd6e10e8caa1ea5870765c392acc240366fe
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/export_utils.py
@@ -0,0 +1,388 @@
+from __future__ import annotations
+
+import io
+import random
+import struct
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from fractions import Fraction
+from itertools import chain
+from typing import TYPE_CHECKING
+
+import numpy as np
+import PIL.Image
+import PIL.ImageOps
+
+from .import_utils import BACKENDS_MAPPING, is_imageio_available, is_opencv_available
+from .logging import get_logger
+
+
+if TYPE_CHECKING:
+ import torch
+
+
+global_rng = random.Random()
+
+logger = get_logger(__name__)
+
+
+@contextmanager
+def buffered_writer(raw_f):
+ f = io.BufferedWriter(raw_f)
+ yield f
+ f.flush()
+
+
+def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None, fps: int = 10) -> str:
+ if output_gif_path is None:
+ output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name
+
+ image[0].save(
+ output_gif_path,
+ save_all=True,
+ append_images=image[1:],
+ optimize=False,
+ duration=1000 // fps,
+ loop=0,
+ )
+ return output_gif_path
+
+
+def export_to_ply(mesh, output_ply_path: str = None):
+ """
+ Write a PLY file for a mesh.
+ """
+ if output_ply_path is None:
+ output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply").name
+
+ coords = mesh.verts.detach().cpu().numpy()
+ faces = mesh.faces.cpu().numpy()
+ rgb = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1)
+
+ with buffered_writer(open(output_ply_path, "wb")) as f:
+ f.write(b"ply\n")
+ f.write(b"format binary_little_endian 1.0\n")
+ f.write(bytes(f"element vertex {len(coords)}\n", "ascii"))
+ f.write(b"property float x\n")
+ f.write(b"property float y\n")
+ f.write(b"property float z\n")
+ if rgb is not None:
+ f.write(b"property uchar red\n")
+ f.write(b"property uchar green\n")
+ f.write(b"property uchar blue\n")
+ if faces is not None:
+ f.write(bytes(f"element face {len(faces)}\n", "ascii"))
+ f.write(b"property list uchar int vertex_index\n")
+ f.write(b"end_header\n")
+
+ if rgb is not None:
+ rgb = (rgb * 255.499).round().astype(int)
+ vertices = [
+ (*coord, *rgb)
+ for coord, rgb in zip(
+ coords.tolist(),
+ rgb.tolist(),
+ )
+ ]
+ format = struct.Struct("<3f3B")
+ for item in vertices:
+ f.write(format.pack(*item))
+ else:
+ format = struct.Struct("<3f")
+ for vertex in coords.tolist():
+ f.write(format.pack(*vertex))
+
+ if faces is not None:
+ format = struct.Struct(" str:
+ """
+ quality:
+ Video output quality. Default is 5. Uses variable bit rate. Highest quality is 10, lowest is 0. Set to None to
+ prevent variable bitrate flags to FFMPEG so you can manually specify them using output_params instead.
+ Specifying a fixed bitrate using `bitrate` disables this parameter.
+
+ bitrate:
+ Set a constant bitrate for the video encoding. Default is None causing `quality` parameter to be used instead.
+ Better quality videos with smaller file sizes will result from using the `quality` variable bitrate parameter
+ rather than specifying a fixed bitrate with this parameter.
+
+ macro_block_size:
+ Size constraint for video. Width and height, must be divisible by this number. If not divisible by this number
+ imageio will tell ffmpeg to scale the image up to the next closest size divisible by this number. Most codecs
+ are compatible with a macroblock size of 16 (default), some can go smaller (4, 8). To disable this automatic
+ feature set it to None or 1, however be warned many players can't decode videos that are odd in size and some
+ codecs will produce poor results or fail. See https://en.wikipedia.org/wiki/Macroblock.
+ """
+ # TODO: Dhruv. Remove by Diffusers release 0.33.0
+ # Added to prevent breaking existing code
+ if not is_imageio_available():
+ logger.warning(
+ (
+ "It is recommended to use `export_to_video` with `imageio` and `imageio-ffmpeg` as a backend. \n"
+ "These libraries are not present in your environment. Attempting to use legacy OpenCV backend to export video. \n"
+ "Support for the OpenCV backend will be deprecated in a future Diffusers version"
+ )
+ )
+ return _legacy_export_to_video(video_frames, output_video_path, fps)
+
+ if is_imageio_available():
+ import imageio
+ else:
+ raise ImportError(BACKENDS_MAPPING["imageio"][1].format("export_to_video"))
+
+ try:
+ imageio.plugins.ffmpeg.get_exe()
+ except AttributeError:
+ raise AttributeError(
+ (
+ "Found an existing imageio backend in your environment. Attempting to export video with imageio. \n"
+ "Unable to find a compatible ffmpeg installation in your environment to use with imageio. Please install via `pip install imageio-ffmpeg"
+ )
+ )
+
+ if output_video_path is None:
+ output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+
+ if isinstance(video_frames[0], np.ndarray):
+ video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames]
+
+ elif isinstance(video_frames[0], PIL.Image.Image):
+ video_frames = [np.array(frame) for frame in video_frames]
+
+ with imageio.get_writer(
+ output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size
+ ) as writer:
+ for frame in video_frames:
+ writer.append_data(frame)
+
+ return output_video_path
+
+
+def _import_av():
+ try:
+ import av
+ except ImportError as error:
+ raise ImportError(
+ "PyAV is required to use `encode_video`. You can install it with `pip install av`."
+ ) from error
+
+ return av
+
+
+def _prepare_audio_stream(container, audio_sample_rate: int):
+ """
+ Prepare the audio stream for writing.
+ """
+ audio_stream = container.add_stream("aac", rate=audio_sample_rate)
+ audio_stream.codec_context.sample_rate = audio_sample_rate
+ audio_stream.codec_context.layout = "stereo"
+ audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
+ return audio_stream
+
+
+def _resample_audio(container, audio_stream, frame_in, av_module) -> None:
+ cc = audio_stream.codec_context
+
+ # Use the encoder's format/layout/rate as the target.
+ target_format = cc.format or "fltp" # AAC is usually fltp.
+ target_layout = cc.layout or "stereo"
+ target_rate = cc.sample_rate or frame_in.sample_rate
+
+ audio_resampler = av_module.audio.resampler.AudioResampler(
+ format=target_format,
+ layout=target_layout,
+ rate=target_rate,
+ )
+
+ audio_next_pts = 0
+ for rframe in audio_resampler.resample(frame_in):
+ if rframe.pts is None:
+ rframe.pts = audio_next_pts
+ audio_next_pts += rframe.samples
+ rframe.sample_rate = frame_in.sample_rate
+ container.mux(audio_stream.encode(rframe))
+
+ # Flush audio encoder.
+ for packet in audio_stream.encode():
+ container.mux(packet)
+
+
+def _write_audio(
+ container,
+ audio_stream,
+ samples: "torch.Tensor",
+ audio_sample_rate: int,
+ av_module,
+) -> None:
+ import torch
+
+ if samples.ndim == 1:
+ samples = samples[:, None]
+
+ if samples.shape[1] != 2 and samples.shape[0] == 2:
+ samples = samples.T
+
+ if samples.shape[1] != 2:
+ raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
+
+ # Convert to int16 packed for ingestion; resampler converts to the encoder format.
+ if samples.dtype != torch.int16:
+ samples = torch.clip(samples, -1.0, 1.0)
+ samples = (samples * 32767.0).to(torch.int16)
+
+ frame_in = av_module.AudioFrame.from_ndarray(
+ samples.contiguous().reshape(1, -1).cpu().numpy(),
+ format="s16",
+ layout="stereo",
+ )
+ frame_in.sample_rate = audio_sample_rate
+
+ _resample_audio(container, audio_stream, frame_in, av_module)
+
+
+def encode_video(
+ video: list[PIL.Image.Image] | np.ndarray | "torch.Tensor" | Iterator["torch.Tensor"],
+ fps: int,
+ output_path: str,
+ audio: "torch.Tensor" | None = None,
+ audio_sample_rate: int | None = None,
+ video_chunks_number: int = 1,
+) -> None:
+ """
+ Encodes a video with optional audio using the PyAV library. Based on code from the original LTX-2 repo:
+ https://github.com/Lightricks/LTX-2/blob/4f410820b198e05074a1e92de793e3b59e9ab5a0/packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py#L182
+
+ Args:
+ video (`List[PIL.Image.Image]` or `np.ndarray` or `torch.Tensor`):
+ A video tensor of shape [frames, height, width, channels] with integer pixel values in [0, 255]. If the
+ input is a `np.ndarray`, it is expected to be a float array with values in [0, 1] (which is what pipelines
+ usually return with `output_type="np"`).
+ fps (`int`)
+ The frames per second (FPS) of the encoded video.
+ output_path (`str`):
+ The path to save the encoded video to.
+ audio (`torch.Tensor`, *optional*):
+ An audio waveform of shape [audio_channels, samples].
+ audio_sample_rate: (`int`, *optional*):
+ The sampling rate of the audio waveform.
+ video_chunks_number (`int`, *optional*, defaults to `1`):
+ The number of chunks to split the video into for encoding. Each chunk will be encoded separately. The
+ number of chunks to use often depends on the tiling config for the video VAE.
+ """
+ av = _import_av()
+ import torch
+ from tqdm import tqdm
+
+ if isinstance(video, list) and isinstance(video[0], PIL.Image.Image):
+ # Pipeline output_type="pil"; assumes each image is in "RGB" mode.
+ video_frames = [np.array(frame) for frame in video]
+ video = np.stack(video_frames, axis=0)
+ video = torch.from_numpy(video)
+ elif isinstance(video, np.ndarray):
+ # Pipeline output_type="np".
+ is_denormalized = np.logical_and(np.zeros_like(video) <= video, video <= np.ones_like(video))
+ if np.all(is_denormalized):
+ video = (video * 255).round().astype("uint8")
+ else:
+ logger.warning(
+ "Supplied `numpy.ndarray` does not have values in [0, 1]. The values will be assumed to be pixel "
+ "values in [0, ..., 255] and will be used as is."
+ )
+ video = torch.from_numpy(video)
+
+ if isinstance(video, torch.Tensor):
+ # Split into video_chunks_number along the frame dimension.
+ video = torch.tensor_split(video, video_chunks_number, dim=0)
+ video = iter(video)
+
+ first_chunk = next(video)
+
+ _, height, width, _ = first_chunk.shape
+
+ container = av.open(output_path, mode="w")
+ stream = container.add_stream("libx264", rate=int(fps))
+ stream.width = width
+ stream.height = height
+ stream.pix_fmt = "yuv420p"
+
+ if audio is not None:
+ if audio_sample_rate is None:
+ raise ValueError("audio_sample_rate is required when audio is provided")
+
+ audio_stream = _prepare_audio_stream(container, audio_sample_rate)
+
+ for video_chunk in tqdm(chain([first_chunk], video), total=video_chunks_number, desc="Encoding video chunks"):
+ video_chunk_cpu = video_chunk.to("cpu").numpy()
+ for frame_array in video_chunk_cpu:
+ frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
+ for packet in stream.encode(frame):
+ container.mux(packet)
+
+ # Flush video encoder.
+ for packet in stream.encode():
+ container.mux(packet)
+
+ if audio is not None:
+ _write_audio(container, audio_stream, audio, audio_sample_rate, av)
+
+ container.close()
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/hub_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/hub_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c20653cd70bf747c1ad834c996acd6f9ffc97a5c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/hub_utils.py
@@ -0,0 +1,593 @@
+# coding=utf-8
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import json
+import os
+import re
+import sys
+import tempfile
+import warnings
+from pathlib import Path
+from uuid import uuid4
+
+from huggingface_hub import (
+ DDUFEntry,
+ ModelCard,
+ ModelCardData,
+ create_repo,
+ hf_hub_download,
+ model_info,
+ snapshot_download,
+ upload_folder,
+)
+from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY, HF_HUB_OFFLINE
+from huggingface_hub.file_download import REGEX_COMMIT_HASH
+from huggingface_hub.utils import (
+ EntryNotFoundError,
+ HfHubHTTPError,
+ RepositoryNotFoundError,
+ RevisionNotFoundError,
+ is_jinja_available,
+ validate_hf_hub_args,
+)
+from packaging import version
+
+from .. import __version__
+from .constants import (
+ DEPRECATED_REVISION_ARGS,
+ HUGGINGFACE_CO_RESOLVE_ENDPOINT,
+ SAFETENSORS_WEIGHTS_NAME,
+ WEIGHTS_NAME,
+)
+from .import_utils import (
+ ENV_VARS_TRUE_VALUES,
+ _onnxruntime_version,
+ _torch_version,
+ is_onnx_available,
+ is_torch_available,
+)
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "model_card_template.md"
+SESSION_ID = uuid4().hex
+
+
+def http_user_agent(user_agent: dict | str | None = None) -> str:
+ """
+ Formats a user-agent string with basic info about a request.
+ """
+ ua = f"diffusers/{__version__}; python/{sys.version.split()[0]}; session_id/{SESSION_ID}"
+ if HF_HUB_DISABLE_TELEMETRY or HF_HUB_OFFLINE:
+ return ua + "; telemetry/off"
+ if is_torch_available():
+ ua += f"; torch/{_torch_version}"
+ if is_onnx_available():
+ ua += f"; onnxruntime/{_onnxruntime_version}"
+ # CI will set this value to True
+ if os.environ.get("DIFFUSERS_IS_CI", "").upper() in ENV_VARS_TRUE_VALUES:
+ ua += "; is_ci/true"
+ if isinstance(user_agent, dict):
+ ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items())
+ elif isinstance(user_agent, str):
+ ua += "; " + user_agent
+ return ua
+
+
+def load_or_create_model_card(
+ repo_id_or_path: str = None,
+ token: str | None = None,
+ is_pipeline: bool = False,
+ from_training: bool = False,
+ model_description: str | None = None,
+ base_model: str = None,
+ prompt: str | None = None,
+ license: str | None = None,
+ widget: list[dict] | None = None,
+ inference: bool | None = None,
+ is_modular: bool = False,
+ update_model_card: bool = False,
+) -> ModelCard:
+ """
+ Loads or creates a model card.
+
+ Args:
+ repo_id_or_path (`str`):
+ The repo id (e.g., "stable-diffusion-v1-5/stable-diffusion-v1-5") or local path where to look for the model
+ card.
+ token (`str`, *optional*):
+ Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more
+ details.
+ is_pipeline (`bool`):
+ Boolean to indicate if we're adding tag to a [`DiffusionPipeline`].
+ from_training: (`bool`): Boolean flag to denote if the model card is being created from a training script.
+ model_description (`str`, *optional*): Model description to add to the model card. Helpful when using
+ `load_or_create_model_card` from a training script.
+ base_model (`str`): Base model identifier (e.g., "stabilityai/stable-diffusion-xl-base-1.0"). Useful
+ for DreamBooth-like training.
+ prompt (`str`, *optional*): Prompt used for training. Useful for DreamBooth-like training.
+ license: (`str`, *optional*): License of the output artifact. Helpful when using
+ `load_or_create_model_card` from a training script.
+ widget (`list[dict]`, *optional*): Widget to accompany a gallery template.
+ inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using
+ `load_or_create_model_card` from a training script.
+ is_modular: (`bool`, optional): Boolean flag to denote if the model card is for a modular pipeline.
+ When True, uses model_description as-is without additional template formatting.
+ update_model_card: (`bool`, optional): When True, regenerates the model card content even if one
+ already exists on the remote repo. Existing card metadata (tags, license, etc.) is preserved. Only
+ supported for modular pipelines (i.e., `is_modular=True`).
+ """
+ if not is_jinja_available():
+ raise ValueError(
+ "Modelcard rendering is based on Jinja templates."
+ " Please make sure to have `jinja` installed before using `load_or_create_model_card`."
+ " To install it, please run `pip install Jinja2`."
+ )
+
+ if update_model_card and not is_modular:
+ raise ValueError("`update_model_card=True` is only supported for modular pipelines (`is_modular=True`).")
+
+ try:
+ # Check if the model card is present on the remote repo
+ model_card = ModelCard.load(repo_id_or_path, token=token)
+ # For modular pipelines, regenerate card content when requested (preserve existing metadata)
+ if update_model_card and is_modular and model_description is not None:
+ existing_data = model_card.data
+ model_card = ModelCard(model_description)
+ model_card.data = existing_data
+ except (EntryNotFoundError, RepositoryNotFoundError):
+ # Otherwise create a model card from template
+ if from_training:
+ model_card = ModelCard.from_template(
+ card_data=ModelCardData( # Card metadata object that will be converted to YAML block
+ license=license,
+ library_name="diffusers",
+ inference=inference,
+ base_model=base_model,
+ instance_prompt=prompt,
+ widget=widget,
+ ),
+ template_path=MODEL_CARD_TEMPLATE_PATH,
+ model_description=model_description,
+ )
+ else:
+ card_data = ModelCardData()
+ if is_modular and model_description is not None:
+ model_card = ModelCard(model_description)
+ model_card.data = card_data
+ else:
+ component = "pipeline" if is_pipeline else "model"
+ if model_description is None:
+ model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated."
+ model_card = ModelCard.from_template(card_data, model_description=model_description)
+
+ return model_card
+
+
+def populate_model_card(model_card: ModelCard, tags: str | list[str] | None = None) -> ModelCard:
+ """Populates the `model_card` with library name and optional tags."""
+ if model_card.data.library_name is None:
+ model_card.data.library_name = "diffusers"
+
+ if tags is not None:
+ if isinstance(tags, str):
+ tags = [tags]
+ if model_card.data.tags is None:
+ model_card.data.tags = []
+ for tag in tags:
+ model_card.data.tags.append(tag)
+
+ return model_card
+
+
+def extract_commit_hash(resolved_file: str | None, commit_hash: str | None = None):
+ """
+ Extracts the commit hash from a resolved filename toward a cache file.
+ """
+ if resolved_file is None or commit_hash is not None:
+ return commit_hash
+ resolved_file = str(Path(resolved_file).as_posix())
+ search = re.search(r"snapshots/([^/]+)/", resolved_file)
+ if search is None:
+ return None
+ commit_hash = search.groups()[0]
+ return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None
+
+
+def _add_variant(weights_name: str, variant: str | None = None) -> str:
+ if variant is not None:
+ splits = weights_name.split(".")
+ splits = splits[:-1] + [variant] + splits[-1:]
+ weights_name = ".".join(splits)
+
+ return weights_name
+
+
+@validate_hf_hub_args
+def _get_model_file(
+ pretrained_model_name_or_path: str | Path,
+ *,
+ weights_name: str,
+ subfolder: str | None = None,
+ cache_dir: str | None = None,
+ force_download: bool = False,
+ proxies: dict | None = None,
+ local_files_only: bool = False,
+ token: str | None = None,
+ user_agent: dict | str | None = None,
+ revision: str | None = None,
+ commit_hash: str | None = None,
+ dduf_entries: dict[str, DDUFEntry] | None = None,
+):
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+
+ if dduf_entries:
+ if subfolder is not None:
+ raise ValueError(
+ "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). "
+ "Please check the DDUF structure"
+ )
+ model_file = (
+ weights_name
+ if pretrained_model_name_or_path == ""
+ else "/".join([pretrained_model_name_or_path, weights_name])
+ )
+ if model_file in dduf_entries:
+ return model_file
+ else:
+ raise EnvironmentError(f"Error no file named {weights_name} found in archive {dduf_entries.keys()}.")
+ elif os.path.isfile(pretrained_model_name_or_path):
+ return pretrained_model_name_or_path
+ elif os.path.isdir(pretrained_model_name_or_path):
+ if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)):
+ # Load from a PyTorch checkpoint
+ model_file = os.path.join(pretrained_model_name_or_path, weights_name)
+ return model_file
+ elif subfolder is not None and os.path.isfile(
+ os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
+ ):
+ model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
+ return model_file
+ else:
+ raise EnvironmentError(
+ f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}."
+ )
+ else:
+ # 1. First check if deprecated way of loading from branches is used
+ if (
+ revision in DEPRECATED_REVISION_ARGS
+ and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME)
+ and version.parse(version.parse(__version__).base_version) >= version.parse("0.22.0")
+ ):
+ try:
+ model_file = hf_hub_download(
+ pretrained_model_name_or_path,
+ filename=_add_variant(weights_name, revision),
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ subfolder=subfolder,
+ revision=revision or commit_hash,
+ )
+ warnings.warn(
+ f"Loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` is deprecated. Loading instead from `revision='main'` with `variant={revision}`. Loading model variants via `revision='{revision}'` will be removed in diffusers v1. Please use `variant='{revision}'` instead.",
+ FutureWarning,
+ )
+ return model_file
+ except: # noqa: E722
+ warnings.warn(
+ f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(weights_name, revision)} file in the 'main' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {_add_variant(weights_name, revision)}' so that the correct variant file can be added.",
+ FutureWarning,
+ )
+ try:
+ # 2. Load model file as usual
+ model_file = hf_hub_download(
+ pretrained_model_name_or_path,
+ filename=weights_name,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ subfolder=subfolder,
+ revision=revision or commit_hash,
+ )
+ return model_file
+
+ except RepositoryNotFoundError as e:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "
+ "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "
+ "token having permission to this repo with `token` or log in with `hf auth login`."
+ ) from e
+ except RevisionNotFoundError as e:
+ raise EnvironmentError(
+ f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "
+ "this model name. Check the model page at "
+ f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
+ ) from e
+ except EntryNotFoundError as e:
+ raise EnvironmentError(
+ f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}."
+ ) from e
+ except HfHubHTTPError as e:
+ raise EnvironmentError(
+ f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{e}"
+ ) from e
+ except ValueError as e:
+ raise EnvironmentError(
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
+ f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
+ f" directory containing a file named {weights_name} or"
+ " \nCheckout your internet connection or see how to run the library in"
+ " offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
+ ) from e
+ except EnvironmentError as e:
+ raise EnvironmentError(
+ f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "
+ "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
+ f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
+ f"containing a file named {weights_name}"
+ ) from e
+
+
+def _get_checkpoint_shard_files(
+ pretrained_model_name_or_path,
+ index_filename,
+ cache_dir=None,
+ proxies=None,
+ local_files_only=False,
+ token=None,
+ user_agent=None,
+ revision=None,
+ subfolder="",
+ dduf_entries: dict[str, DDUFEntry] | None = None,
+):
+ """
+ For a given model:
+
+ - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the
+ Hub
+ - returns the list of paths to all the shards, as well as some metadata.
+
+ For the description of each arg, see [`PreTrainedModel.from_pretrained`]. `index_filename` is the full path to the
+ index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub).
+ """
+ if dduf_entries:
+ if index_filename not in dduf_entries:
+ raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.")
+ else:
+ if not os.path.isfile(index_filename):
+ raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.")
+
+ if dduf_entries:
+ index = json.loads(dduf_entries[index_filename].read_text())
+ else:
+ with open(index_filename, "r") as f:
+ index = json.loads(f.read())
+
+ original_shard_filenames = sorted(set(index["weight_map"].values()))
+ for shard_filename in original_shard_filenames:
+ if os.path.basename(shard_filename) != shard_filename:
+ raise ValueError(
+ f"The shard filename {shard_filename!r} in the checkpoint index contains a path separator or a "
+ "parent-directory reference, which is not allowed. Shard filenames must be plain filenames located "
+ "in the model directory."
+ )
+ sharded_metadata = index["metadata"]
+ sharded_metadata["all_checkpoint_keys"] = list(index["weight_map"].keys())
+ sharded_metadata["weight_map"] = index["weight_map"].copy()
+ shards_path = os.path.join(pretrained_model_name_or_path, subfolder)
+
+ # First, let's deal with local folder.
+ if os.path.isdir(pretrained_model_name_or_path) or dduf_entries:
+ shard_filenames = [os.path.join(shards_path, f) for f in original_shard_filenames]
+ for shard_file in shard_filenames:
+ if dduf_entries:
+ if shard_file not in dduf_entries:
+ raise FileNotFoundError(
+ f"{shards_path} does not appear to have a file named {shard_file} which is "
+ "required according to the checkpoint index."
+ )
+ else:
+ if not os.path.exists(shard_file):
+ raise FileNotFoundError(
+ f"{shards_path} does not appear to have a file named {shard_file} which is "
+ "required according to the checkpoint index."
+ )
+ return shard_filenames, sharded_metadata
+
+ # At this stage pretrained_model_name_or_path is a model identifier on the Hub
+ allow_patterns = original_shard_filenames
+ if subfolder is not None:
+ allow_patterns = [os.path.join(subfolder, p) for p in allow_patterns]
+
+ ignore_patterns = ["*.json", "*.md"]
+
+ # If the repo doesn't have the required shards, error out early even before downloading anything.
+ if not local_files_only:
+ model_files_info = model_info(pretrained_model_name_or_path, revision=revision, token=token)
+ for shard_file in original_shard_filenames:
+ shard_file_present = any(shard_file in k.rfilename for k in model_files_info.siblings)
+ if not shard_file_present:
+ raise EnvironmentError(
+ f"{shards_path} does not appear to have a file named {shard_file} which is "
+ "required according to the checkpoint index."
+ )
+
+ try:
+ # Load from URL
+ cached_folder = snapshot_download(
+ pretrained_model_name_or_path,
+ cache_dir=cache_dir,
+ proxies=proxies,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ allow_patterns=allow_patterns,
+ ignore_patterns=ignore_patterns,
+ user_agent=user_agent,
+ )
+ if subfolder is not None:
+ cached_folder = os.path.join(cached_folder, subfolder)
+
+ # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index, so
+ # we don't have to catch them here. We have also dealt with EntryNotFoundError.
+ except HfHubHTTPError as e:
+ raise EnvironmentError(
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {pretrained_model_name_or_path}. You should try"
+ " again after checking your internet connection."
+ ) from e
+
+ cached_filenames = [os.path.join(cached_folder, f) for f in original_shard_filenames]
+ for cached_file in cached_filenames:
+ if not os.path.isfile(cached_file):
+ raise EnvironmentError(
+ f"{cached_folder} does not have a file named {cached_file} which is required according to the checkpoint index."
+ )
+
+ return cached_filenames, sharded_metadata
+
+
+def _check_legacy_sharding_variant_format(folder: str = None, filenames: list[str] = None, variant: str = None):
+ if filenames and folder:
+ raise ValueError("Both `filenames` and `folder` cannot be provided.")
+ if not filenames:
+ filenames = []
+ for _, _, files in os.walk(folder):
+ for file in files:
+ filenames.append(os.path.basename(file))
+ transformers_index_format = r"\d{5}-of-\d{5}"
+ variant_file_re = re.compile(rf".*-{transformers_index_format}\.{variant}\.[a-z]+$")
+ return any(variant_file_re.match(f) is not None for f in filenames)
+
+
+class PushToHubMixin:
+ """
+ A Mixin to push a model, scheduler, or pipeline to the Hugging Face Hub.
+ """
+
+ def _upload_folder(
+ self,
+ working_dir: str | os.PathLike,
+ repo_id: str,
+ token: str | None = None,
+ commit_message: str | None = None,
+ create_pr: bool = False,
+ subfolder: str | None = None,
+ ):
+ """
+ Uploads all files in `working_dir` to `repo_id`.
+ """
+ if commit_message is None:
+ if "Model" in self.__class__.__name__:
+ commit_message = "Upload model"
+ elif "Scheduler" in self.__class__.__name__:
+ commit_message = "Upload scheduler"
+ else:
+ commit_message = f"Upload {self.__class__.__name__}"
+
+ logger.info(f"Uploading the files of {working_dir} to {repo_id}.")
+ return upload_folder(
+ repo_id=repo_id,
+ folder_path=working_dir,
+ token=token,
+ commit_message=commit_message,
+ create_pr=create_pr,
+ path_in_repo=subfolder,
+ )
+
+ def push_to_hub(
+ self,
+ repo_id: str,
+ commit_message: str | None = None,
+ private: bool | None = None,
+ token: str | None = None,
+ create_pr: bool = False,
+ safe_serialization: bool = True,
+ variant: str | None = None,
+ subfolder: str | None = None,
+ ) -> str:
+ """
+ Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub.
+
+ Parameters:
+ repo_id (`str`):
+ The name of the repository you want to push your model, scheduler, or pipeline files to. It should
+ contain your organization name when pushing to an organization. `repo_id` can also be a path to a local
+ directory.
+ commit_message (`str`, *optional*):
+ Message to commit while pushing. Default to `"Upload {object}"`.
+ private (`bool`, *optional*):
+ Whether to make the repo private. If `None` (default), the repo will be public unless the
+ organization's default is private. This value is ignored if the repo already exists.
+ token (`str`, *optional*):
+ The token to use as HTTP bearer authorization for remote files. The token generated when running `hf
+ auth login` (stored in `~/.huggingface`).
+ create_pr (`bool`, *optional*, defaults to `False`):
+ Whether or not to create a PR with the uploaded files or directly commit.
+ safe_serialization (`bool`, *optional*, defaults to `True`):
+ Whether or not to convert the model weights to the `safetensors` format.
+ variant (`str`, *optional*):
+ If specified, weights are saved in the format `pytorch_model..bin`.
+
+ Examples:
+
+ ```python
+ from diffusers import UNet2DConditionModel
+
+ unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet")
+
+ # Push the `unet` to your namespace with the name "my-finetuned-unet".
+ unet.push_to_hub("my-finetuned-unet")
+
+ # Push the `unet` to an organization with the name "my-finetuned-unet".
+ unet.push_to_hub("your-org/my-finetuned-unet")
+ ```
+ """
+ repo_id = create_repo(repo_id, private=private, token=token, exist_ok=True).repo_id
+
+ # Create a new empty model card and eventually tag it
+ if not subfolder:
+ model_card = load_or_create_model_card(repo_id, token=token)
+ model_card = populate_model_card(model_card)
+
+ # Save all files.
+ save_kwargs = {"safe_serialization": safe_serialization}
+ if "Scheduler" not in self.__class__.__name__:
+ save_kwargs.update({"variant": variant})
+
+ with tempfile.TemporaryDirectory() as tmpdir:
+ self.save_pretrained(tmpdir, **save_kwargs)
+
+ # Update model card if needed:
+ if not subfolder:
+ model_card.save(os.path.join(tmpdir, "README.md"))
+
+ return self._upload_folder(
+ tmpdir,
+ repo_id,
+ token=token,
+ commit_message=commit_message,
+ create_pr=create_pr,
+ subfolder=subfolder,
+ )
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/import_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/import_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d7d310937ef7ca14066aa7affa8f40385679d05
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/import_utils.py
@@ -0,0 +1,1044 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Import utilities: Utilities related to imports and our lazy inits.
+"""
+
+import importlib.util
+import inspect
+import operator as op
+import os
+import sys
+from collections import OrderedDict, defaultdict
+from functools import lru_cache as cache
+from itertools import chain
+from types import ModuleType
+from typing import Any
+
+from huggingface_hub.utils import is_jinja_available # noqa: F401
+from packaging.version import Version, parse
+
+from . import logging
+
+
+# The package importlib_metadata is in a different place, depending on the python version.
+if sys.version_info < (3, 8):
+ import importlib_metadata
+else:
+ import importlib.metadata as importlib_metadata
+try:
+ _package_map = importlib_metadata.packages_distributions() # load-once to avoid expensive calls
+except Exception:
+ _package_map = None
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
+ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
+
+USE_TF = os.environ.get("USE_TF", "AUTO").upper()
+USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
+USE_SAFETENSORS = os.environ.get("USE_SAFETENSORS", "AUTO").upper()
+DIFFUSERS_SLOW_IMPORT = os.environ.get("DIFFUSERS_SLOW_IMPORT", "FALSE").upper()
+DIFFUSERS_SLOW_IMPORT = DIFFUSERS_SLOW_IMPORT in ENV_VARS_TRUE_VALUES
+
+STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
+
+_is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ)
+
+
+def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[bool, str]:
+ global _package_map
+ pkg_exists = importlib.util.find_spec(pkg_name) is not None
+ pkg_version = "N/A"
+
+ if pkg_exists:
+ if _package_map is None:
+ _package_map = defaultdict(list)
+ try:
+ # Fallback for Python < 3.10
+ for dist in importlib_metadata.distributions():
+ _top_level_declared = (dist.read_text("top_level.txt") or "").split()
+ # Infer top-level package names from file structure
+ _inferred_opt_names = {
+ f.parts[0] if len(f.parts) > 1 else inspect.getmodulename(f) for f in (dist.files or [])
+ } - {None}
+ _top_level_inferred = filter(lambda name: "." not in name, _inferred_opt_names)
+ for pkg in _top_level_declared or _top_level_inferred:
+ _package_map[pkg].append(dist.metadata["Name"])
+ except Exception as _:
+ pass
+ try:
+ if get_dist_name and pkg_name in _package_map and _package_map[pkg_name]:
+ if len(_package_map[pkg_name]) > 1:
+ logger.warning(
+ f"Multiple distributions found for package {pkg_name}. Picked distribution: {_package_map[pkg_name][0]}"
+ )
+ pkg_name = _package_map[pkg_name][0]
+ pkg_version = importlib_metadata.version(pkg_name)
+ logger.debug(f"Successfully imported {pkg_name} version {pkg_version}")
+ except (ImportError, importlib_metadata.PackageNotFoundError):
+ pkg_exists = False
+
+ return pkg_exists, pkg_version
+
+
+if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
+ _torch_available, _torch_version = _is_package_available("torch")
+
+else:
+ logger.info("Disabling PyTorch because USE_TORCH is set")
+ _torch_available = False
+ _torch_version = "N/A"
+
+if USE_SAFETENSORS in ENV_VARS_TRUE_AND_AUTO_VALUES:
+ _safetensors_available, _safetensors_version = _is_package_available("safetensors")
+
+else:
+ logger.info("Disabling Safetensors because USE_SAFETENSORS is set")
+ _safetensors_available = False
+
+_onnxruntime_version = "N/A"
+_onnx_available = importlib.util.find_spec("onnxruntime") is not None
+if _onnx_available:
+ candidates = (
+ "onnxruntime",
+ "onnxruntime-cann",
+ "onnxruntime-directml",
+ "ort_nightly_directml",
+ "onnxruntime-gpu",
+ "ort_nightly_gpu",
+ "onnxruntime-migraphx",
+ "onnxruntime-openvino",
+ "onnxruntime-qnn",
+ "onnxruntime-rocm",
+ "onnxruntime-training",
+ "onnxruntime-vitisai",
+ )
+ _onnxruntime_version = None
+ # For the metadata, we have to look for both onnxruntime and onnxruntime-x
+ for pkg in candidates:
+ try:
+ _onnxruntime_version = importlib_metadata.version(pkg)
+ break
+ except importlib_metadata.PackageNotFoundError:
+ pass
+ _onnx_available = _onnxruntime_version is not None
+ if _onnx_available:
+ logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}")
+
+# (sayakpaul): importlib.util.find_spec("opencv-python") returns None even when it's installed.
+# _opencv_available = importlib.util.find_spec("opencv-python") is not None
+try:
+ candidates = (
+ "opencv-python",
+ "opencv-contrib-python",
+ "opencv-python-headless",
+ "opencv-contrib-python-headless",
+ )
+ _opencv_version = None
+ for pkg in candidates:
+ try:
+ _opencv_version = importlib_metadata.version(pkg)
+ break
+ except importlib_metadata.PackageNotFoundError:
+ pass
+ _opencv_available = _opencv_version is not None
+ if _opencv_available:
+ logger.debug(f"Successfully imported cv2 version {_opencv_version}")
+except importlib_metadata.PackageNotFoundError:
+ _opencv_available = False
+
+_bs4_available = importlib.util.find_spec("bs4") is not None
+try:
+ # importlib metadata under different name
+ _bs4_version = importlib_metadata.version("beautifulsoup4")
+ logger.debug(f"Successfully imported ftfy version {_bs4_version}")
+except importlib_metadata.PackageNotFoundError:
+ _bs4_available = False
+
+_invisible_watermark_available = importlib.util.find_spec("imwatermark") is not None
+try:
+ _invisible_watermark_version = importlib_metadata.version("invisible-watermark")
+ logger.debug(f"Successfully imported invisible-watermark version {_invisible_watermark_version}")
+except importlib_metadata.PackageNotFoundError:
+ _invisible_watermark_available = False
+
+_torch_xla_available, _torch_xla_version = _is_package_available("torch_xla")
+_torch_npu_available, _torch_npu_version = _is_package_available("torch_npu")
+_torch_mlu_available, _torch_mlu_version = _is_package_available("torch_mlu")
+_torch_neuronx_available, _torch_neuronx_version = _is_package_available("torch_neuronx")
+_transformers_available, _transformers_version = _is_package_available("transformers")
+_hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub")
+_kernels_available, _kernels_version = _is_package_available("kernels")
+_inflect_available, _inflect_version = _is_package_available("inflect")
+_unidecode_available, _unidecode_version = _is_package_available("unidecode")
+
+_note_seq_available, _note_seq_version = _is_package_available("note_seq")
+_wandb_available, _wandb_version = _is_package_available("wandb")
+_tensorboard_available, _tensorboard_version = _is_package_available("tensorboard")
+_compel_available, _compel_version = _is_package_available("compel")
+_sentencepiece_available, _sentencepiece_version = _is_package_available("sentencepiece")
+_outlines_available, _outlines_version = _is_package_available("outlines")
+_torchsde_available, _torchsde_version = _is_package_available("torchsde")
+_peft_available, _peft_version = _is_package_available("peft")
+_torchvision_available, _torchvision_version = _is_package_available("torchvision")
+_matplotlib_available, _matplotlib_version = _is_package_available("matplotlib")
+_timm_available, _timm_version = _is_package_available("timm")
+_bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes")
+_imageio_available, _imageio_version = _is_package_available("imageio")
+_ftfy_available, _ftfy_version = _is_package_available("ftfy")
+_scipy_available, _scipy_version = _is_package_available("scipy")
+_librosa_available, _librosa_version = _is_package_available("librosa")
+_accelerate_available, _accelerate_version = _is_package_available("accelerate")
+_xformers_available, _xformers_version = _is_package_available("xformers")
+_gguf_available, _gguf_version = _is_package_available("gguf")
+_torchao_available, _torchao_version = _is_package_available("torchao")
+_bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes")
+_optimum_quanto_available, _optimum_quanto_version = _is_package_available("optimum", get_dist_name=True)
+_pytorch_retinaface_available, _pytorch_retinaface_version = _is_package_available("pytorch_retinaface")
+_better_profanity_available, _better_profanity_version = _is_package_available("better_profanity")
+_nltk_available, _nltk_version = _is_package_available("nltk")
+_cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available("cosmos_guardrail")
+_sageattention_available, _sageattention_version = _is_package_available("sageattention")
+_flash_attn_available, _flash_attn_version = _is_package_available("flash_attn")
+_flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3")
+_aiter_available, _aiter_version = _is_package_available("aiter", get_dist_name=True)
+_kornia_available, _kornia_version = _is_package_available("kornia")
+_nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True)
+_auto_round_available, _auto_round_version = _is_package_available("auto_round")
+_sdnq_available, _sdnq_version = _is_package_available("sdnq")
+_flashpack_available, _flashpack_version = _is_package_available("flashpack")
+_av_available, _av_version = _is_package_available("av")
+
+
+def is_torch_available():
+ return _torch_available
+
+
+def is_torch_xla_available():
+ return _torch_xla_available
+
+
+def is_torch_npu_available():
+ return _torch_npu_available
+
+
+def is_torch_mlu_available():
+ return _torch_mlu_available
+
+
+def is_torch_neuronx_available():
+ return _torch_neuronx_available
+
+
+def is_transformers_available():
+ return _transformers_available
+
+
+def is_inflect_available():
+ return _inflect_available
+
+
+def is_unidecode_available():
+ return _unidecode_available
+
+
+def is_onnx_available():
+ return _onnx_available
+
+
+def is_opencv_available():
+ return _opencv_available
+
+
+def is_scipy_available():
+ return _scipy_available
+
+
+def is_librosa_available():
+ return _librosa_available
+
+
+def is_xformers_available():
+ return _xformers_available
+
+
+def is_accelerate_available():
+ return _accelerate_available
+
+
+def is_kernels_available():
+ return _kernels_available
+
+
+def is_note_seq_available():
+ return _note_seq_available
+
+
+def is_wandb_available():
+ return _wandb_available
+
+
+def is_tensorboard_available():
+ return _tensorboard_available
+
+
+def is_compel_available():
+ return _compel_available
+
+
+def is_ftfy_available():
+ return _ftfy_available
+
+
+def is_bs4_available():
+ return _bs4_available
+
+
+def is_torchsde_available():
+ return _torchsde_available
+
+
+def is_invisible_watermark_available():
+ return _invisible_watermark_available
+
+
+def is_peft_available():
+ return _peft_available
+
+
+def is_torchvision_available():
+ return _torchvision_available
+
+
+def is_matplotlib_available():
+ return _matplotlib_available
+
+
+def is_safetensors_available():
+ return _safetensors_available
+
+
+def is_bitsandbytes_available():
+ return _bitsandbytes_available
+
+
+def is_google_colab():
+ return _is_google_colab
+
+
+def is_sentencepiece_available():
+ return _sentencepiece_available
+
+
+def is_outlines_available():
+ return _outlines_available
+
+
+def is_imageio_available():
+ return _imageio_available
+
+
+def is_gguf_available():
+ return _gguf_available
+
+
+def is_flashpack_available():
+ return _flashpack_available
+
+
+def is_torchao_available():
+ return _torchao_available
+
+
+def is_optimum_quanto_available():
+ return _optimum_quanto_available
+
+
+def is_nvidia_modelopt_available():
+ return _nvidia_modelopt_available
+
+
+def is_auto_round_available():
+ return _auto_round_available
+
+
+def is_sdnq_available():
+ return _sdnq_available
+
+
+def is_timm_available():
+ return _timm_available
+
+
+def is_pytorch_retinaface_available():
+ return _pytorch_retinaface_available
+
+
+def is_better_profanity_available():
+ return _better_profanity_available
+
+
+def is_nltk_available():
+ return _nltk_available
+
+
+def is_cosmos_guardrail_available():
+ return _cosmos_guardrail_available
+
+
+def is_hpu_available():
+ return all(importlib.util.find_spec(lib) for lib in ("habana_frameworks", "habana_frameworks.torch"))
+
+
+def is_sageattention_available():
+ return _sageattention_available
+
+
+def is_flash_attn_available():
+ return _flash_attn_available
+
+
+def is_flash_attn_3_available():
+ return _flash_attn_3_available
+
+
+def is_aiter_available():
+ return _aiter_available
+
+
+def is_kornia_available():
+ return _kornia_available
+
+
+def is_av_available():
+ return _av_available
+
+
+# docstyle-ignore
+INFLECT_IMPORT_ERROR = """
+{0} requires the inflect library but it was not found in your environment. You can install it with pip: `pip install
+inflect`
+"""
+
+# docstyle-ignore
+PYTORCH_IMPORT_ERROR = """
+{0} requires the PyTorch library but it was not found in your environment. Checkout the instructions on the
+installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment.
+"""
+
+# docstyle-ignore
+ONNX_IMPORT_ERROR = """
+{0} requires the onnxruntime library but it was not found in your environment. You can install it with pip: `pip
+install onnxruntime`
+"""
+
+# docstyle-ignore
+OPENCV_IMPORT_ERROR = """
+{0} requires the OpenCV library but it was not found in your environment. You can install it with pip: `pip
+install opencv-python`
+"""
+
+# docstyle-ignore
+SCIPY_IMPORT_ERROR = """
+{0} requires the scipy library but it was not found in your environment. You can install it with pip: `pip install
+scipy`
+"""
+
+# docstyle-ignore
+LIBROSA_IMPORT_ERROR = """
+{0} requires the librosa library but it was not found in your environment. Checkout the instructions on the
+installation page: https://librosa.org/doc/latest/install.html and follow the ones that match your environment.
+"""
+
+# docstyle-ignore
+TRANSFORMERS_IMPORT_ERROR = """
+{0} requires the transformers library but it was not found in your environment. You can install it with pip: `pip
+install transformers`
+"""
+
+# docstyle-ignore
+UNIDECODE_IMPORT_ERROR = """
+{0} requires the unidecode library but it was not found in your environment. You can install it with pip: `pip install
+Unidecode`
+"""
+
+# docstyle-ignore
+NOTE_SEQ_IMPORT_ERROR = """
+{0} requires the note-seq library but it was not found in your environment. You can install it with pip: `pip
+install note-seq`
+"""
+
+# docstyle-ignore
+WANDB_IMPORT_ERROR = """
+{0} requires the wandb library but it was not found in your environment. You can install it with pip: `pip
+install wandb`
+"""
+
+# docstyle-ignore
+TENSORBOARD_IMPORT_ERROR = """
+{0} requires the tensorboard library but it was not found in your environment. You can install it with pip: `pip
+install tensorboard`
+"""
+
+
+# docstyle-ignore
+COMPEL_IMPORT_ERROR = """
+{0} requires the compel library but it was not found in your environment. You can install it with pip: `pip install compel`
+"""
+
+# docstyle-ignore
+BS4_IMPORT_ERROR = """
+{0} requires the Beautiful Soup library but it was not found in your environment. You can install it with pip:
+`pip install beautifulsoup4`. Please note that you may need to restart your runtime after installation.
+"""
+
+# docstyle-ignore
+FTFY_IMPORT_ERROR = """
+{0} requires the ftfy library but it was not found in your environment. Checkout the instructions on the
+installation section: https://github.com/rspeer/python-ftfy/tree/master#installing and follow the ones
+that match your environment. Please note that you may need to restart your runtime after installation.
+"""
+
+# docstyle-ignore
+TORCHSDE_IMPORT_ERROR = """
+{0} requires the torchsde library but it was not found in your environment. You can install it with pip: `pip install torchsde`
+"""
+
+# docstyle-ignore
+INVISIBLE_WATERMARK_IMPORT_ERROR = """
+{0} requires the invisible-watermark library but it was not found in your environment. You can install it with pip: `pip install invisible-watermark>=0.2.0`
+"""
+
+# docstyle-ignore
+PEFT_IMPORT_ERROR = """
+{0} requires the peft library but it was not found in your environment. You can install it with pip: `pip install peft`
+"""
+
+# docstyle-ignore
+SAFETENSORS_IMPORT_ERROR = """
+{0} requires the safetensors library but it was not found in your environment. You can install it with pip: `pip install safetensors`
+"""
+
+# docstyle-ignore
+SENTENCEPIECE_IMPORT_ERROR = """
+{0} requires the sentencepiece library but it was not found in your environment. You can install it with pip: `pip install sentencepiece`
+"""
+
+
+# docstyle-ignore
+BITSANDBYTES_IMPORT_ERROR = """
+{0} requires the bitsandbytes library but it was not found in your environment. You can install it with pip: `pip install bitsandbytes`
+"""
+
+# docstyle-ignore
+IMAGEIO_IMPORT_ERROR = """
+{0} requires the imageio library and ffmpeg but it was not found in your environment. You can install it with pip: `pip install imageio imageio-ffmpeg`
+"""
+
+# docstyle-ignore
+GGUF_IMPORT_ERROR = """
+{0} requires the gguf library but it was not found in your environment. You can install it with pip: `pip install gguf`
+"""
+
+TORCHAO_IMPORT_ERROR = """
+{0} requires the torchao library but it was not found in your environment. You can install it with pip: `pip install
+torchao`
+"""
+
+QUANTO_IMPORT_ERROR = """
+{0} requires the optimum-quanto library but it was not found in your environment. You can install it with pip: `pip
+install optimum-quanto`
+"""
+
+SDNQ_IMPORT_ERROR = """
+{0} requires the sdnq library but it was not found in your environment. You can install it with pip: `pip install sdnq`
+"""
+
+# docstyle-ignore
+PYTORCH_RETINAFACE_IMPORT_ERROR = """
+{0} requires the pytorch_retinaface library but it was not found in your environment. You can install it with pip: `pip install pytorch_retinaface`
+"""
+
+# docstyle-ignore
+BETTER_PROFANITY_IMPORT_ERROR = """
+{0} requires the better_profanity library but it was not found in your environment. You can install it with pip: `pip install better_profanity`
+"""
+
+# docstyle-ignore
+NLTK_IMPORT_ERROR = """
+{0} requires the nltk library but it was not found in your environment. You can install it with pip: `pip install nltk`
+"""
+
+
+TORCH_NEURONX_IMPORT_ERROR = """
+{0} requires the torch_neuronx library (AWS Neuron SDK) but it was not found in your environment. Please install it
+following the AWS Neuron documentation: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/
+"""
+
+BACKENDS_MAPPING = OrderedDict(
+ [
+ ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)),
+ ("inflect", (is_inflect_available, INFLECT_IMPORT_ERROR)),
+ ("onnx", (is_onnx_available, ONNX_IMPORT_ERROR)),
+ ("opencv", (is_opencv_available, OPENCV_IMPORT_ERROR)),
+ ("scipy", (is_scipy_available, SCIPY_IMPORT_ERROR)),
+ ("torch", (is_torch_available, PYTORCH_IMPORT_ERROR)),
+ ("transformers", (is_transformers_available, TRANSFORMERS_IMPORT_ERROR)),
+ ("unidecode", (is_unidecode_available, UNIDECODE_IMPORT_ERROR)),
+ ("librosa", (is_librosa_available, LIBROSA_IMPORT_ERROR)),
+ ("note_seq", (is_note_seq_available, NOTE_SEQ_IMPORT_ERROR)),
+ ("wandb", (is_wandb_available, WANDB_IMPORT_ERROR)),
+ ("tensorboard", (is_tensorboard_available, TENSORBOARD_IMPORT_ERROR)),
+ ("compel", (is_compel_available, COMPEL_IMPORT_ERROR)),
+ ("ftfy", (is_ftfy_available, FTFY_IMPORT_ERROR)),
+ ("torchsde", (is_torchsde_available, TORCHSDE_IMPORT_ERROR)),
+ ("invisible_watermark", (is_invisible_watermark_available, INVISIBLE_WATERMARK_IMPORT_ERROR)),
+ ("peft", (is_peft_available, PEFT_IMPORT_ERROR)),
+ ("safetensors", (is_safetensors_available, SAFETENSORS_IMPORT_ERROR)),
+ ("bitsandbytes", (is_bitsandbytes_available, BITSANDBYTES_IMPORT_ERROR)),
+ ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)),
+ ("imageio", (is_imageio_available, IMAGEIO_IMPORT_ERROR)),
+ ("gguf", (is_gguf_available, GGUF_IMPORT_ERROR)),
+ ("torchao", (is_torchao_available, TORCHAO_IMPORT_ERROR)),
+ ("quanto", (is_optimum_quanto_available, QUANTO_IMPORT_ERROR)),
+ ("sdnq", (is_sdnq_available, SDNQ_IMPORT_ERROR)),
+ ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)),
+ ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)),
+ ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)),
+ ("torch_neuronx", (is_torch_neuronx_available, TORCH_NEURONX_IMPORT_ERROR)),
+ ]
+)
+
+
+def requires_backends(obj, backends):
+ if not isinstance(backends, (list, tuple)):
+ backends = [backends]
+
+ name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__
+ checks = (BACKENDS_MAPPING[backend] for backend in backends)
+ failed = [msg.format(name) for available, msg in checks if not available()]
+ if failed:
+ raise ImportError("".join(failed))
+
+ if name in [
+ "VersatileDiffusionTextToImagePipeline",
+ "VersatileDiffusionPipeline",
+ "VersatileDiffusionDualGuidedPipeline",
+ "StableDiffusionImageVariationPipeline",
+ "UnCLIPPipeline",
+ ] and is_transformers_version("<", "4.25.0"):
+ raise ImportError(
+ f"You need to install `transformers>=4.25` in order to use {name}: \n```\n pip install"
+ " --upgrade transformers \n```"
+ )
+
+ if name in ["StableDiffusionDepth2ImgPipeline", "StableDiffusionPix2PixZeroPipeline"] and is_transformers_version(
+ "<", "4.26.0"
+ ):
+ raise ImportError(
+ f"You need to install `transformers>=4.26` in order to use {name}: \n```\n pip install"
+ " --upgrade transformers \n```"
+ )
+
+
+class DummyObject(type):
+ """
+ Metaclass for the dummy objects. Any class inheriting from it will return the ImportError generated by
+ `requires_backend` each time a user tries to access any method of that class.
+ """
+
+ def __getattr__(cls, key):
+ if key.startswith("_") and key not in ["_load_connected_pipes", "_is_onnx"]:
+ return super().__getattr__(cls, key)
+ requires_backends(cls, cls._backends)
+
+
+# This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319
+def compare_versions(library_or_version: str | Version, operation: str, requirement_version: str):
+ """
+ Compares a library version to some requirement using a given operation.
+
+ Args:
+ library_or_version (`str` or `packaging.version.Version`):
+ A library name or a version to check.
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`.
+ requirement_version (`str`):
+ The version to compare the library version against
+ """
+ if operation not in STR_OPERATION_TO_FUNC.keys():
+ raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+ operation = STR_OPERATION_TO_FUNC[operation]
+ if isinstance(library_or_version, str):
+ library_or_version = parse(importlib_metadata.version(library_or_version))
+ return operation(library_or_version, parse(requirement_version))
+
+
+# This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L338
+@cache
+def is_torch_version(operation: str, version: str):
+ """
+ Compares the current PyTorch version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A string version of PyTorch
+ """
+ return compare_versions(parse(_torch_version), operation, version)
+
+
+@cache
+def is_torch_xla_version(operation: str, version: str):
+ """
+ Compares the current torch_xla version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A string version of torch_xla
+ """
+ if not is_torch_xla_available:
+ return False
+ return compare_versions(parse(_torch_xla_version), operation, version)
+
+
+@cache
+def is_transformers_version(operation: str, version: str):
+ """
+ Compares the current Transformers version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _transformers_available:
+ return False
+ return compare_versions(parse(_transformers_version), operation, version)
+
+
+@cache
+def is_kernels_version(operation: str, version: str):
+ """
+ Compares the current Kernels version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _kernels_available:
+ return False
+ return compare_versions(parse(_kernels_version), operation, version)
+
+
+@cache
+def is_hf_hub_version(operation: str, version: str):
+ """
+ Compares the current Hugging Face Hub version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _hf_hub_available:
+ return False
+ return compare_versions(parse(_hf_hub_version), operation, version)
+
+
+@cache
+def is_accelerate_version(operation: str, version: str):
+ """
+ Compares the current Accelerate version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _accelerate_available:
+ return False
+ return compare_versions(parse(_accelerate_version), operation, version)
+
+
+@cache
+def is_peft_version(operation: str, version: str):
+ """
+ Compares the current PEFT version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _peft_available:
+ return False
+ return compare_versions(parse(_peft_version), operation, version)
+
+
+@cache
+def is_bitsandbytes_version(operation: str, version: str):
+ """
+ Args:
+ Compares the current bitsandbytes version to a given reference with an operation.
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _bitsandbytes_available:
+ return False
+ return compare_versions(parse(_bitsandbytes_version), operation, version)
+
+
+@cache
+def is_gguf_version(operation: str, version: str):
+ """
+ Compares the current Accelerate version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _gguf_available:
+ return False
+ return compare_versions(parse(_gguf_version), operation, version)
+
+
+@cache
+def is_sdnq_version(operation: str, version: str):
+ """
+ Compares the current sdnq version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _sdnq_available:
+ return False
+ return compare_versions(parse(_sdnq_version), operation, version)
+
+
+@cache
+def is_torchao_version(operation: str, version: str):
+ """
+ Compares the current torchao version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _torchao_available:
+ return False
+ return compare_versions(parse(_torchao_version), operation, version)
+
+
+@cache
+def is_optimum_quanto_version(operation: str, version: str):
+ """
+ Compares the current Accelerate version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _optimum_quanto_available:
+ return False
+ return compare_versions(parse(_optimum_quanto_version), operation, version)
+
+
+@cache
+def is_nvidia_modelopt_version(operation: str, version: str):
+ """
+ Compares the current Nvidia ModelOpt version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _nvidia_modelopt_available:
+ return False
+ return compare_versions(parse(_nvidia_modelopt_version), operation, version)
+
+
+@cache
+def is_xformers_version(operation: str, version: str):
+ """
+ Compares the current xformers version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _xformers_available:
+ return False
+ return compare_versions(parse(_xformers_version), operation, version)
+
+
+@cache
+def is_sageattention_version(operation: str, version: str):
+ """
+ Compares the current sageattention version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _sageattention_available:
+ return False
+ return compare_versions(parse(_sageattention_version), operation, version)
+
+
+@cache
+def is_flash_attn_version(operation: str, version: str):
+ """
+ Compares the current flash-attention version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _flash_attn_available:
+ return False
+ return compare_versions(parse(_flash_attn_version), operation, version)
+
+
+@cache
+def is_aiter_version(operation: str, version: str):
+ """
+ Compares the current aiter version to a given reference with an operation.
+
+ Args:
+ operation (`str`):
+ A string representation of an operator, such as `">"` or `"<="`
+ version (`str`):
+ A version string
+ """
+ if not _aiter_available:
+ return False
+ return compare_versions(parse(_aiter_version), operation, version)
+
+
+def get_objects_from_module(module):
+ """
+ Returns a dict of object names and values in a module, while skipping private/internal objects
+
+ Args:
+ module (ModuleType):
+ Module to extract the objects from.
+
+ Returns:
+ dict: Dictionary of object names and corresponding values
+ """
+
+ objects = {}
+ for name in dir(module):
+ if name.startswith("_"):
+ continue
+ objects[name] = getattr(module, name)
+
+ return objects
+
+
+class OptionalDependencyNotAvailable(BaseException):
+ """
+ An error indicating that an optional dependency of Diffusers was not found in the environment.
+ """
+
+
+class _LazyModule(ModuleType):
+ """
+ Module class that surfaces all objects but only performs associated imports when the objects are requested.
+ """
+
+ # Very heavily inspired by optuna.integration._IntegrationModule
+ # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
+ def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):
+ super().__init__(name)
+ self._modules = set(import_structure.keys())
+ self._class_to_module = {}
+ for key, values in import_structure.items():
+ for value in values:
+ self._class_to_module[value] = key
+ # Needed for autocompletion in an IDE
+ self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values()))
+ self.__file__ = module_file
+ self.__spec__ = module_spec
+ self.__path__ = [os.path.dirname(module_file)]
+ self._objects = {} if extra_objects is None else extra_objects
+ self._name = name
+ self._import_structure = import_structure
+
+ # Needed for autocompletion in an IDE
+ def __dir__(self):
+ result = super().__dir__()
+ # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether
+ # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir.
+ for attr in self.__all__:
+ if attr not in result:
+ result.append(attr)
+ return result
+
+ def __getattr__(self, name: str) -> Any:
+ if name in self._objects:
+ return self._objects[name]
+ if name in self._modules:
+ value = self._get_module(name)
+ elif name in self._class_to_module.keys():
+ module = self._get_module(self._class_to_module[name])
+ value = getattr(module, name)
+ else:
+ raise AttributeError(f"module {self.__name__} has no attribute {name}")
+
+ setattr(self, name, value)
+ return value
+
+ def _get_module(self, module_name: str):
+ try:
+ return importlib.import_module("." + module_name, self.__name__)
+ except Exception as e:
+ raise RuntimeError(
+ f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
+ f" traceback):\n{e}"
+ ) from e
+
+ def __reduce__(self):
+ return (self.__class__, (self._name, self.__file__, self._import_structure))
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/loading_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/loading_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..95d27a12616dc91d59d4fdf9dadf76bd55b8b403
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/loading_utils.py
@@ -0,0 +1,163 @@
+import os
+import tempfile
+from typing import Any, Callable
+from urllib.parse import unquote, urlparse
+
+import PIL.Image
+import PIL.ImageOps
+import requests
+
+from .constants import DIFFUSERS_REQUEST_TIMEOUT
+from .import_utils import BACKENDS_MAPPING, is_imageio_available
+
+
+def load_image(
+ image: str | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None
+) -> PIL.Image.Image:
+ """
+ Loads `image` to a PIL Image.
+
+ Args:
+ image (`str` or `PIL.Image.Image`):
+ The image to convert to the PIL Image format.
+ convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):
+ A conversion method to apply to the image after loading it. When set to `None` the image will be converted
+ "RGB".
+
+ Returns:
+ `PIL.Image.Image`:
+ A PIL Image.
+ """
+ if isinstance(image, str):
+ if image.startswith("http://") or image.startswith("https://"):
+ image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)
+ elif os.path.isfile(image):
+ image = PIL.Image.open(image)
+ else:
+ raise ValueError(
+ f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path."
+ )
+ elif isinstance(image, PIL.Image.Image):
+ image = image
+ else:
+ raise ValueError(
+ "Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image."
+ )
+
+ image = PIL.ImageOps.exif_transpose(image)
+
+ if convert_method is not None:
+ image = convert_method(image)
+ else:
+ image = image.convert("RGB")
+
+ return image
+
+
+def load_video(
+ video: str,
+ convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None = None,
+) -> list[PIL.Image.Image]:
+ """
+ Loads `video` to a list of PIL Image.
+
+ Args:
+ video (`str`):
+ A URL or Path to a video to convert to a list of PIL Image format.
+ convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], *optional*):
+ A conversion method to apply to the video after loading it. When set to `None` the images will be converted
+ to "RGB".
+
+ Returns:
+ `list[PIL.Image.Image]`:
+ The video as a list of PIL images.
+ """
+ is_url = video.startswith("http://") or video.startswith("https://")
+ is_file = os.path.isfile(video)
+ was_tempfile_created = False
+
+ if not (is_url or is_file):
+ raise ValueError(
+ f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path."
+ )
+
+ if is_url:
+ response = requests.get(video, stream=True)
+ if response.status_code != 200:
+ raise ValueError(f"Failed to download video. Status code: {response.status_code}")
+
+ parsed_url = urlparse(video)
+ file_name = os.path.basename(unquote(parsed_url.path))
+
+ suffix = os.path.splitext(file_name)[1] or ".mp4"
+ video_path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name
+
+ was_tempfile_created = True
+
+ video_data = response.iter_content(chunk_size=8192)
+ with open(video_path, "wb") as f:
+ for chunk in video_data:
+ f.write(chunk)
+
+ video = video_path
+
+ pil_images = []
+ if video.endswith(".gif"):
+ gif = PIL.Image.open(video)
+ try:
+ while True:
+ pil_images.append(gif.copy())
+ gif.seek(gif.tell() + 1)
+ except EOFError:
+ pass
+
+ else:
+ if is_imageio_available():
+ import imageio
+ else:
+ raise ImportError(BACKENDS_MAPPING["imageio"][1].format("load_video"))
+
+ try:
+ imageio.plugins.ffmpeg.get_exe()
+ except AttributeError:
+ raise AttributeError(
+ "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
+ )
+
+ with imageio.get_reader(video) as reader:
+ # Read all frames
+ for frame in reader:
+ pil_images.append(PIL.Image.fromarray(frame))
+
+ if was_tempfile_created:
+ os.remove(video_path)
+
+ if convert_method is not None:
+ pil_images = convert_method(pil_images)
+
+ return pil_images
+
+
+# Taken from `transformers`.
+def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]:
+ if "." in tensor_name:
+ splits = tensor_name.split(".")
+ for split in splits[:-1]:
+ new_module = getattr(module, split)
+ if new_module is None:
+ raise ValueError(f"{module} has no attribute {split}.")
+ module = new_module
+ tensor_name = splits[-1]
+ return module, tensor_name
+
+
+def get_submodule_by_name(root_module, module_path: str):
+ current = root_module
+ parts = module_path.split(".")
+ for part in parts:
+ if part.isdigit():
+ idx = int(part)
+ current = current[idx] # e.g., for nn.ModuleList or nn.Sequential
+ else:
+ current = getattr(current, part)
+ return current
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/logging.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..a302f6b5dc41b9ade255be5cc57844a71cb8da6c
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/logging.py
@@ -0,0 +1,361 @@
+# coding=utf-8
+# Copyright 2025 Optuna, Hugging Face
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Logging utilities."""
+
+import logging
+import os
+import sys
+import threading
+from logging import (
+ CRITICAL, # NOQA
+ DEBUG, # NOQA
+ ERROR, # NOQA
+ FATAL, # NOQA
+ INFO, # NOQA
+ NOTSET, # NOQA
+ WARN, # NOQA
+ WARNING, # NOQA
+)
+
+from tqdm import auto as tqdm_lib
+
+from .distributed_utils import is_torch_dist_rank_zero
+
+
+_lock = threading.Lock()
+_default_handler: logging.Handler | None = None
+
+log_levels = {
+ "debug": logging.DEBUG,
+ "info": logging.INFO,
+ "warning": logging.WARNING,
+ "error": logging.ERROR,
+ "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.WARNING
+
+_tqdm_active = True
+_rank_zero_filter = None
+
+
+class _RankZeroFilter(logging.Filter):
+ def filter(self, record):
+ # Always allow rank-zero logs, but keep debug-level messages from all ranks for troubleshooting.
+ return is_torch_dist_rank_zero() or record.levelno <= logging.DEBUG
+
+
+def _ensure_rank_zero_filter(logger: logging.Logger) -> None:
+ global _rank_zero_filter
+
+ if _rank_zero_filter is None:
+ _rank_zero_filter = _RankZeroFilter()
+
+ if not any(isinstance(f, _RankZeroFilter) for f in logger.filters):
+ logger.addFilter(_rank_zero_filter)
+
+
+def _get_default_logging_level() -> int:
+ """
+ If DIFFUSERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+ not - fall back to `_default_log_level`
+ """
+ env_level_str = os.getenv("DIFFUSERS_VERBOSITY", None)
+ if env_level_str:
+ if env_level_str in log_levels:
+ return log_levels[env_level_str]
+ else:
+ logging.getLogger().warning(
+ f"Unknown option DIFFUSERS_VERBOSITY={env_level_str}, has to be one of: {', '.join(log_levels.keys())}"
+ )
+ return _default_log_level
+
+
+def _get_library_name() -> str:
+ return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+ return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if _default_handler:
+ # This library has already configured the library root logger.
+ return
+ _default_handler = logging.StreamHandler() # Set sys.stderr as stream.
+
+ if sys.stderr: # only if sys.stderr exists, e.g. when not using pythonw in windows
+ _default_handler.flush = sys.stderr.flush
+
+ # Apply our default configuration to the library root logger.
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.addHandler(_default_handler)
+ library_root_logger.setLevel(_get_default_logging_level())
+ library_root_logger.propagate = False
+ _ensure_rank_zero_filter(library_root_logger)
+
+
+def _reset_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if not _default_handler:
+ return
+
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.removeHandler(_default_handler)
+ library_root_logger.setLevel(logging.NOTSET)
+ _default_handler = None
+
+
+def get_log_levels_dict() -> dict[str, int]:
+ return log_levels
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+ """
+ Return a logger with the specified name.
+
+ This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
+ """
+
+ if name is None:
+ name = _get_library_name()
+
+ _configure_library_root_logger()
+ logger = logging.getLogger(name)
+ _ensure_rank_zero_filter(logger)
+ return logger
+
+
+def get_verbosity() -> int:
+ """
+ Return the current level for the 🤗 Diffusers' root logger as an `int`.
+
+ Returns:
+ `int`:
+ Logging level integers which can be one of:
+
+ - `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+ - `40`: `diffusers.logging.ERROR`
+ - `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+ - `20`: `diffusers.logging.INFO`
+ - `10`: `diffusers.logging.DEBUG`
+
+ """
+
+ _configure_library_root_logger()
+ return _get_library_root_logger().getEffectiveLevel()
+
+
+def set_verbosity(verbosity: int) -> None:
+ """
+ Set the verbosity level for the 🤗 Diffusers' root logger.
+
+ Args:
+ verbosity (`int`):
+ Logging level which can be one of:
+
+ - `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
+ - `diffusers.logging.ERROR`
+ - `diffusers.logging.WARNING` or `diffusers.logging.WARN`
+ - `diffusers.logging.INFO`
+ - `diffusers.logging.DEBUG`
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().setLevel(verbosity)
+
+
+def set_verbosity_info() -> None:
+ """Set the verbosity to the `INFO` level."""
+ return set_verbosity(INFO)
+
+
+def set_verbosity_warning() -> None:
+ """Set the verbosity to the `WARNING` level."""
+ return set_verbosity(WARNING)
+
+
+def set_verbosity_debug() -> None:
+ """Set the verbosity to the `DEBUG` level."""
+ return set_verbosity(DEBUG)
+
+
+def set_verbosity_error() -> None:
+ """Set the verbosity to the `ERROR` level."""
+ return set_verbosity(ERROR)
+
+
+def disable_default_handler() -> None:
+ """Disable the default handler of the 🤗 Diffusers' root logger."""
+
+ _configure_library_root_logger()
+
+ assert _default_handler is not None
+ _get_library_root_logger().removeHandler(_default_handler)
+
+
+def enable_default_handler() -> None:
+ """Enable the default handler of the 🤗 Diffusers' root logger."""
+
+ _configure_library_root_logger()
+
+ assert _default_handler is not None
+ _get_library_root_logger().addHandler(_default_handler)
+
+
+def add_handler(handler: logging.Handler) -> None:
+ """adds a handler to the HuggingFace Diffusers' root logger."""
+
+ _configure_library_root_logger()
+
+ assert handler is not None
+ _get_library_root_logger().addHandler(handler)
+
+
+def remove_handler(handler: logging.Handler) -> None:
+ """removes given handler from the HuggingFace Diffusers' root logger."""
+
+ _configure_library_root_logger()
+
+ assert handler is not None and handler in _get_library_root_logger().handlers
+ _get_library_root_logger().removeHandler(handler)
+
+
+def disable_propagation() -> None:
+ """
+ Disable propagation of the library log outputs. Note that log propagation is disabled by default.
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().propagate = False
+
+
+def enable_propagation() -> None:
+ """
+ Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent
+ double logging if the root logger has been configured.
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().propagate = True
+
+
+def enable_explicit_format() -> None:
+ """
+ Enable explicit formatting for every 🤗 Diffusers' logger. The explicit formatter is as follows:
+ ```
+ [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
+ ```
+ All handlers currently bound to the root logger are affected by this method.
+ """
+ handlers = _get_library_root_logger().handlers
+
+ for handler in handlers:
+ formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s")
+ handler.setFormatter(formatter)
+
+
+def reset_format() -> None:
+ """
+ Resets the formatting for 🤗 Diffusers' loggers.
+
+ All handlers currently bound to the root logger are affected by this method.
+ """
+ handlers = _get_library_root_logger().handlers
+
+ for handler in handlers:
+ handler.setFormatter(None)
+
+
+def warning_advice(self, *args, **kwargs) -> None:
+ """
+ This method is identical to `logger.warning()`, but if env var DIFFUSERS_NO_ADVISORY_WARNINGS=1 is set, this
+ warning will not be printed
+ """
+ no_advisory_warnings = os.getenv("DIFFUSERS_NO_ADVISORY_WARNINGS", False)
+ if no_advisory_warnings:
+ return
+ self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_advice = warning_advice
+
+
+class EmptyTqdm:
+ """Dummy tqdm which doesn't do anything."""
+
+ def __init__(self, *args, **kwargs): # pylint: disable=unused-argument
+ self._iterator = args[0] if args else None
+
+ def __iter__(self):
+ return iter(self._iterator)
+
+ def __getattr__(self, _):
+ """Return empty function."""
+
+ def empty_fn(*args, **kwargs): # pylint: disable=unused-argument
+ return
+
+ return empty_fn
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, type_, value, traceback):
+ return
+
+
+class _tqdm_cls:
+ def __call__(self, *args, **kwargs):
+ if _tqdm_active:
+ return tqdm_lib.tqdm(*args, **kwargs)
+ else:
+ return EmptyTqdm(*args, **kwargs)
+
+ def set_lock(self, *args, **kwargs):
+ self._lock = None
+ if _tqdm_active:
+ return tqdm_lib.tqdm.set_lock(*args, **kwargs)
+
+ def get_lock(self):
+ if _tqdm_active:
+ return tqdm_lib.tqdm.get_lock()
+
+
+tqdm = _tqdm_cls()
+
+
+def is_progress_bar_enabled() -> bool:
+ """Return a boolean indicating whether tqdm progress bars are enabled."""
+ global _tqdm_active
+ return bool(_tqdm_active)
+
+
+def enable_progress_bar() -> None:
+ """Enable tqdm progress bar."""
+ global _tqdm_active
+ _tqdm_active = True
+
+
+def disable_progress_bar() -> None:
+ """Disable tqdm progress bar."""
+ global _tqdm_active
+ _tqdm_active = False
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/model_card_template.md b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/model_card_template.md
new file mode 100644
index 0000000000000000000000000000000000000000..0a2c1d8fe69dd1577ce86a07257099ded968fdcb
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/model_card_template.md
@@ -0,0 +1,24 @@
+---
+{{ card_data }}
+---
+
+
+
+{{ model_description }}
+
+## Intended uses & limitations
+
+#### How to use
+
+```python
+# TODO: add an example code snippet for running this diffusion pipeline
+```
+
+#### Limitations and bias
+
+[TODO: provide examples of latent issues and potential remediations]
+
+## Training details
+
+[TODO: describe the data used to train the model]
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/outputs.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/outputs.py
new file mode 100644
index 0000000000000000000000000000000000000000..915e43d222086f8849ffe5d309ce459cdde0c38d
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/outputs.py
@@ -0,0 +1,134 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Generic utilities
+"""
+
+from collections import OrderedDict
+from dataclasses import fields, is_dataclass
+from typing import Any
+
+import numpy as np
+
+from .import_utils import is_torch_available, is_torch_version
+
+
+def is_tensor(x) -> bool:
+ """
+ Tests if `x` is a `torch.Tensor` or `np.ndarray`.
+ """
+ if is_torch_available():
+ import torch
+
+ if isinstance(x, torch.Tensor):
+ return True
+
+ return isinstance(x, np.ndarray)
+
+
+class BaseOutput(OrderedDict):
+ """
+ Base class for all model outputs as dataclass. Has a `__getitem__` that allows indexing by integer or slice (like a
+ tuple) or strings (like a dictionary) that will ignore the `None` attributes. Otherwise behaves like a regular
+ Python dictionary.
+
+ > [!WARNING] > You can't unpack a [`BaseOutput`] directly. Use the [`~utils.BaseOutput.to_tuple`] method to convert
+ it to a tuple > first.
+ """
+
+ def __init_subclass__(cls) -> None:
+ """Register subclasses as pytree nodes.
+
+ This is necessary to synchronize gradients when using `torch.nn.parallel.DistributedDataParallel` with
+ `static_graph=True` with modules that output `ModelOutput` subclasses.
+ """
+ if is_torch_available():
+ import torch.utils._pytree
+
+ if is_torch_version("<", "2.2"):
+ torch.utils._pytree._register_pytree_node(
+ cls,
+ torch.utils._pytree._dict_flatten,
+ lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+ )
+ else:
+ torch.utils._pytree.register_pytree_node(
+ cls,
+ torch.utils._pytree._dict_flatten,
+ lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+ serialized_type_name=f"{cls.__module__}.{cls.__name__}",
+ )
+
+ def __post_init__(self) -> None:
+ class_fields = fields(self)
+
+ # Safety and consistency checks
+ if not len(class_fields):
+ raise ValueError(f"{self.__class__.__name__} has no fields.")
+
+ first_field = getattr(self, class_fields[0].name)
+ other_fields_are_none = all(getattr(self, field.name) is None for field in class_fields[1:])
+
+ if other_fields_are_none and isinstance(first_field, dict):
+ for key, value in first_field.items():
+ self[key] = value
+ else:
+ for field in class_fields:
+ v = getattr(self, field.name)
+ if v is not None:
+ self[field.name] = v
+
+ def __delitem__(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
+
+ def setdefault(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
+
+ def pop(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
+
+ def update(self, *args, **kwargs):
+ raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
+
+ def __getitem__(self, k: Any) -> Any:
+ if isinstance(k, str):
+ inner_dict = dict(self.items())
+ return inner_dict[k]
+ else:
+ return self.to_tuple()[k]
+
+ def __setattr__(self, name: Any, value: Any) -> None:
+ if name in self.keys() and value is not None:
+ # Don't call self.__setitem__ to avoid recursion errors
+ super().__setitem__(name, value)
+ super().__setattr__(name, value)
+
+ def __setitem__(self, key, value):
+ # Will raise a KeyException if needed
+ super().__setitem__(key, value)
+ # Don't call self.__setattr__ to avoid recursion errors
+ super().__setattr__(key, value)
+
+ def __reduce__(self):
+ if not is_dataclass(self):
+ return super().__reduce__()
+ callable, _args, *remaining = super().__reduce__()
+ args = tuple(getattr(self, field.name) for field in fields(self))
+ return callable, args, *remaining
+
+ def to_tuple(self) -> tuple[Any, ...]:
+ """
+ Convert self to a tuple containing all the attributes/keys that are not `None`.
+ """
+ return tuple(self[k] for k in self.keys())
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/peft_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/peft_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..42d6b1f0fd3903b2233fbea686795fcb2bead1d4
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/peft_utils.py
@@ -0,0 +1,425 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+PEFT utilities: Utilities related to peft library
+"""
+
+import collections
+import functools
+import importlib
+
+from packaging import version
+
+from . import logging
+from .import_utils import is_peft_available, is_peft_version, is_torch_available
+from .torch_utils import empty_device_cache
+
+
+logger = logging.get_logger(__name__)
+
+if is_torch_available():
+ import torch
+
+
+def recurse_remove_peft_layers(model):
+ r"""
+ Recursively replace all instances of `LoraLayer` with corresponding new layers in `model`.
+ """
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ has_base_layer_pattern = False
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ has_base_layer_pattern = hasattr(module, "base_layer")
+ break
+
+ if has_base_layer_pattern:
+ from peft.utils import _get_submodules
+
+ key_list = [key for key, _ in model.named_modules() if "lora" not in key]
+ for key in key_list:
+ try:
+ parent, target, target_name = _get_submodules(model, key)
+ except AttributeError:
+ continue
+ if hasattr(target, "base_layer"):
+ setattr(parent, target_name, target.get_base_layer())
+ else:
+ # This is for backwards compatibility with PEFT <= 0.6.2.
+ # TODO can be removed once that PEFT version is no longer supported.
+ from peft.tuners.lora import LoraLayer
+
+ for name, module in model.named_children():
+ if len(list(module.children())) > 0:
+ ## compound module, go inside it
+ recurse_remove_peft_layers(module)
+
+ module_replaced = False
+
+ if isinstance(module, LoraLayer) and isinstance(module, torch.nn.Linear):
+ new_module = torch.nn.Linear(
+ module.in_features,
+ module.out_features,
+ bias=module.bias is not None,
+ ).to(module.weight.device)
+ new_module.weight = module.weight
+ if module.bias is not None:
+ new_module.bias = module.bias
+
+ module_replaced = True
+ elif isinstance(module, LoraLayer) and isinstance(module, torch.nn.Conv2d):
+ new_module = torch.nn.Conv2d(
+ module.in_channels,
+ module.out_channels,
+ module.kernel_size,
+ module.stride,
+ module.padding,
+ module.dilation,
+ module.groups,
+ ).to(module.weight.device)
+
+ new_module.weight = module.weight
+ if module.bias is not None:
+ new_module.bias = module.bias
+
+ module_replaced = True
+
+ if module_replaced:
+ setattr(model, name, new_module)
+ del module
+
+ empty_device_cache()
+ return model
+
+
+def scale_lora_layers(model, weight):
+ """
+ Adjust the weightage given to the LoRA layers of the model.
+
+ Args:
+ model (`torch.nn.Module`):
+ The model to scale.
+ weight (`float`):
+ The weight to be given to the LoRA layers.
+ """
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ if weight == 1.0:
+ return
+
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ module.scale_layer(weight)
+
+
+def unscale_lora_layers(model, weight: float | None = None):
+ """
+ Removes the previously passed weight given to the LoRA layers of the model.
+
+ Args:
+ model (`torch.nn.Module`):
+ The model to scale.
+ weight (`float`, *optional*):
+ The weight to be given to the LoRA layers. If no scale is passed the scale of the lora layer will be
+ re-initialized to the correct value. If 0.0 is passed, we will re-initialize the scale with the correct
+ value.
+ """
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ if weight is None or weight == 1.0:
+ return
+
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ if weight != 0:
+ module.unscale_layer(weight)
+ else:
+ for adapter_name in module.active_adapters:
+ # if weight == 0 unscale should re-set the scale to the original value.
+ module.set_scale(adapter_name, 1.0)
+
+
+def get_peft_kwargs(
+ rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, model_state_dict=None, adapter_name=None
+):
+ rank_pattern = {}
+ alpha_pattern = {}
+ r = lora_alpha = list(rank_dict.values())[0]
+
+ if len(set(rank_dict.values())) > 1:
+ # get the rank occurring the most number of times
+ r = collections.Counter(rank_dict.values()).most_common()[0][0]
+
+ # for modules with rank different from the most occurring rank, add it to the `rank_pattern`
+ rank_pattern = dict(filter(lambda x: x[1] != r, rank_dict.items()))
+ rank_pattern = {k.split(".lora_B.")[0]: v for k, v in rank_pattern.items()}
+
+ if network_alpha_dict is not None and len(network_alpha_dict) > 0:
+ if len(set(network_alpha_dict.values())) > 1:
+ # get the alpha occurring the most number of times
+ lora_alpha = collections.Counter(network_alpha_dict.values()).most_common()[0][0]
+
+ # for modules with alpha different from the most occurring alpha, add it to the `alpha_pattern`
+ alpha_pattern = dict(filter(lambda x: x[1] != lora_alpha, network_alpha_dict.items()))
+ if is_unet:
+ alpha_pattern = {
+ ".".join(k.split(".lora_A.")[0].split(".")).replace(".alpha", ""): v
+ for k, v in alpha_pattern.items()
+ }
+ else:
+ alpha_pattern = {".".join(k.split(".down.")[0].split(".")[:-1]): v for k, v in alpha_pattern.items()}
+ else:
+ lora_alpha = set(network_alpha_dict.values()).pop()
+
+ target_modules = list({name.split(".lora")[0] for name in peft_state_dict.keys()})
+ use_dora = any("lora_magnitude_vector" in k for k in peft_state_dict)
+ # for now we know that the "bias" keys are only associated with `lora_B`.
+ lora_bias = any("lora_B" in k and k.endswith(".bias") for k in peft_state_dict)
+
+ lora_config_kwargs = {
+ "r": r,
+ "lora_alpha": lora_alpha,
+ "rank_pattern": rank_pattern,
+ "alpha_pattern": alpha_pattern,
+ "target_modules": target_modules,
+ "use_dora": use_dora,
+ "lora_bias": lora_bias,
+ }
+
+ return lora_config_kwargs
+
+
+def get_adapter_name(model):
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ return f"default_{len(module.r)}"
+ return "default_0"
+
+
+def set_adapter_layers(model, enabled=True):
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ # The recent version of PEFT needs to call `enable_adapters` instead
+ if hasattr(module, "enable_adapters"):
+ module.enable_adapters(enabled=enabled)
+ else:
+ module.disable_adapters = not enabled
+
+
+def delete_adapter_layers(model, adapter_name):
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ for module in model.modules():
+ if isinstance(module, BaseTunerLayer):
+ if hasattr(module, "delete_adapter"):
+ module.delete_adapter(adapter_name)
+ else:
+ raise ValueError(
+ "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1"
+ )
+
+ # For transformers integration - we need to pop the adapter from the config
+ if getattr(model, "_hf_peft_config_loaded", False) and hasattr(model, "peft_config"):
+ model.peft_config.pop(adapter_name, None)
+ # In case all adapters are deleted, we need to delete the config
+ # and make sure to set the flag to False
+ if len(model.peft_config) == 0:
+ del model.peft_config
+ model._hf_peft_config_loaded = None
+
+
+def set_weights_and_activate_adapters(model, adapter_names, weights):
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ def get_module_weight(weight_for_adapter, module_name):
+ if not isinstance(weight_for_adapter, dict):
+ # If weight_for_adapter is a single number, always return it.
+ return weight_for_adapter
+
+ for layer_name, weight_ in weight_for_adapter.items():
+ if layer_name in module_name:
+ return weight_
+
+ parts = module_name.split(".")
+ # e.g. key = "down_blocks.1.attentions.0"
+ key = f"{parts[0]}.{parts[1]}.attentions.{parts[3]}"
+ block_weight = weight_for_adapter.get(key, 1.0)
+
+ return block_weight
+
+ for module_name, module in model.named_modules():
+ if isinstance(module, BaseTunerLayer):
+ # For backward compatibility with previous PEFT versions, set multiple active adapters
+ if hasattr(module, "set_adapter"):
+ module.set_adapter(adapter_names)
+ else:
+ module.active_adapter = adapter_names
+
+ # Set the scaling weight for each adapter for this module
+ for adapter_name, weight in zip(adapter_names, weights):
+ module.set_scale(adapter_name, get_module_weight(weight, module_name))
+
+
+def apply_lora_scale(kwargs_name: str = "joint_attention_kwargs"):
+ """
+ Decorator to automatically handle LoRA layer scaling/unscaling in forward methods.
+
+ This decorator extracts the `lora_scale` from the specified kwargs parameter, applies scaling before the forward
+ pass, and ensures unscaling happens after, even if an exception occurs.
+
+ Args:
+ kwargs_name (`str`, defaults to `"joint_attention_kwargs"`):
+ The name of the keyword argument that contains the LoRA scale. Common values include
+ "joint_attention_kwargs", "attention_kwargs", "cross_attention_kwargs", etc.
+ """
+
+ def decorator(forward_fn):
+ @functools.wraps(forward_fn)
+ def wrapper(self, *args, **kwargs):
+ from . import USE_PEFT_BACKEND
+
+ lora_scale = 1.0
+ attention_kwargs = kwargs.get(kwargs_name)
+
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ kwargs[kwargs_name] = attention_kwargs
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+
+ if not USE_PEFT_BACKEND and lora_scale != 1.0:
+ logger.warning(
+ f"Passing `scale` via `{kwargs_name}` when not using the PEFT backend is ineffective."
+ )
+
+ # Apply LoRA scaling if using PEFT backend
+ if USE_PEFT_BACKEND:
+ scale_lora_layers(self, lora_scale)
+
+ try:
+ # Execute the forward pass
+ result = forward_fn(self, *args, **kwargs)
+ return result
+ finally:
+ # Always unscale, even if forward pass raises an exception
+ if USE_PEFT_BACKEND:
+ unscale_lora_layers(self, lora_scale)
+
+ return wrapper
+
+ return decorator
+
+
+def check_peft_version(min_version: str) -> None:
+ r"""
+ Checks if the version of PEFT is compatible.
+
+ Args:
+ version (`str`):
+ The version of PEFT to check against.
+ """
+ if not is_peft_available():
+ raise ValueError("PEFT is not installed. Please install it with `pip install peft`")
+
+ is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) > version.parse(min_version)
+
+ if not is_peft_version_compatible:
+ raise ValueError(
+ f"The version of PEFT you are using is not compatible, please use a version that is greater"
+ f" than {min_version}"
+ )
+
+
+def _create_lora_config(
+ state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None
+):
+ from peft import LoraConfig
+
+ if metadata is not None:
+ lora_config_kwargs = metadata
+ else:
+ lora_config_kwargs = get_peft_kwargs(
+ rank_pattern_dict,
+ network_alpha_dict=network_alphas,
+ peft_state_dict=state_dict,
+ is_unet=is_unet,
+ model_state_dict=model_state_dict,
+ adapter_name=adapter_name,
+ )
+
+ _maybe_raise_error_for_ambiguous_keys(lora_config_kwargs)
+
+ # Version checks for DoRA and lora_bias
+ if "use_dora" in lora_config_kwargs and lora_config_kwargs["use_dora"]:
+ if is_peft_version("<", "0.9.0"):
+ raise ValueError("DoRA requires PEFT >= 0.9.0. Please upgrade.")
+
+ if "lora_bias" in lora_config_kwargs and lora_config_kwargs["lora_bias"]:
+ if is_peft_version("<=", "0.13.2"):
+ raise ValueError("lora_bias requires PEFT >= 0.14.0. Please upgrade.")
+
+ try:
+ return LoraConfig(**lora_config_kwargs)
+ except TypeError as e:
+ raise TypeError("`LoraConfig` class could not be instantiated.") from e
+
+
+def _maybe_raise_error_for_ambiguous_keys(config):
+ rank_pattern = config["rank_pattern"].copy()
+ target_modules = config["target_modules"]
+
+ for key in list(rank_pattern.keys()):
+ # try to detect ambiguity
+ # `target_modules` can also be a str, in which case this loop would loop
+ # over the chars of the str. The technically correct way to match LoRA keys
+ # in PEFT is to use LoraModel._check_target_module_exists (lora_config, key).
+ # But this cuts it for now.
+ exact_matches = [mod for mod in target_modules if mod == key]
+ substring_matches = [mod for mod in target_modules if key in mod and mod != key]
+
+ if exact_matches and substring_matches:
+ if is_peft_version("<", "0.14.1"):
+ raise ValueError(
+ "There are ambiguous keys present in this LoRA. To load it, please update your `peft` installation - `pip install -U peft`."
+ )
+
+
+def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name):
+ warn_msg = ""
+ if incompatible_keys is not None:
+ # Check only for unexpected keys.
+ unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+ if unexpected_keys:
+ lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k]
+ if lora_unexpected_keys:
+ warn_msg = (
+ f"Loading adapter weights from state_dict led to unexpected keys found in the model:"
+ f" {', '.join(lora_unexpected_keys)}. "
+ )
+
+ # Filter missing keys specific to the current adapter.
+ missing_keys = getattr(incompatible_keys, "missing_keys", None)
+ if missing_keys:
+ lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k]
+ if lora_missing_keys:
+ warn_msg += (
+ f"Loading adapter weights from state_dict led to missing keys in the model:"
+ f" {', '.join(lora_missing_keys)}."
+ )
+
+ if warn_msg:
+ logger.warning(warn_msg)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/pil_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/pil_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..54db98dafd6f8975508f584da5a146969894c007
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/pil_utils.py
@@ -0,0 +1,65 @@
+import PIL.Image
+import PIL.ImageOps
+from packaging import version
+from PIL import Image
+
+
+if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
+ PIL_INTERPOLATION = {
+ "linear": PIL.Image.Resampling.BILINEAR,
+ "bilinear": PIL.Image.Resampling.BILINEAR,
+ "bicubic": PIL.Image.Resampling.BICUBIC,
+ "lanczos": PIL.Image.Resampling.LANCZOS,
+ "nearest": PIL.Image.Resampling.NEAREST,
+ }
+else:
+ PIL_INTERPOLATION = {
+ "linear": PIL.Image.LINEAR,
+ "bilinear": PIL.Image.BILINEAR,
+ "bicubic": PIL.Image.BICUBIC,
+ "lanczos": PIL.Image.LANCZOS,
+ "nearest": PIL.Image.NEAREST,
+ }
+
+
+def pt_to_pil(images):
+ """
+ Convert a torch image to a PIL image.
+ """
+ images = (images / 2 + 0.5).clamp(0, 1)
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+ images = numpy_to_pil(images)
+ return images
+
+
+def numpy_to_pil(images):
+ """
+ Convert a numpy image or a batch of images to a PIL image.
+ """
+ if images.ndim == 3:
+ images = images[None, ...]
+ images = (images * 255).round().astype("uint8")
+ if images.shape[-1] == 1:
+ # special case for grayscale (single channel) images
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+ else:
+ pil_images = [Image.fromarray(image) for image in images]
+
+ return pil_images
+
+
+def make_image_grid(images: list[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image:
+ """
+ Prepares a single grid of images. Useful for visualization purposes.
+ """
+ assert len(images) == rows * cols
+
+ if resize is not None:
+ images = [img.resize((resize, resize)) for img in images]
+
+ w, h = images[0].size
+ grid = Image.new("RGB", size=(cols * w, rows * h))
+
+ for i, img in enumerate(images):
+ grid.paste(img, box=(i % cols * w, i // cols * h))
+ return grid
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/remote_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/remote_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..69a6226274213e9607c3188b742fee0deb842069
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/remote_utils.py
@@ -0,0 +1,427 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import io
+import json
+from typing import Literal, cast
+
+import requests
+
+from .deprecation_utils import deprecate
+from .import_utils import is_safetensors_available, is_torch_available
+
+
+if is_torch_available():
+ import torch
+
+ from ..image_processor import VaeImageProcessor
+ from ..video_processor import VideoProcessor
+
+ if is_safetensors_available():
+ import safetensors.torch
+
+ DTYPE_MAP = {
+ "float16": torch.float16,
+ "float32": torch.float32,
+ "bfloat16": torch.bfloat16,
+ "uint8": torch.uint8,
+ }
+
+
+from PIL import Image
+
+
+def detect_image_type(data: bytes) -> str:
+ if data.startswith(b"\xff\xd8"):
+ return "jpeg"
+ elif data.startswith(b"\x89PNG\r\n\x1a\n"):
+ return "png"
+ elif data.startswith(b"GIF87a") or data.startswith(b"GIF89a"):
+ return "gif"
+ elif data.startswith(b"BM"):
+ return "bmp"
+ return "unknown"
+
+
+def check_inputs_decode(
+ endpoint: str,
+ tensor: "torch.Tensor",
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
+ do_scaling: bool = True,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+ output_type: Literal["mp4", "pil", "pt"] = "pil",
+ return_type: Literal["mp4", "pil", "pt"] = "pil",
+ image_format: Literal["png", "jpg"] = "jpg",
+ partial_postprocess: bool = False,
+ input_tensor_type: Literal["binary"] = "binary",
+ output_tensor_type: Literal["binary"] = "binary",
+ height: int | None = None,
+ width: int | None = None,
+):
+ if tensor.ndim == 3 and height is None and width is None:
+ raise ValueError("`height` and `width` required for packed latents.")
+ if (
+ output_type == "pt"
+ and return_type == "pil"
+ and not partial_postprocess
+ and not isinstance(processor, (VaeImageProcessor, VideoProcessor))
+ ):
+ raise ValueError("`processor` is required.")
+ if do_scaling and scaling_factor is None:
+ deprecate(
+ "do_scaling",
+ "1.0.0",
+ "`do_scaling` is deprecated, pass `scaling_factor` and `shift_factor` if required.",
+ standard_warn=False,
+ )
+
+
+def postprocess_decode(
+ response: requests.Response,
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
+ output_type: Literal["mp4", "pil", "pt"] = "pil",
+ return_type: Literal["mp4", "pil", "pt"] = "pil",
+ partial_postprocess: bool = False,
+):
+ if output_type == "pt" or (output_type == "pil" and processor is not None):
+ output_tensor = response.content
+ parameters = response.headers
+ shape = json.loads(parameters["shape"])
+ dtype = parameters["dtype"]
+ torch_dtype = DTYPE_MAP[dtype]
+ output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape)
+ if output_type == "pt":
+ if partial_postprocess:
+ if return_type == "pil":
+ output = [Image.fromarray(image.numpy()) for image in output_tensor]
+ if len(output) == 1:
+ output = output[0]
+ elif return_type == "pt":
+ output = output_tensor
+ else:
+ if processor is None or return_type == "pt":
+ output = output_tensor
+ else:
+ if isinstance(processor, VideoProcessor):
+ output = cast(
+ list[Image.Image],
+ processor.postprocess_video(output_tensor, output_type="pil")[0],
+ )
+ else:
+ output = cast(
+ Image.Image,
+ processor.postprocess(output_tensor, output_type="pil")[0],
+ )
+ elif output_type == "pil" and return_type == "pil" and processor is None:
+ output = Image.open(io.BytesIO(response.content)).convert("RGB")
+ detected_format = detect_image_type(response.content)
+ output.format = detected_format
+ elif output_type == "pil" and processor is not None:
+ if return_type == "pil":
+ output = [
+ Image.fromarray(image)
+ for image in (output_tensor.permute(0, 2, 3, 1).float().numpy() * 255).round().astype("uint8")
+ ]
+ elif return_type == "pt":
+ output = output_tensor
+ elif output_type == "mp4" and return_type == "mp4":
+ output = response.content
+ return output
+
+
+def prepare_decode(
+ tensor: "torch.Tensor",
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
+ do_scaling: bool = True,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+ output_type: Literal["mp4", "pil", "pt"] = "pil",
+ image_format: Literal["png", "jpg"] = "jpg",
+ partial_postprocess: bool = False,
+ height: int | None = None,
+ width: int | None = None,
+):
+ headers = {}
+ parameters = {
+ "image_format": image_format,
+ "output_type": output_type,
+ "partial_postprocess": partial_postprocess,
+ "shape": list(tensor.shape),
+ "dtype": str(tensor.dtype).split(".")[-1],
+ }
+ if do_scaling and scaling_factor is not None:
+ parameters["scaling_factor"] = scaling_factor
+ if do_scaling and shift_factor is not None:
+ parameters["shift_factor"] = shift_factor
+ if do_scaling and scaling_factor is None:
+ parameters["do_scaling"] = do_scaling
+ elif do_scaling and scaling_factor is None and shift_factor is None:
+ parameters["do_scaling"] = do_scaling
+ if height is not None and width is not None:
+ parameters["height"] = height
+ parameters["width"] = width
+ headers["Content-Type"] = "tensor/binary"
+ headers["Accept"] = "tensor/binary"
+ if output_type == "pil" and image_format == "jpg" and processor is None:
+ headers["Accept"] = "image/jpeg"
+ elif output_type == "pil" and image_format == "png" and processor is None:
+ headers["Accept"] = "image/png"
+ elif output_type == "mp4":
+ headers["Accept"] = "text/plain"
+ tensor_data = safetensors.torch._to_ndarray(tensor)[0].tobytes()
+ return {"data": tensor_data, "params": parameters, "headers": headers}
+
+
+def remote_decode(
+ endpoint: str,
+ tensor: "torch.Tensor",
+ processor: "VaeImageProcessor" | "VideoProcessor" | None = None,
+ do_scaling: bool = True,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+ output_type: Literal["mp4", "pil", "pt"] = "pil",
+ return_type: Literal["mp4", "pil", "pt"] = "pil",
+ image_format: Literal["png", "jpg"] = "jpg",
+ partial_postprocess: bool = False,
+ input_tensor_type: Literal["binary"] = "binary",
+ output_tensor_type: Literal["binary"] = "binary",
+ height: int | None = None,
+ width: int | None = None,
+) -> Image.Image | list[Image.Image] | bytes | "torch.Tensor":
+ """
+ Hugging Face Hybrid Inference that allow running VAE decode remotely.
+
+ Args:
+ endpoint (`str`):
+ Endpoint for Remote Decode.
+ tensor (`torch.Tensor`):
+ Tensor to be decoded.
+ processor (`VaeImageProcessor` or `VideoProcessor`, *optional*):
+ Used with `return_type="pt"`, and `return_type="pil"` for Video models.
+ do_scaling (`bool`, default `True`, *optional*):
+ **DEPRECATED**. **pass `scaling_factor`/`shift_factor` instead.** **still set
+ do_scaling=None/do_scaling=False for no scaling until option is removed** When `True` scaling e.g. `latents
+ / self.vae.config.scaling_factor` is applied remotely. If `False`, input must be passed with scaling
+ applied.
+ scaling_factor (`float`, *optional*):
+ Scaling is applied when passed e.g. [`latents /
+ self.vae.config.scaling_factor`](https://github.com/huggingface/diffusers/blob/7007febae5cff000d4df9059d9cf35133e8b2ca9/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L1083C37-L1083C77).
+ - SD v1: 0.18215
+ - SD XL: 0.13025
+ - Flux: 0.3611
+ If `None`, input must be passed with scaling applied.
+ shift_factor (`float`, *optional*):
+ Shift is applied when passed e.g. `latents + self.vae.config.shift_factor`.
+ - Flux: 0.1159
+ If `None`, input must be passed with scaling applied.
+ output_type (`"mp4"` or `"pil"` or `"pt", default `"pil"):
+ **Endpoint** output type. Subject to change. Report feedback on preferred type.
+
+ `"mp4": Supported by video models. Endpoint returns `bytes` of video. `"pil"`: Supported by image and video
+ models.
+ Image models: Endpoint returns `bytes` of an image in `image_format`. Video models: Endpoint returns
+ `torch.Tensor` with partial `postprocessing` applied.
+ Requires `processor` as a flag (any `None` value will work).
+ `"pt"`: Support by image and video models. Endpoint returns `torch.Tensor`.
+ With `partial_postprocess=True` the tensor is postprocessed `uint8` image tensor.
+
+ Recommendations:
+ `"pt"` with `partial_postprocess=True` is the smallest transfer for full quality. `"pt"` with
+ `partial_postprocess=False` is the most compatible with third party code. `"pil"` with
+ `image_format="jpg"` is the smallest transfer overall.
+
+ return_type (`"mp4"` or `"pil"` or `"pt", default `"pil"):
+ **Function** return type.
+
+ `"mp4": Function returns `bytes` of video. `"pil"`: Function returns `PIL.Image.Image`.
+ With `output_type="pil" no further processing is applied. With `output_type="pt" a `PIL.Image.Image` is
+ created.
+ `partial_postprocess=False` `processor` is required. `partial_postprocess=True` `processor` is
+ **not** required.
+ `"pt"`: Function returns `torch.Tensor`.
+ `processor` is **not** required. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without
+ denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized.
+
+ image_format (`"png"` or `"jpg"`, default `jpg`):
+ Used with `output_type="pil"`. Endpoint returns `jpg` or `png`.
+
+ partial_postprocess (`bool`, default `False`):
+ Used with `output_type="pt"`. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without
+ denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized.
+
+ input_tensor_type (`"binary"`, default `"binary"`):
+ Tensor transfer type.
+
+ output_tensor_type (`"binary"`, default `"binary"`):
+ Tensor transfer type.
+
+ height (`int`, **optional**):
+ Required for `"packed"` latents.
+
+ width (`int`, **optional**):
+ Required for `"packed"` latents.
+
+ Returns:
+ output (`Image.Image` or `list[Image.Image]` or `bytes` or `torch.Tensor`).
+ """
+ if input_tensor_type == "base64":
+ deprecate(
+ "input_tensor_type='base64'",
+ "1.0.0",
+ "input_tensor_type='base64' is deprecated. Using `binary`.",
+ standard_warn=False,
+ )
+ input_tensor_type = "binary"
+ if output_tensor_type == "base64":
+ deprecate(
+ "output_tensor_type='base64'",
+ "1.0.0",
+ "output_tensor_type='base64' is deprecated. Using `binary`.",
+ standard_warn=False,
+ )
+ output_tensor_type = "binary"
+ check_inputs_decode(
+ endpoint,
+ tensor,
+ processor,
+ do_scaling,
+ scaling_factor,
+ shift_factor,
+ output_type,
+ return_type,
+ image_format,
+ partial_postprocess,
+ input_tensor_type,
+ output_tensor_type,
+ height,
+ width,
+ )
+ kwargs = prepare_decode(
+ tensor=tensor,
+ processor=processor,
+ do_scaling=do_scaling,
+ scaling_factor=scaling_factor,
+ shift_factor=shift_factor,
+ output_type=output_type,
+ image_format=image_format,
+ partial_postprocess=partial_postprocess,
+ height=height,
+ width=width,
+ )
+ response = requests.post(endpoint, **kwargs)
+ if not response.ok:
+ raise RuntimeError(response.json())
+ output = postprocess_decode(
+ response=response,
+ processor=processor,
+ output_type=output_type,
+ return_type=return_type,
+ partial_postprocess=partial_postprocess,
+ )
+ return output
+
+
+def check_inputs_encode(
+ endpoint: str,
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+):
+ pass
+
+
+def postprocess_encode(
+ response: requests.Response,
+):
+ output_tensor = response.content
+ parameters = response.headers
+ shape = json.loads(parameters["shape"])
+ dtype = parameters["dtype"]
+ torch_dtype = DTYPE_MAP[dtype]
+ output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape)
+ return output_tensor
+
+
+def prepare_encode(
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+):
+ headers = {}
+ parameters = {}
+ if scaling_factor is not None:
+ parameters["scaling_factor"] = scaling_factor
+ if shift_factor is not None:
+ parameters["shift_factor"] = shift_factor
+ if isinstance(image, torch.Tensor):
+ data = safetensors.torch._to_ndarray(image.contiguous())[0].tobytes()
+ parameters["shape"] = list(image.shape)
+ parameters["dtype"] = str(image.dtype).split(".")[-1]
+ else:
+ buffer = io.BytesIO()
+ image.save(buffer, format="PNG")
+ data = buffer.getvalue()
+ return {"data": data, "params": parameters, "headers": headers}
+
+
+def remote_encode(
+ endpoint: str,
+ image: "torch.Tensor" | Image.Image,
+ scaling_factor: float | None = None,
+ shift_factor: float | None = None,
+) -> "torch.Tensor":
+ """
+ Hugging Face Hybrid Inference that allow running VAE encode remotely.
+
+ Args:
+ endpoint (`str`):
+ Endpoint for Remote Decode.
+ image (`torch.Tensor` or `PIL.Image.Image`):
+ Image to be encoded.
+ scaling_factor (`float`, *optional*):
+ Scaling is applied when passed e.g. [`latents * self.vae.config.scaling_factor`].
+ - SD v1: 0.18215
+ - SD XL: 0.13025
+ - Flux: 0.3611
+ If `None`, input must be passed with scaling applied.
+ shift_factor (`float`, *optional*):
+ Shift is applied when passed e.g. `latents - self.vae.config.shift_factor`.
+ - Flux: 0.1159
+ If `None`, input must be passed with scaling applied.
+
+ Returns:
+ output (`torch.Tensor`).
+ """
+ check_inputs_encode(
+ endpoint,
+ image,
+ scaling_factor,
+ shift_factor,
+ )
+ kwargs = prepare_encode(
+ image=image,
+ scaling_factor=scaling_factor,
+ shift_factor=shift_factor,
+ )
+ response = requests.post(endpoint, **kwargs)
+ if not response.ok:
+ raise RuntimeError(response.json())
+ output = postprocess_encode(
+ response=response,
+ )
+ return output
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/source_code_parsing_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/source_code_parsing_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..2943984e85e68a2db55622caf1817833df63a057
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/source_code_parsing_utils.py
@@ -0,0 +1,52 @@
+import ast
+import importlib
+import inspect
+import textwrap
+
+
+class ReturnNameVisitor(ast.NodeVisitor):
+ """Thanks to ChatGPT for pairing."""
+
+ def __init__(self):
+ self.return_names = []
+
+ def visit_Return(self, node):
+ # Check if the return value is a tuple.
+ if isinstance(node.value, ast.Tuple):
+ for elt in node.value.elts:
+ if isinstance(elt, ast.Name):
+ self.return_names.append(elt.id)
+ else:
+ try:
+ self.return_names.append(ast.unparse(elt))
+ except Exception:
+ self.return_names.append(str(elt))
+ else:
+ if isinstance(node.value, ast.Name):
+ self.return_names.append(node.value.id)
+ else:
+ try:
+ self.return_names.append(ast.unparse(node.value))
+ except Exception:
+ self.return_names.append(str(node.value))
+ self.generic_visit(node)
+
+ def _determine_parent_module(self, cls):
+ from diffusers import DiffusionPipeline
+ from diffusers.models.modeling_utils import ModelMixin
+
+ if issubclass(cls, DiffusionPipeline):
+ return "pipelines"
+ elif issubclass(cls, ModelMixin):
+ return "models"
+ else:
+ raise NotImplementedError
+
+ def get_ast_tree(self, cls, attribute_name="encode_prompt"):
+ parent_module_name = self._determine_parent_module(cls)
+ main_module = importlib.import_module(f"diffusers.{parent_module_name}")
+ current_cls_module = getattr(main_module, cls.__name__)
+ source_code = inspect.getsource(getattr(current_cls_module, attribute_name))
+ source_code = textwrap.dedent(source_code)
+ tree = ast.parse(source_code)
+ return tree
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/state_dict_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/state_dict_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5f0f2349368b4b09057fe472ca78a418481a65a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/state_dict_utils.py
@@ -0,0 +1,545 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+State dict utilities: utility methods for converting state dicts easily
+"""
+
+import enum
+import json
+
+from .import_utils import is_torch_available
+from .logging import get_logger
+
+
+if is_torch_available():
+ import torch
+
+
+logger = get_logger(__name__)
+
+
+class StateDictType(enum.Enum):
+ """
+ The mode to use when converting state dicts.
+ """
+
+ DIFFUSERS_OLD = "diffusers_old"
+ KOHYA_SS = "kohya_ss"
+ PEFT = "peft"
+ DIFFUSERS = "diffusers"
+
+
+# We need to define a proper mapping for Unet since it uses different output keys than text encoder
+# e.g. to_q_lora -> q_proj / to_q
+UNET_TO_DIFFUSERS = {
+ ".to_out_lora.up": ".to_out.0.lora_B",
+ ".to_out_lora.down": ".to_out.0.lora_A",
+ ".to_q_lora.down": ".to_q.lora_A",
+ ".to_q_lora.up": ".to_q.lora_B",
+ ".to_k_lora.down": ".to_k.lora_A",
+ ".to_k_lora.up": ".to_k.lora_B",
+ ".to_v_lora.down": ".to_v.lora_A",
+ ".to_v_lora.up": ".to_v.lora_B",
+ ".lora.up": ".lora_B",
+ ".lora.down": ".lora_A",
+ ".to_out.lora_magnitude_vector": ".to_out.0.lora_magnitude_vector",
+}
+
+CONTROL_LORA_TO_DIFFUSERS = {
+ ".to_q.down": ".to_q.lora_A.weight",
+ ".to_q.up": ".to_q.lora_B.weight",
+ ".to_k.down": ".to_k.lora_A.weight",
+ ".to_k.up": ".to_k.lora_B.weight",
+ ".to_v.down": ".to_v.lora_A.weight",
+ ".to_v.up": ".to_v.lora_B.weight",
+ ".to_out.0.down": ".to_out.0.lora_A.weight",
+ ".to_out.0.up": ".to_out.0.lora_B.weight",
+ ".ff.net.0.proj.down": ".ff.net.0.proj.lora_A.weight",
+ ".ff.net.0.proj.up": ".ff.net.0.proj.lora_B.weight",
+ ".ff.net.2.down": ".ff.net.2.lora_A.weight",
+ ".ff.net.2.up": ".ff.net.2.lora_B.weight",
+ ".proj_in.down": ".proj_in.lora_A.weight",
+ ".proj_in.up": ".proj_in.lora_B.weight",
+ ".proj_out.down": ".proj_out.lora_A.weight",
+ ".proj_out.up": ".proj_out.lora_B.weight",
+ ".conv.down": ".conv.lora_A.weight",
+ ".conv.up": ".conv.lora_B.weight",
+ **{f".conv{i}.down": f".conv{i}.lora_A.weight" for i in range(1, 3)},
+ **{f".conv{i}.up": f".conv{i}.lora_B.weight" for i in range(1, 3)},
+ "conv_in.down": "conv_in.lora_A.weight",
+ "conv_in.up": "conv_in.lora_B.weight",
+ ".conv_shortcut.down": ".conv_shortcut.lora_A.weight",
+ ".conv_shortcut.up": ".conv_shortcut.lora_B.weight",
+ **{f".linear_{i}.down": f".linear_{i}.lora_A.weight" for i in range(1, 3)},
+ **{f".linear_{i}.up": f".linear_{i}.lora_B.weight" for i in range(1, 3)},
+ "time_emb_proj.down": "time_emb_proj.lora_A.weight",
+ "time_emb_proj.up": "time_emb_proj.lora_B.weight",
+}
+
+DIFFUSERS_TO_PEFT = {
+ ".q_proj.lora_linear_layer.up": ".q_proj.lora_B",
+ ".q_proj.lora_linear_layer.down": ".q_proj.lora_A",
+ ".k_proj.lora_linear_layer.up": ".k_proj.lora_B",
+ ".k_proj.lora_linear_layer.down": ".k_proj.lora_A",
+ ".v_proj.lora_linear_layer.up": ".v_proj.lora_B",
+ ".v_proj.lora_linear_layer.down": ".v_proj.lora_A",
+ ".out_proj.lora_linear_layer.up": ".out_proj.lora_B",
+ ".out_proj.lora_linear_layer.down": ".out_proj.lora_A",
+ ".lora_linear_layer.up": ".lora_B",
+ ".lora_linear_layer.down": ".lora_A",
+ "text_projection.lora.down.weight": "text_projection.lora_A.weight",
+ "text_projection.lora.up.weight": "text_projection.lora_B.weight",
+}
+
+DIFFUSERS_OLD_TO_PEFT = {
+ ".to_q_lora.up": ".q_proj.lora_B",
+ ".to_q_lora.down": ".q_proj.lora_A",
+ ".to_k_lora.up": ".k_proj.lora_B",
+ ".to_k_lora.down": ".k_proj.lora_A",
+ ".to_v_lora.up": ".v_proj.lora_B",
+ ".to_v_lora.down": ".v_proj.lora_A",
+ ".to_out_lora.up": ".out_proj.lora_B",
+ ".to_out_lora.down": ".out_proj.lora_A",
+ ".lora_linear_layer.up": ".lora_B",
+ ".lora_linear_layer.down": ".lora_A",
+}
+
+PEFT_TO_DIFFUSERS = {
+ ".q_proj.lora_B": ".q_proj.lora_linear_layer.up",
+ ".q_proj.lora_A": ".q_proj.lora_linear_layer.down",
+ ".k_proj.lora_B": ".k_proj.lora_linear_layer.up",
+ ".k_proj.lora_A": ".k_proj.lora_linear_layer.down",
+ ".v_proj.lora_B": ".v_proj.lora_linear_layer.up",
+ ".v_proj.lora_A": ".v_proj.lora_linear_layer.down",
+ ".out_proj.lora_B": ".out_proj.lora_linear_layer.up",
+ ".out_proj.lora_A": ".out_proj.lora_linear_layer.down",
+ "to_k.lora_A": "to_k.lora.down",
+ "to_k.lora_B": "to_k.lora.up",
+ "to_q.lora_A": "to_q.lora.down",
+ "to_q.lora_B": "to_q.lora.up",
+ "to_v.lora_A": "to_v.lora.down",
+ "to_v.lora_B": "to_v.lora.up",
+ "to_out.0.lora_A": "to_out.0.lora.down",
+ "to_out.0.lora_B": "to_out.0.lora.up",
+}
+
+DIFFUSERS_OLD_TO_DIFFUSERS = {
+ ".to_q_lora.up": ".q_proj.lora_linear_layer.up",
+ ".to_q_lora.down": ".q_proj.lora_linear_layer.down",
+ ".to_k_lora.up": ".k_proj.lora_linear_layer.up",
+ ".to_k_lora.down": ".k_proj.lora_linear_layer.down",
+ ".to_v_lora.up": ".v_proj.lora_linear_layer.up",
+ ".to_v_lora.down": ".v_proj.lora_linear_layer.down",
+ ".to_out_lora.up": ".out_proj.lora_linear_layer.up",
+ ".to_out_lora.down": ".out_proj.lora_linear_layer.down",
+ ".to_k.lora_magnitude_vector": ".k_proj.lora_magnitude_vector",
+ ".to_v.lora_magnitude_vector": ".v_proj.lora_magnitude_vector",
+ ".to_q.lora_magnitude_vector": ".q_proj.lora_magnitude_vector",
+ ".to_out.lora_magnitude_vector": ".out_proj.lora_magnitude_vector",
+}
+
+PEFT_TO_KOHYA_SS = {
+ "lora_A": "lora_down",
+ "lora_B": "lora_up",
+ # This is not a comprehensive dict as kohya format requires replacing `.` with `_` in keys,
+ # adding prefixes and adding alpha values
+ # Check `convert_state_dict_to_kohya` for more
+}
+
+PEFT_STATE_DICT_MAPPINGS = {
+ StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_PEFT,
+ StateDictType.DIFFUSERS: DIFFUSERS_TO_PEFT,
+}
+
+DIFFUSERS_STATE_DICT_MAPPINGS = {
+ StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_DIFFUSERS,
+ StateDictType.PEFT: PEFT_TO_DIFFUSERS,
+}
+
+KOHYA_STATE_DICT_MAPPINGS = {StateDictType.PEFT: PEFT_TO_KOHYA_SS}
+
+KEYS_TO_ALWAYS_REPLACE = {
+ ".processor.": ".",
+}
+
+
+def convert_state_dict(state_dict, mapping):
+ r"""
+ Simply iterates over the state dict and replaces the patterns in `mapping` with the corresponding values.
+
+ Args:
+ state_dict (`dict[str, torch.Tensor]`):
+ The state dict to convert.
+ mapping (`dict[str, str]`):
+ The mapping to use for conversion, the mapping should be a dictionary with the following structure:
+ - key: the pattern to replace
+ - value: the pattern to replace with
+
+ Returns:
+ converted_state_dict (`dict`)
+ The converted state dict.
+ """
+ converted_state_dict = {}
+ for k, v in state_dict.items():
+ # First, filter out the keys that we always want to replace
+ for pattern in KEYS_TO_ALWAYS_REPLACE.keys():
+ if pattern in k:
+ new_pattern = KEYS_TO_ALWAYS_REPLACE[pattern]
+ k = k.replace(pattern, new_pattern)
+
+ for pattern in mapping.keys():
+ if pattern in k:
+ new_pattern = mapping[pattern]
+ k = k.replace(pattern, new_pattern)
+ break
+ converted_state_dict[k] = v
+ return converted_state_dict
+
+
+def convert_state_dict_to_peft(state_dict, original_type=None, **kwargs):
+ r"""
+ Converts a state dict to the PEFT format The state dict can be from previous diffusers format (`OLD_DIFFUSERS`), or
+ new diffusers format (`DIFFUSERS`). The method only supports the conversion from diffusers old/new to PEFT for now.
+
+ Args:
+ state_dict (`dict[str, torch.Tensor]`):
+ The state dict to convert.
+ original_type (`StateDictType`, *optional*):
+ The original type of the state dict, if not provided, the method will try to infer it automatically.
+ """
+ if original_type is None:
+ # Old diffusers to PEFT
+ if any("to_out_lora" in k for k in state_dict.keys()):
+ original_type = StateDictType.DIFFUSERS_OLD
+ elif any("lora_linear_layer" in k for k in state_dict.keys()):
+ original_type = StateDictType.DIFFUSERS
+ else:
+ raise ValueError("Could not automatically infer state dict type")
+
+ if original_type not in PEFT_STATE_DICT_MAPPINGS.keys():
+ raise ValueError(f"Original type {original_type} is not supported")
+
+ mapping = PEFT_STATE_DICT_MAPPINGS[original_type]
+ return convert_state_dict(state_dict, mapping)
+
+
+def convert_state_dict_to_diffusers(state_dict, original_type=None, **kwargs):
+ r"""
+ Converts a state dict to new diffusers format. The state dict can be from previous diffusers format
+ (`OLD_DIFFUSERS`), or PEFT format (`PEFT`) or new diffusers format (`DIFFUSERS`). In the last case the method will
+ return the state dict as is.
+
+ The method only supports the conversion from diffusers old, PEFT to diffusers new for now.
+
+ Args:
+ state_dict (`dict[str, torch.Tensor]`):
+ The state dict to convert.
+ original_type (`StateDictType`, *optional*):
+ The original type of the state dict, if not provided, the method will try to infer it automatically.
+ kwargs (`dict`, *args*):
+ Additional arguments to pass to the method.
+
+ - **adapter_name**: For example, in case of PEFT, some keys will be prepended
+ with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in
+ `get_peft_model_state_dict` method:
+ https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92
+ but we add it here in case we don't want to rely on that method.
+ """
+ peft_adapter_name = kwargs.pop("adapter_name", None)
+ if peft_adapter_name is not None:
+ peft_adapter_name = "." + peft_adapter_name
+ else:
+ peft_adapter_name = ""
+
+ if original_type is None:
+ # Old diffusers to PEFT
+ if any("to_out_lora" in k for k in state_dict.keys()):
+ original_type = StateDictType.DIFFUSERS_OLD
+ elif any(f".lora_A{peft_adapter_name}.weight" in k for k in state_dict.keys()):
+ original_type = StateDictType.PEFT
+ elif any("lora_linear_layer" in k for k in state_dict.keys()):
+ # nothing to do
+ return state_dict
+ else:
+ raise ValueError("Could not automatically infer state dict type")
+
+ if original_type not in DIFFUSERS_STATE_DICT_MAPPINGS.keys():
+ raise ValueError(f"Original type {original_type} is not supported")
+
+ mapping = DIFFUSERS_STATE_DICT_MAPPINGS[original_type]
+ return convert_state_dict(state_dict, mapping)
+
+
+def convert_unet_state_dict_to_peft(state_dict):
+ r"""
+ Converts a state dict from UNet format to diffusers format - i.e. by removing some keys
+ """
+ mapping = UNET_TO_DIFFUSERS
+ return convert_state_dict(state_dict, mapping)
+
+
+def convert_sai_sd_control_lora_state_dict_to_peft(state_dict):
+ def _convert_controlnet_to_diffusers(state_dict):
+ is_sdxl = "input_blocks.11.0.in_layers.0.weight" not in state_dict
+ logger.info(f"Using ControlNet lora ({'SDXL' if is_sdxl else 'SD15'})")
+
+ # Retrieves the keys for the input blocks only
+ num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in state_dict if "input_blocks" in layer})
+ input_blocks = {
+ layer_id: [key for key in state_dict if f"input_blocks.{layer_id}" in key]
+ for layer_id in range(num_input_blocks)
+ }
+ layers_per_block = 2
+
+ # op blocks
+ op_blocks = [key for key in state_dict if "0.op" in key]
+
+ converted_state_dict = {}
+ # Conv in layers
+ for key in input_blocks[0]:
+ diffusers_key = key.replace("input_blocks.0.0", "conv_in")
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ # controlnet time embedding blocks
+ time_embedding_blocks = [key for key in state_dict if "time_embed" in key]
+ for key in time_embedding_blocks:
+ diffusers_key = key.replace("time_embed.0", "time_embedding.linear_1").replace(
+ "time_embed.2", "time_embedding.linear_2"
+ )
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ # controlnet label embedding blocks
+ label_embedding_blocks = [key for key in state_dict if "label_emb" in key]
+ for key in label_embedding_blocks:
+ diffusers_key = key.replace("label_emb.0.0", "add_embedding.linear_1").replace(
+ "label_emb.0.2", "add_embedding.linear_2"
+ )
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ # Down blocks
+ for i in range(1, num_input_blocks):
+ block_id = (i - 1) // (layers_per_block + 1)
+ layer_in_block_id = (i - 1) % (layers_per_block + 1)
+
+ resnets = [
+ key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+ ]
+ for key in resnets:
+ diffusers_key = (
+ key.replace("in_layers.0", "norm1")
+ .replace("in_layers.2", "conv1")
+ .replace("out_layers.0", "norm2")
+ .replace("out_layers.3", "conv2")
+ .replace("emb_layers.1", "time_emb_proj")
+ .replace("skip_connection", "conv_shortcut")
+ )
+ diffusers_key = diffusers_key.replace(
+ f"input_blocks.{i}.0", f"down_blocks.{block_id}.resnets.{layer_in_block_id}"
+ )
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ if f"input_blocks.{i}.0.op.bias" in state_dict:
+ for key in [key for key in op_blocks if f"input_blocks.{i}.0.op" in key]:
+ diffusers_key = key.replace(
+ f"input_blocks.{i}.0.op", f"down_blocks.{block_id}.downsamplers.0.conv"
+ )
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+ if attentions:
+ for key in attentions:
+ diffusers_key = key.replace(
+ f"input_blocks.{i}.1", f"down_blocks.{block_id}.attentions.{layer_in_block_id}"
+ )
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ # controlnet down blocks
+ for i in range(num_input_blocks):
+ converted_state_dict[f"controlnet_down_blocks.{i}.weight"] = state_dict.get(f"zero_convs.{i}.0.weight")
+ converted_state_dict[f"controlnet_down_blocks.{i}.bias"] = state_dict.get(f"zero_convs.{i}.0.bias")
+
+ # Retrieves the keys for the middle blocks only
+ num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in state_dict if "middle_block" in layer})
+ middle_blocks = {
+ layer_id: [key for key in state_dict if f"middle_block.{layer_id}" in key]
+ for layer_id in range(num_middle_blocks)
+ }
+
+ # Mid blocks
+ for key in middle_blocks.keys():
+ diffusers_key = max(key - 1, 0)
+ if key % 2 == 0:
+ for k in middle_blocks[key]:
+ diffusers_key_hf = (
+ k.replace("in_layers.0", "norm1")
+ .replace("in_layers.2", "conv1")
+ .replace("out_layers.0", "norm2")
+ .replace("out_layers.3", "conv2")
+ .replace("emb_layers.1", "time_emb_proj")
+ .replace("skip_connection", "conv_shortcut")
+ )
+ diffusers_key_hf = diffusers_key_hf.replace(
+ f"middle_block.{key}", f"mid_block.resnets.{diffusers_key}"
+ )
+ converted_state_dict[diffusers_key_hf] = state_dict.get(k)
+ else:
+ for k in middle_blocks[key]:
+ diffusers_key_hf = k.replace(f"middle_block.{key}", f"mid_block.attentions.{diffusers_key}")
+ converted_state_dict[diffusers_key_hf] = state_dict.get(k)
+
+ # mid block
+ converted_state_dict["controlnet_mid_block.weight"] = state_dict.get("middle_block_out.0.weight")
+ converted_state_dict["controlnet_mid_block.bias"] = state_dict.get("middle_block_out.0.bias")
+
+ # controlnet cond embedding blocks
+ cond_embedding_blocks = {
+ ".".join(layer.split(".")[:2])
+ for layer in state_dict
+ if "input_hint_block" in layer
+ and ("input_hint_block.0" not in layer)
+ and ("input_hint_block.14" not in layer)
+ }
+ num_cond_embedding_blocks = len(cond_embedding_blocks)
+
+ for idx in range(1, num_cond_embedding_blocks + 1):
+ diffusers_idx = idx - 1
+ cond_block_id = 2 * idx
+
+ converted_state_dict[f"controlnet_cond_embedding.blocks.{diffusers_idx}.weight"] = state_dict.get(
+ f"input_hint_block.{cond_block_id}.weight"
+ )
+ converted_state_dict[f"controlnet_cond_embedding.blocks.{diffusers_idx}.bias"] = state_dict.get(
+ f"input_hint_block.{cond_block_id}.bias"
+ )
+
+ for key in [key for key in state_dict if "input_hint_block.0" in key]:
+ diffusers_key = key.replace("input_hint_block.0", "controlnet_cond_embedding.conv_in")
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ for key in [key for key in state_dict if "input_hint_block.14" in key]:
+ diffusers_key = key.replace("input_hint_block.14", "controlnet_cond_embedding.conv_out")
+ converted_state_dict[diffusers_key] = state_dict.get(key)
+
+ return converted_state_dict
+
+ state_dict = _convert_controlnet_to_diffusers(state_dict)
+ mapping = CONTROL_LORA_TO_DIFFUSERS
+ return convert_state_dict(state_dict, mapping)
+
+
+def convert_all_state_dict_to_peft(state_dict):
+ r"""
+ Attempts to first `convert_state_dict_to_peft`, and if it doesn't detect `lora_linear_layer` for a valid
+ `DIFFUSERS` LoRA for example, attempts to exclusively convert the Unet `convert_unet_state_dict_to_peft`
+ """
+ try:
+ peft_dict = convert_state_dict_to_peft(state_dict)
+ except Exception as e:
+ if str(e) == "Could not automatically infer state dict type":
+ peft_dict = convert_unet_state_dict_to_peft(state_dict)
+ else:
+ raise
+
+ if not any("lora_A" in key or "lora_B" in key for key in peft_dict.keys()):
+ raise ValueError("Your LoRA was not converted to PEFT")
+
+ return peft_dict
+
+
+def convert_state_dict_to_kohya(state_dict, original_type=None, **kwargs):
+ r"""
+ Converts a `PEFT` state dict to `Kohya` format that can be used in AUTOMATIC1111, ComfyUI, SD.Next, InvokeAI, etc.
+ The method only supports the conversion from PEFT to Kohya for now.
+
+ Args:
+ state_dict (`dict[str, torch.Tensor]`):
+ The state dict to convert.
+ original_type (`StateDictType`, *optional*):
+ The original type of the state dict, if not provided, the method will try to infer it automatically.
+ kwargs (`dict`, *args*):
+ Additional arguments to pass to the method.
+
+ - **adapter_name**: For example, in case of PEFT, some keys will be prepended
+ with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in
+ `get_peft_model_state_dict` method:
+ https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92
+ but we add it here in case we don't want to rely on that method.
+ """
+ try:
+ import torch
+ except ImportError:
+ logger.error("Converting PEFT state dicts to Kohya requires torch to be installed.")
+ raise
+
+ peft_adapter_name = kwargs.pop("adapter_name", None)
+ if peft_adapter_name is not None:
+ peft_adapter_name = "." + peft_adapter_name
+ else:
+ peft_adapter_name = ""
+
+ if original_type is None:
+ if any(f".lora_A{peft_adapter_name}.weight" in k for k in state_dict.keys()):
+ original_type = StateDictType.PEFT
+
+ if original_type not in KOHYA_STATE_DICT_MAPPINGS.keys():
+ raise ValueError(f"Original type {original_type} is not supported")
+
+ # Use the convert_state_dict function with the appropriate mapping
+ kohya_ss_partial_state_dict = convert_state_dict(state_dict, KOHYA_STATE_DICT_MAPPINGS[StateDictType.PEFT])
+ kohya_ss_state_dict = {}
+
+ # Additional logic for replacing header, alpha parameters `.` with `_` in all keys
+ for kohya_key, weight in kohya_ss_partial_state_dict.items():
+ if "text_encoder_2." in kohya_key:
+ kohya_key = kohya_key.replace("text_encoder_2.", "lora_te2.")
+ elif "text_encoder." in kohya_key:
+ kohya_key = kohya_key.replace("text_encoder.", "lora_te1.")
+ elif "unet" in kohya_key:
+ kohya_key = kohya_key.replace("unet", "lora_unet")
+ elif "lora_magnitude_vector" in kohya_key:
+ kohya_key = kohya_key.replace("lora_magnitude_vector", "dora_scale")
+
+ kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
+ kohya_key = kohya_key.replace(peft_adapter_name, "") # Kohya doesn't take names
+ kohya_ss_state_dict[kohya_key] = weight
+ if "lora_down" in kohya_key:
+ alpha_key = f"{kohya_key.split('.')[0]}.alpha"
+ kohya_ss_state_dict[alpha_key] = torch.tensor(len(weight))
+
+ return kohya_ss_state_dict
+
+
+def state_dict_all_zero(state_dict, filter_str=None):
+ if filter_str is not None:
+ if isinstance(filter_str, str):
+ filter_str = [filter_str]
+ state_dict = {k: v for k, v in state_dict.items() if any(f in k for f in filter_str)}
+
+ return all(torch.all(param == 0).item() for param in state_dict.values())
+
+
+def _load_sft_state_dict_metadata(model_file: str):
+ import safetensors.torch
+
+ from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY
+
+ with safetensors.torch.safe_open(model_file, framework="pt", device="cpu") as f:
+ metadata = f.metadata() or {}
+
+ metadata.pop("format", None)
+ if metadata:
+ raw = metadata.get(LORA_ADAPTER_METADATA_KEY)
+ return json.loads(raw) if raw else None
+ else:
+ return None
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/testing_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/testing_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4701527b166ba5bd9db71d6e617c5519720d1465
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/testing_utils.py
@@ -0,0 +1,1620 @@
+import functools
+import glob
+import importlib
+import importlib.metadata
+import inspect
+import io
+import logging
+import multiprocessing
+import os
+import random
+import re
+import struct
+import sys
+import tempfile
+import time
+import unittest
+import urllib.parse
+from collections import UserDict
+from contextlib import contextmanager
+from io import BytesIO, StringIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Set
+
+import numpy as np
+import PIL.Image
+import PIL.ImageOps
+import requests
+from numpy.linalg import norm
+from packaging import version
+
+from .constants import DIFFUSERS_REQUEST_TIMEOUT
+from .deprecation_utils import deprecate
+from .import_utils import (
+ BACKENDS_MAPPING,
+ is_accelerate_available,
+ is_bitsandbytes_available,
+ is_compel_available,
+ is_gguf_available,
+ is_kernels_available,
+ is_note_seq_available,
+ is_nvidia_modelopt_available,
+ is_onnx_available,
+ is_opencv_available,
+ is_optimum_quanto_available,
+ is_peft_available,
+ is_sdnq_available,
+ is_timm_available,
+ is_torch_available,
+ is_torch_neuronx_available,
+ is_torch_version,
+ is_torchao_available,
+ is_torchsde_available,
+ is_transformers_available,
+)
+from .logging import get_logger
+
+
+if is_torch_available():
+ import torch
+
+ IS_ROCM_SYSTEM = torch.version.hip is not None
+ IS_CUDA_SYSTEM = torch.version.cuda is not None
+ IS_XPU_SYSTEM = getattr(torch.version, "xpu", None) is not None
+else:
+ IS_ROCM_SYSTEM = False
+ IS_CUDA_SYSTEM = False
+ IS_XPU_SYSTEM = False
+
+global_rng = random.Random()
+
+logger = get_logger(__name__)
+deprecate(
+ "diffusers.utils.testing_utils",
+ "1.0.0",
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version. "
+ "Determinism and device backend utilities have been moved to `diffusers.utils.torch_utils`. ",
+)
+_required_peft_version = is_peft_available() and version.parse(
+ version.parse(importlib.metadata.version("peft")).base_version
+) > version.parse("0.5")
+_required_transformers_version = is_transformers_available() and version.parse(
+ version.parse(importlib.metadata.version("transformers")).base_version
+) > version.parse("4.33")
+
+USE_PEFT_BACKEND = _required_peft_version and _required_transformers_version
+BIG_GPU_MEMORY = int(os.getenv("BIG_GPU_MEMORY", 40))
+
+if is_torch_available():
+ import torch
+
+ # Set a backend environment variable for any extra module import required for a custom accelerator
+ if "DIFFUSERS_TEST_BACKEND" in os.environ:
+ backend = os.environ["DIFFUSERS_TEST_BACKEND"]
+ try:
+ _ = importlib.import_module(backend)
+ except ModuleNotFoundError as e:
+ raise ModuleNotFoundError(
+ f"Failed to import `DIFFUSERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module \
+ to enable a specified backend.):\n{e}"
+ ) from e
+
+ if "DIFFUSERS_TEST_DEVICE" in os.environ:
+ torch_device = os.environ["DIFFUSERS_TEST_DEVICE"]
+ try:
+ # try creating device to see if provided device is valid
+ _ = torch.device(torch_device)
+ except RuntimeError as e:
+ raise RuntimeError(
+ f"Unknown testing device specified by environment variable `DIFFUSERS_TEST_DEVICE`: {torch_device}"
+ ) from e
+ logger.info(f"torch_device overrode to {torch_device}")
+ else:
+ if torch.cuda.is_available():
+ torch_device = "cuda"
+ elif torch.xpu.is_available():
+ torch_device = "xpu"
+ elif is_torch_neuronx_available() and hasattr(torch, "neuron") and torch.neuron.is_available():
+ torch_device = torch.neuron.current_device()
+ else:
+ torch_device = "cpu"
+ is_torch_higher_equal_than_1_12 = version.parse(
+ version.parse(torch.__version__).base_version
+ ) >= version.parse("1.12")
+
+ if is_torch_higher_equal_than_1_12:
+ # Some builds of torch 1.12 don't have the mps backend registered. See #892 for more details
+ mps_backend_registered = hasattr(torch.backends, "mps")
+ torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device
+
+ from .torch_utils import get_torch_cuda_device_capability
+
+
+def torch_all_close(a, b, *args, **kwargs):
+ if not is_torch_available():
+ raise ValueError("PyTorch needs to be installed to use this function.")
+ if not torch.allclose(a, b, *args, **kwargs):
+ assert False, f"Max diff is absolute {(a - b).abs().max()}. Diff tensor is {(a - b).abs()}."
+ return True
+
+
+def numpy_cosine_similarity_distance(a, b):
+ similarity = np.dot(a, b) / (norm(a) * norm(b))
+ distance = 1.0 - similarity.mean()
+
+ return distance
+
+
+def check_if_dicts_are_equal(dict1, dict2):
+ dict1, dict2 = dict1.copy(), dict2.copy()
+
+ for key, value in dict1.items():
+ if isinstance(value, set):
+ dict1[key] = sorted(value)
+ for key, value in dict2.items():
+ if isinstance(value, set):
+ dict2[key] = sorted(value)
+
+ for key in dict1:
+ if key not in dict2:
+ return False
+ if dict1[key] != dict2[key]:
+ return False
+
+ for key in dict2:
+ if key not in dict1:
+ return False
+
+ return True
+
+
+def print_tensor_test(
+ tensor,
+ limit_to_slices=None,
+ max_torch_print=None,
+ filename="test_corrections.txt",
+ expected_tensor_name="expected_slice",
+):
+ if max_torch_print:
+ torch.set_printoptions(threshold=10_000)
+
+ test_name = os.environ.get("PYTEST_CURRENT_TEST")
+ if not torch.is_tensor(tensor):
+ tensor = torch.from_numpy(tensor)
+ if limit_to_slices:
+ tensor = tensor[0, -3:, -3:, -1]
+
+ tensor_str = str(tensor.detach().cpu().flatten().to(torch.float32)).replace("\n", "")
+ # format is usually:
+ # expected_slice = np.array([-0.5713, -0.3018, -0.9814, 0.04663, -0.879, 0.76, -1.734, 0.1044, 1.161])
+ output_str = tensor_str.replace("tensor", f"{expected_tensor_name} = np.array")
+ test_file, test_class, test_fn = test_name.split("::")
+ test_fn = test_fn.split()[0]
+ with open(filename, "a") as f:
+ print("::".join([test_file, test_class, test_fn, output_str]), file=f)
+
+
+def get_tests_dir(append_path=None):
+ """
+ Args:
+ append_path: optional path to append to the tests dir path
+ Return:
+ The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is
+ joined after the `tests` dir the former is provided.
+ """
+ # this function caller's __file__
+ caller__file__ = inspect.stack()[1][1]
+ tests_dir = os.path.abspath(os.path.dirname(caller__file__))
+
+ while not tests_dir.endswith("tests"):
+ tests_dir = os.path.dirname(tests_dir)
+
+ if append_path:
+ return Path(tests_dir, append_path).as_posix()
+ else:
+ return tests_dir
+
+
+# Taken from the following PR:
+# https://github.com/huggingface/accelerate/pull/1964
+def str_to_bool(value) -> int:
+ """
+ Converts a string representation of truth to `True` (1) or `False` (0). True values are `y`, `yes`, `t`, `true`,
+ `on`, and `1`; False value are `n`, `no`, `f`, `false`, `off`, and `0`;
+ """
+ value = value.lower()
+ if value in ("y", "yes", "t", "true", "on", "1"):
+ return 1
+ elif value in ("n", "no", "f", "false", "off", "0"):
+ return 0
+ else:
+ raise ValueError(f"invalid truth value {value}")
+
+
+def parse_flag_from_env(key, default=False):
+ try:
+ value = os.environ[key]
+ except KeyError:
+ # KEY isn't set, default to `default`.
+ _value = default
+ else:
+ # KEY is set, convert it to True or False.
+ try:
+ _value = str_to_bool(value)
+ except ValueError:
+ # More values are supported, but let's keep the message simple.
+ raise ValueError(f"If set, {key} must be yes or no.")
+ return _value
+
+
+_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
+_run_nightly_tests = parse_flag_from_env("RUN_NIGHTLY", default=False)
+_run_compile_tests = parse_flag_from_env("RUN_COMPILE", default=False)
+
+
+def floats_tensor(shape, scale=1.0, rng=None, name=None):
+ """Creates a random float32 tensor"""
+ if rng is None:
+ rng = global_rng
+
+ total_dims = 1
+ for dim in shape:
+ total_dims *= dim
+
+ values = []
+ for _ in range(total_dims):
+ values.append(rng.random() * scale)
+
+ return torch.tensor(data=values, dtype=torch.float).view(shape).contiguous()
+
+
+def slow(test_case):
+ """
+ Decorator marking a test as slow.
+
+ Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.
+
+ """
+ return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)
+
+
+def nightly(test_case):
+ """
+ Decorator marking a test that runs nightly in the diffusers CI.
+
+ Slow tests are skipped by default. Set the RUN_NIGHTLY environment variable to a truthy value to run them.
+
+ """
+ return unittest.skipUnless(_run_nightly_tests, "test is nightly")(test_case)
+
+
+def is_torch_compile(test_case):
+ """
+ Decorator marking a test that runs compile tests in the diffusers CI.
+
+ Compile tests are skipped by default. Set the RUN_COMPILE environment variable to a truthy value to run them.
+
+ """
+ return unittest.skipUnless(_run_compile_tests, "test is torch compile")(test_case)
+
+
+def require_torch(test_case):
+ """
+ Decorator marking a test that requires PyTorch. These tests are skipped when PyTorch isn't installed.
+ """
+ return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case)
+
+
+def require_torch_2(test_case):
+ """
+ Decorator marking a test that requires PyTorch 2. These tests are skipped when it isn't installed.
+ """
+ return unittest.skipUnless(is_torch_available() and is_torch_version(">=", "2.0.0"), "test requires PyTorch 2")(
+ test_case
+ )
+
+
+def require_torch_version_greater_equal(torch_version):
+ """Decorator marking a test that requires torch with a specific version or greater."""
+
+ def decorator(test_case):
+ correct_torch_version = is_torch_available() and is_torch_version(">=", torch_version)
+ return unittest.skipUnless(
+ correct_torch_version, f"test requires torch with the version greater than or equal to {torch_version}"
+ )(test_case)
+
+ return decorator
+
+
+def require_torch_version_greater(torch_version):
+ """Decorator marking a test that requires torch with a specific version greater."""
+
+ def decorator(test_case):
+ correct_torch_version = is_torch_available() and is_torch_version(">", torch_version)
+ return unittest.skipUnless(
+ correct_torch_version, f"test requires torch with the version greater than {torch_version}"
+ )(test_case)
+
+ return decorator
+
+
+def require_torch_gpu(test_case):
+ """Decorator marking a test that requires CUDA and PyTorch."""
+ return unittest.skipUnless(is_torch_available() and torch_device == "cuda", "test requires PyTorch+CUDA")(
+ test_case
+ )
+
+
+def require_torch_cuda_compatibility(expected_compute_capability):
+ def decorator(test_case):
+ if torch.cuda.is_available():
+ current_compute_capability = get_torch_cuda_device_capability()
+ return unittest.skipUnless(
+ float(current_compute_capability) == float(expected_compute_capability),
+ "Test not supported for this compute capability.",
+ )
+
+ return decorator
+
+
+# These decorators are for accelerator-specific behaviours that are not GPU-specific
+def require_torch_accelerator(test_case):
+ """Decorator marking a test that requires an accelerator backend and PyTorch."""
+ return unittest.skipUnless(is_torch_available() and torch_device != "cpu", "test requires accelerator+PyTorch")(
+ test_case
+ )
+
+
+def require_torch_multi_gpu(test_case):
+ """
+ Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without
+ multiple GPUs. To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests
+ -k "multi_gpu"
+ """
+ if not is_torch_available():
+ return unittest.skip("test requires PyTorch")(test_case)
+
+ import torch
+
+ return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
+
+
+def require_torch_multi_accelerator(test_case):
+ """
+ Decorator marking a test that requires a multi-accelerator setup (in PyTorch). These tests are skipped on a machine
+ without multiple hardware accelerators.
+ """
+ if not is_torch_available():
+ return unittest.skip("test requires PyTorch")(test_case)
+
+ import torch
+
+ return unittest.skipUnless(
+ torch.cuda.device_count() > 1 or torch.xpu.device_count() > 1, "test requires multiple hardware accelerators"
+ )(test_case)
+
+
+def require_torch_accelerator_with_fp16(test_case):
+ """Decorator marking a test that requires an accelerator with support for the FP16 data type."""
+ return unittest.skipUnless(_is_torch_fp16_available(torch_device), "test requires accelerator with fp16 support")(
+ test_case
+ )
+
+
+def require_torch_accelerator_with_fp64(test_case):
+ """Decorator marking a test that requires an accelerator with support for the FP64 data type."""
+ return unittest.skipUnless(_is_torch_fp64_available(torch_device), "test requires accelerator with fp64 support")(
+ test_case
+ )
+
+
+def require_big_gpu_with_torch_cuda(test_case):
+ """
+ Decorator marking a test that requires a bigger GPU (24GB) for execution. Some example pipelines: Flux, SD3, Cog,
+ etc.
+ """
+ if not is_torch_available():
+ return unittest.skip("test requires PyTorch")(test_case)
+
+ import torch
+
+ if not torch.cuda.is_available():
+ return unittest.skip("test requires PyTorch CUDA")(test_case)
+
+ device_properties = torch.cuda.get_device_properties(0)
+ total_memory = device_properties.total_memory / (1024**3)
+ return unittest.skipUnless(
+ total_memory >= BIG_GPU_MEMORY, f"test requires a GPU with at least {BIG_GPU_MEMORY} GB memory"
+ )(test_case)
+
+
+def require_big_accelerator(test_case):
+ """
+ Decorator marking a test that requires a bigger hardware accelerator (24GB) for execution. Some example pipelines:
+ Flux, SD3, Cog, etc.
+ """
+ import pytest
+
+ test_case = pytest.mark.big_accelerator(test_case)
+
+ if not is_torch_available():
+ return unittest.skip("test requires PyTorch")(test_case)
+
+ import torch
+
+ if not (torch.cuda.is_available() or torch.xpu.is_available()):
+ return unittest.skip("test requires PyTorch CUDA")(test_case)
+
+ if torch.xpu.is_available():
+ device_properties = torch.xpu.get_device_properties(0)
+ else:
+ device_properties = torch.cuda.get_device_properties(0)
+
+ total_memory = device_properties.total_memory / (1024**3)
+ return unittest.skipUnless(
+ total_memory >= BIG_GPU_MEMORY,
+ f"test requires a hardware accelerator with at least {BIG_GPU_MEMORY} GB memory",
+ )(test_case)
+
+
+def require_torch_accelerator_with_training(test_case):
+ """Decorator marking a test that requires an accelerator with support for training."""
+ return unittest.skipUnless(
+ is_torch_available() and backend_supports_training(torch_device),
+ "test requires accelerator with training support",
+ )(test_case)
+
+
+def skip_mps(test_case):
+ """Decorator marking a test to skip if torch_device is 'mps'"""
+ return unittest.skipUnless(torch_device != "mps", "test requires non 'mps' device")(test_case)
+
+
+def require_compel(test_case):
+ """
+ Decorator marking a test that requires compel: https://github.com/damian0815/compel. These tests are skipped when
+ the library is not installed.
+ """
+ return unittest.skipUnless(is_compel_available(), "test requires compel")(test_case)
+
+
+def require_onnxruntime(test_case):
+ """
+ Decorator marking a test that requires onnxruntime. These tests are skipped when onnxruntime isn't installed.
+ """
+ return unittest.skipUnless(is_onnx_available(), "test requires onnxruntime")(test_case)
+
+
+def require_note_seq(test_case):
+ """
+ Decorator marking a test that requires note_seq. These tests are skipped when note_seq isn't installed.
+ """
+ return unittest.skipUnless(is_note_seq_available(), "test requires note_seq")(test_case)
+
+
+def require_accelerator(test_case):
+ """
+ Decorator marking a test that requires a hardware accelerator backend. These tests are skipped when there are no
+ hardware accelerator available.
+ """
+ return unittest.skipUnless(torch_device != "cpu", "test requires a hardware accelerator")(test_case)
+
+
+def require_torchsde(test_case):
+ """
+ Decorator marking a test that requires torchsde. These tests are skipped when torchsde isn't installed.
+ """
+ return unittest.skipUnless(is_torchsde_available(), "test requires torchsde")(test_case)
+
+
+def require_peft_backend(test_case):
+ """
+ Decorator marking a test that requires PEFT backend, this would require some specific versions of PEFT and
+ transformers.
+ """
+ return unittest.skipUnless(USE_PEFT_BACKEND, "test requires PEFT backend")(test_case)
+
+
+def require_timm(test_case):
+ """
+ Decorator marking a test that requires timm. These tests are skipped when timm isn't installed.
+ """
+ return unittest.skipUnless(is_timm_available(), "test requires timm")(test_case)
+
+
+def require_bitsandbytes(test_case):
+ """
+ Decorator marking a test that requires bitsandbytes. These tests are skipped when bitsandbytes isn't installed.
+ """
+ return unittest.skipUnless(is_bitsandbytes_available(), "test requires bitsandbytes")(test_case)
+
+
+def require_quanto(test_case):
+ """
+ Decorator marking a test that requires quanto. These tests are skipped when quanto isn't installed.
+ """
+ return unittest.skipUnless(is_optimum_quanto_available(), "test requires quanto")(test_case)
+
+
+def require_sdnq(test_case):
+ """
+ Decorator marking a test that requires sdnq. These tests are skipped when sdnq isn't installed.
+ """
+ return unittest.skipUnless(is_sdnq_available(), "test requires sdnq")(test_case)
+
+
+def require_accelerate(test_case):
+ """
+ Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed.
+ """
+ return unittest.skipUnless(is_accelerate_available(), "test requires accelerate")(test_case)
+
+
+def require_peft_version_greater(peft_version):
+ """
+ Decorator marking a test that requires PEFT backend with a specific version, this would require some specific
+ versions of PEFT and transformers.
+ """
+
+ def decorator(test_case):
+ correct_peft_version = is_peft_available() and version.parse(
+ version.parse(importlib.metadata.version("peft")).base_version
+ ) > version.parse(peft_version)
+ return unittest.skipUnless(
+ correct_peft_version, f"test requires PEFT backend with the version greater than {peft_version}"
+ )(test_case)
+
+ return decorator
+
+
+def require_transformers_version_greater(transformers_version):
+ """
+ Decorator marking a test that requires transformers with a specific version, this would require some specific
+ versions of PEFT and transformers.
+ """
+
+ def decorator(test_case):
+ correct_transformers_version = is_transformers_available() and version.parse(
+ version.parse(importlib.metadata.version("transformers")).base_version
+ ) > version.parse(transformers_version)
+ return unittest.skipUnless(
+ correct_transformers_version,
+ f"test requires transformers with the version greater than {transformers_version}",
+ )(test_case)
+
+ return decorator
+
+
+def require_accelerate_version_greater(accelerate_version):
+ def decorator(test_case):
+ correct_accelerate_version = is_accelerate_available() and version.parse(
+ version.parse(importlib.metadata.version("accelerate")).base_version
+ ) > version.parse(accelerate_version)
+ return unittest.skipUnless(
+ correct_accelerate_version, f"Test requires accelerate with the version greater than {accelerate_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_bitsandbytes_version_greater(bnb_version):
+ def decorator(test_case):
+ correct_bnb_version = is_bitsandbytes_available() and version.parse(
+ version.parse(importlib.metadata.version("bitsandbytes")).base_version
+ ) > version.parse(bnb_version)
+ return unittest.skipUnless(
+ correct_bnb_version, f"Test requires bitsandbytes with the version greater than {bnb_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_hf_hub_version_greater(hf_hub_version):
+ def decorator(test_case):
+ correct_hf_hub_version = version.parse(
+ version.parse(importlib.metadata.version("huggingface_hub")).base_version
+ ) > version.parse(hf_hub_version)
+ return unittest.skipUnless(
+ correct_hf_hub_version, f"Test requires huggingface_hub with the version greater than {hf_hub_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_gguf_version_greater_or_equal(gguf_version):
+ def decorator(test_case):
+ correct_gguf_version = is_gguf_available() and version.parse(
+ version.parse(importlib.metadata.version("gguf")).base_version
+ ) >= version.parse(gguf_version)
+ return unittest.skipUnless(
+ correct_gguf_version, f"Test requires gguf with the version greater than {gguf_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_torchao_version_greater_or_equal(torchao_version):
+ def decorator(test_case):
+ correct_torchao_version = is_torchao_available() and version.parse(
+ version.parse(importlib.metadata.version("torchao")).base_version
+ ) >= version.parse(torchao_version)
+ return unittest.skipUnless(
+ correct_torchao_version, f"Test requires torchao with version greater than {torchao_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_modelopt_version_greater_or_equal(modelopt_version):
+ def decorator(test_case):
+ correct_nvidia_modelopt_version = is_nvidia_modelopt_available() and version.parse(
+ version.parse(importlib.metadata.version("nvidia-modelopt")).base_version
+ ) >= version.parse(modelopt_version)
+ return unittest.skipUnless(
+ correct_nvidia_modelopt_version, f"Test requires modelopt with version greater than {modelopt_version}."
+ )(test_case)
+
+ return decorator
+
+
+def require_kernels_version_greater_or_equal(kernels_version):
+ def decorator(test_case):
+ correct_kernels_version = is_kernels_available() and version.parse(
+ version.parse(importlib.metadata.version("kernels")).base_version
+ ) >= version.parse(kernels_version)
+ return unittest.skipUnless(
+ correct_kernels_version, f"Test requires kernels with version greater than {kernels_version}."
+ )(test_case)
+
+ return decorator
+
+
+def deprecate_after_peft_backend(test_case):
+ """
+ Decorator marking a test that will be skipped after PEFT backend
+ """
+ return unittest.skipUnless(not USE_PEFT_BACKEND, "test skipped in favor of PEFT backend")(test_case)
+
+
+def get_python_version():
+ sys_info = sys.version_info
+ major, minor = sys_info.major, sys_info.minor
+ return major, minor
+
+
+def load_numpy(arry: str | np.ndarray, local_path: str | None = None) -> np.ndarray:
+ if isinstance(arry, str):
+ if local_path is not None:
+ # local_path can be passed to correct images of tests
+ return Path(local_path, arry.split("/")[-5], arry.split("/")[-2], arry.split("/")[-1]).as_posix()
+ elif arry.startswith("http://") or arry.startswith("https://"):
+ response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT)
+ response.raise_for_status()
+ arry = np.load(BytesIO(response.content))
+ elif os.path.isfile(arry):
+ arry = np.load(arry)
+ else:
+ raise ValueError(
+ f"Incorrect path or url, URLs must start with `http://` or `https://`, and {arry} is not a valid path"
+ )
+ elif isinstance(arry, np.ndarray):
+ pass
+ else:
+ raise ValueError(
+ "Incorrect format used for numpy ndarray. Should be an url linking to an image, a local path, or a"
+ " ndarray."
+ )
+
+ return arry
+
+
+def load_pt(url: str, map_location: str | None = None, weights_only: bool | None = True):
+ response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
+ response.raise_for_status()
+ arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only)
+ return arry
+
+
+def load_image(image: str | PIL.Image.Image) -> PIL.Image.Image:
+ """
+ Loads `image` to a PIL Image.
+
+ Args:
+ image (`str` or `PIL.Image.Image`):
+ The image to convert to the PIL Image format.
+ Returns:
+ `PIL.Image.Image`:
+ A PIL Image.
+ """
+ if isinstance(image, str):
+ if image.startswith("http://") or image.startswith("https://"):
+ image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)
+ elif os.path.isfile(image):
+ image = PIL.Image.open(image)
+ else:
+ raise ValueError(
+ f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path"
+ )
+ elif isinstance(image, PIL.Image.Image):
+ image = image
+ else:
+ raise ValueError(
+ "Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image."
+ )
+ image = PIL.ImageOps.exif_transpose(image)
+ image = image.convert("RGB")
+ return image
+
+
+def preprocess_image(image: PIL.Image, batch_size: int):
+ w, h = image.size
+ w, h = (x - x % 8 for x in (w, h)) # resize to integer multiple of 8
+ image = image.resize((w, h), resample=PIL.Image.LANCZOS)
+ image = np.array(image).astype(np.float32) / 255.0
+ image = np.vstack([image[None].transpose(0, 3, 1, 2)] * batch_size)
+ image = torch.from_numpy(image)
+ return 2.0 * image - 1.0
+
+
+def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> str:
+ if output_gif_path is None:
+ output_gif_path = tempfile.NamedTemporaryFile(suffix=".gif").name
+
+ image[0].save(
+ output_gif_path,
+ save_all=True,
+ append_images=image[1:],
+ optimize=False,
+ duration=100,
+ loop=0,
+ )
+ return output_gif_path
+
+
+@contextmanager
+def buffered_writer(raw_f):
+ f = io.BufferedWriter(raw_f)
+ yield f
+ f.flush()
+
+
+def export_to_ply(mesh, output_ply_path: str = None):
+ """
+ Write a PLY file for a mesh.
+ """
+ if output_ply_path is None:
+ output_ply_path = tempfile.NamedTemporaryFile(suffix=".ply").name
+
+ coords = mesh.verts.detach().cpu().numpy()
+ faces = mesh.faces.cpu().numpy()
+ rgb = np.stack([mesh.vertex_channels[x].detach().cpu().numpy() for x in "RGB"], axis=1)
+
+ with buffered_writer(open(output_ply_path, "wb")) as f:
+ f.write(b"ply\n")
+ f.write(b"format binary_little_endian 1.0\n")
+ f.write(bytes(f"element vertex {len(coords)}\n", "ascii"))
+ f.write(b"property float x\n")
+ f.write(b"property float y\n")
+ f.write(b"property float z\n")
+ if rgb is not None:
+ f.write(b"property uchar red\n")
+ f.write(b"property uchar green\n")
+ f.write(b"property uchar blue\n")
+ if faces is not None:
+ f.write(bytes(f"element face {len(faces)}\n", "ascii"))
+ f.write(b"property list uchar int vertex_index\n")
+ f.write(b"end_header\n")
+
+ if rgb is not None:
+ rgb = (rgb * 255.499).round().astype(int)
+ vertices = [
+ (*coord, *rgb)
+ for coord, rgb in zip(
+ coords.tolist(),
+ rgb.tolist(),
+ )
+ ]
+ format = struct.Struct("<3f3B")
+ for item in vertices:
+ f.write(format.pack(*item))
+ else:
+ format = struct.Struct("<3f")
+ for vertex in coords.tolist():
+ f.write(format.pack(*vertex))
+
+ if faces is not None:
+ for tri in faces.tolist():
+ f.write(format.pack(len(tri), *tri))
+ format = struct.Struct(" str:
+ if is_opencv_available():
+ import cv2
+ else:
+ raise ImportError(BACKENDS_MAPPING["opencv"][1].format("export_to_video"))
+ if output_video_path is None:
+ output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+
+ fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+ h, w, c = video_frames[0].shape
+ video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=8, frameSize=(w, h))
+ for i in range(len(video_frames)):
+ img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)
+ video_writer.write(img)
+ return output_video_path
+
+
+def load_hf_numpy(path) -> np.ndarray:
+ base_url = "https://huggingface.co/datasets/fusing/diffusers-testing/resolve/main"
+
+ if not path.startswith("http://") and not path.startswith("https://"):
+ path = os.path.join(base_url, urllib.parse.quote(path))
+
+ return load_numpy(path)
+
+
+# --- pytest conf functions --- #
+
+# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once
+pytest_opt_registered = {}
+
+
+def pytest_addoption_shared(parser):
+ """
+ This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.
+
+ It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`
+ option.
+
+ """
+ option = "--make-reports"
+ if option not in pytest_opt_registered:
+ parser.addoption(
+ option,
+ action="store",
+ default=False,
+ help="generate report files. The value of this option is used as a prefix to report names",
+ )
+ pytest_opt_registered[option] = 1
+
+
+def pytest_terminal_summary_main(tr, id):
+ """
+ Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current
+ directory. The report files are prefixed with the test suite name.
+
+ This function emulates --duration and -rA pytest arguments.
+
+ This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined
+ there.
+
+ Args:
+ - tr: `terminalreporter` passed from `conftest.py`
+ - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is
+ needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other.
+
+ NB: this functions taps into a private _pytest API and while unlikely, it could break should
+ pytest do internal changes - also it calls default internal methods of terminalreporter which
+ can be hijacked by various `pytest-` plugins and interfere.
+
+ """
+ from _pytest.config import create_terminal_writer
+
+ if not len(id):
+ id = "tests"
+
+ config = tr.config
+ orig_writer = config.get_terminal_writer()
+ orig_tbstyle = config.option.tbstyle
+ orig_reportchars = tr.reportchars
+
+ dir = "reports"
+ Path(dir).mkdir(parents=True, exist_ok=True)
+ report_files = {
+ k: f"{dir}/{id}_{k}.txt"
+ for k in [
+ "durations",
+ "errors",
+ "failures_long",
+ "failures_short",
+ "failures_line",
+ "passes",
+ "stats",
+ "summary_short",
+ "warnings",
+ ]
+ }
+
+ # custom durations report
+ # note: there is no need to call pytest --durations=XX to get this separate report
+ # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66
+ dlist = []
+ for replist in tr.stats.values():
+ for rep in replist:
+ if hasattr(rep, "duration"):
+ dlist.append(rep)
+ if dlist:
+ dlist.sort(key=lambda x: x.duration, reverse=True)
+ with open(report_files["durations"], "w") as f:
+ durations_min = 0.05 # sec
+ f.write("slowest durations\n")
+ for i, rep in enumerate(dlist):
+ if rep.duration < durations_min:
+ f.write(f"{len(dlist) - i} durations < {durations_min} secs were omitted")
+ break
+ f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n")
+
+ def summary_failures_short(tr):
+ # expecting that the reports were --tb=long (default) so we chop them off here to the last frame
+ reports = tr.getreports("failed")
+ if not reports:
+ return
+ tr.write_sep("=", "FAILURES SHORT STACK")
+ for rep in reports:
+ msg = tr._getfailureheadline(rep)
+ tr.write_sep("_", msg, red=True, bold=True)
+ # chop off the optional leading extra frames, leaving only the last one
+ longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S)
+ tr._tw.line(longrepr)
+ # note: not printing out any rep.sections to keep the report short
+
+ # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each
+ # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814
+ # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g.
+ # pytest-instafail does that)
+
+ # report failures with line/short/long styles
+ config.option.tbstyle = "auto" # full tb
+ with open(report_files["failures_long"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_failures()
+
+ # config.option.tbstyle = "short" # short tb
+ with open(report_files["failures_short"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ summary_failures_short(tr)
+
+ config.option.tbstyle = "line" # one line per error
+ with open(report_files["failures_line"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_failures()
+
+ with open(report_files["errors"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_errors()
+
+ with open(report_files["warnings"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_warnings() # normal warnings
+ tr.summary_warnings() # final warnings
+
+ tr.reportchars = "wPpsxXEf" # emulate -rA (used in summary_passes() and short_test_summary())
+ with open(report_files["passes"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_passes()
+
+ with open(report_files["summary_short"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.short_test_summary()
+
+ with open(report_files["stats"], "w") as f:
+ tr._tw = create_terminal_writer(config, f)
+ tr.summary_stats()
+
+ # restore:
+ tr._tw = orig_writer
+ tr.reportchars = orig_reportchars
+ config.option.tbstyle = orig_tbstyle
+
+
+# Adapted from https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/testing_utils.py#L1905
+def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None):
+ """
+ To decorate flaky tests (methods or entire classes). They will be retried on failures.
+
+ Args:
+ max_attempts (`int`, *optional*, defaults to 5):
+ The maximum number of attempts to retry the flaky test.
+ wait_before_retry (`float`, *optional*):
+ If provided, will wait that number of seconds before retrying the test.
+ description (`str`, *optional*):
+ A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
+ etc.)
+ """
+
+ def decorator(obj):
+ # If decorating a class, wrap each test method on it
+ if inspect.isclass(obj):
+ for attr_name, attr_value in list(obj.__dict__.items()):
+ if callable(attr_value) and attr_name.startswith("test"):
+ # recursively decorate the method
+ setattr(obj, attr_name, decorator(attr_value))
+ return obj
+
+ # Otherwise we're decorating a single test function / method
+ @functools.wraps(obj)
+ def wrapper(*args, **kwargs):
+ retry_count = 1
+ while retry_count < max_attempts:
+ try:
+ return obj(*args, **kwargs)
+ except Exception as err:
+ msg = (
+ f"[FLAKY] {description or obj.__name__!r} "
+ f"failed on attempt {retry_count}/{max_attempts}: {err}"
+ )
+ print(msg, file=sys.stderr)
+ if wait_before_retry is not None:
+ time.sleep(wait_before_retry)
+ retry_count += 1
+
+ return obj(*args, **kwargs)
+
+ return wrapper
+
+ return decorator
+
+
+# Taken from: https://github.com/huggingface/transformers/blob/3658488ff77ff8d45101293e749263acf437f4d5/src/transformers/testing_utils.py#L1787
+def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None):
+ """
+ To run a test in a subprocess. In particular, this can avoid (GPU) memory issue.
+
+ Args:
+ test_case (`unittest.TestCase`):
+ The test that will run `target_func`.
+ target_func (`Callable`):
+ The function implementing the actual testing logic.
+ inputs (`dict`, *optional*, defaults to `None`):
+ The inputs that will be passed to `target_func` through an (input) queue.
+ timeout (`int`, *optional*, defaults to `None`):
+ The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env.
+ variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`.
+ """
+ if timeout is None:
+ timeout = int(os.environ.get("PYTEST_TIMEOUT", 600))
+
+ start_methohd = "spawn"
+ ctx = multiprocessing.get_context(start_methohd)
+
+ input_queue = ctx.Queue(1)
+ output_queue = ctx.JoinableQueue(1)
+
+ # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle.
+ input_queue.put(inputs, timeout=timeout)
+
+ process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout))
+ process.start()
+ # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents
+ # the test to exit properly.
+ try:
+ results = output_queue.get(timeout=timeout)
+ output_queue.task_done()
+ except Exception as e:
+ process.terminate()
+ test_case.fail(e)
+ process.join(timeout=timeout)
+
+ if results["error"] is not None:
+ test_case.fail(f"{results['error']}")
+
+
+class CaptureLogger:
+ """
+ Args:
+ Context manager to capture `logging` streams
+ logger: 'logging` logger object
+ Returns:
+ The captured output is available via `self.out`
+ Example:
+ ```python
+ >>> from diffusers import logging
+ >>> from diffusers.testing_utils import CaptureLogger
+
+ >>> msg = "Testing 1, 2, 3"
+ >>> logging.set_verbosity_info()
+ >>> logger = logging.get_logger("diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py")
+ >>> with CaptureLogger(logger) as cl:
+ ... logger.info(msg)
+ >>> assert cl.out, msg + "\n"
+ ```
+ """
+
+ def __init__(self, logger):
+ self.logger = logger
+ self.io = StringIO()
+ self.sh = logging.StreamHandler(self.io)
+ self.out = ""
+
+ def __enter__(self):
+ self.logger.addHandler(self.sh)
+ return self
+
+ def __exit__(self, *exc):
+ self.logger.removeHandler(self.sh)
+ self.out = self.io.getvalue()
+
+ def __repr__(self):
+ return f"captured: {self.out}\n"
+
+
+def enable_full_determinism():
+ """
+ Helper function for reproducible behavior during distributed training. See
+ - https://pytorch.org/docs/stable/notes/randomness.html for pytorch
+ """
+ from .torch_utils import enable_full_determinism as _enable_full_determinism
+
+ logger.warning(
+ "enable_full_determinism has been moved to diffusers.utils.torch_utils. "
+ "Importing from diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _enable_full_determinism()
+
+
+def disable_full_determinism():
+ from .torch_utils import disable_full_determinism as _disable_full_determinism
+
+ logger.warning(
+ "disable_full_determinism has been moved to diffusers.utils.torch_utils. "
+ "Importing from diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _disable_full_determinism()
+
+
+# Utils for custom and alternative accelerator devices
+def _is_torch_fp16_available(device):
+ if not is_torch_available():
+ return False
+
+ import torch
+
+ device = torch.device(device)
+
+ try:
+ x = torch.zeros((2, 2), dtype=torch.float16).to(device)
+ _ = torch.mul(x, x)
+ return True
+
+ except Exception as e:
+ if device.type == "cuda":
+ raise ValueError(
+ f"You have passed a device of type 'cuda' which should work with 'fp16', but 'cuda' does not seem to be correctly installed on your machine: {e}"
+ )
+
+ return False
+
+
+def _is_torch_fp64_available(device):
+ if not is_torch_available():
+ return False
+
+ import torch
+
+ device = torch.device(device)
+
+ try:
+ x = torch.zeros((2, 2), dtype=torch.float64).to(device)
+ _ = torch.mul(x, x)
+ return True
+
+ except Exception as e:
+ if device.type == "cuda":
+ raise ValueError(
+ f"You have passed a device of type 'cuda' which should work with 'fp64', but 'cuda' does not seem to be correctly installed on your machine: {e}"
+ )
+
+ return False
+
+
+# Guard these lookups for when Torch is not used - alternative accelerator support is for PyTorch
+if is_torch_available():
+ # Behaviour flags
+ BACKEND_SUPPORTS_TRAINING = {"cuda": True, "xpu": True, "cpu": True, "mps": False, "default": True}
+
+ # Function definitions
+ BACKEND_EMPTY_CACHE = {
+ "cuda": torch.cuda.empty_cache,
+ "xpu": torch.xpu.empty_cache,
+ "cpu": None,
+ "mps": torch.mps.empty_cache,
+ "default": None,
+ }
+ BACKEND_DEVICE_COUNT = {
+ "cuda": torch.cuda.device_count,
+ "xpu": torch.xpu.device_count,
+ "cpu": lambda: 0,
+ "mps": lambda: 0,
+ "default": 0,
+ }
+ BACKEND_MANUAL_SEED = {
+ "cuda": torch.cuda.manual_seed,
+ "xpu": torch.xpu.manual_seed,
+ "cpu": torch.manual_seed,
+ "mps": torch.mps.manual_seed,
+ "default": torch.manual_seed,
+ }
+ BACKEND_RESET_PEAK_MEMORY_STATS = {
+ "cuda": torch.cuda.reset_peak_memory_stats,
+ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
+ "cpu": None,
+ "mps": None,
+ "default": None,
+ }
+ BACKEND_RESET_MAX_MEMORY_ALLOCATED = {
+ "cuda": torch.cuda.reset_max_memory_allocated,
+ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
+ "cpu": None,
+ "mps": None,
+ "default": None,
+ }
+ BACKEND_MAX_MEMORY_ALLOCATED = {
+ "cuda": torch.cuda.max_memory_allocated,
+ "xpu": getattr(torch.xpu, "max_memory_allocated", None),
+ "cpu": 0,
+ "mps": 0,
+ "default": 0,
+ }
+ BACKEND_SYNCHRONIZE = {
+ "cuda": torch.cuda.synchronize,
+ "xpu": getattr(torch.xpu, "synchronize", None),
+ "cpu": None,
+ "mps": None,
+ "default": None,
+ }
+
+
+# This dispatches a defined function according to the accelerator from the function definitions.
+def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs):
+ if device not in dispatch_table:
+ return dispatch_table["default"](*args, **kwargs)
+
+ fn = dispatch_table[device]
+
+ # Some device agnostic functions return values. Need to guard against 'None' instead at
+ # user level
+ if not callable(fn):
+ return fn
+
+ return fn(*args, **kwargs)
+
+
+# These are callables which automatically dispatch the function specific to the accelerator
+def backend_manual_seed(device: str, seed: int):
+ from .torch_utils import backend_manual_seed as _backend_manual_seed
+
+ logger.warning(
+ "backend_manual_seed has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_manual_seed(device, seed)
+
+
+def backend_synchronize(device: str):
+ from .torch_utils import backend_synchronize as _backend_synchronize
+
+ logger.warning(
+ "backend_synchronize has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_synchronize(device)
+
+
+def backend_empty_cache(device: str):
+ from .torch_utils import backend_empty_cache as _backend_empty_cache
+
+ logger.warning(
+ "backend_empty_cache has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_empty_cache(device)
+
+
+def backend_device_count(device: str):
+ from .torch_utils import backend_device_count as _backend_device_count
+
+ logger.warning(
+ "backend_device_count has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_device_count(device)
+
+
+def backend_reset_peak_memory_stats(device: str):
+ from .torch_utils import backend_reset_peak_memory_stats as _backend_reset_peak_memory_stats
+
+ logger.warning(
+ "backend_reset_peak_memory_stats has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_reset_peak_memory_stats(device)
+
+
+def backend_reset_max_memory_allocated(device: str):
+ from .torch_utils import backend_reset_max_memory_allocated as _backend_reset_max_memory_allocated
+
+ logger.warning(
+ "backend_reset_max_memory_allocated has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_reset_max_memory_allocated(device)
+
+
+def backend_max_memory_allocated(device: str):
+ from .torch_utils import backend_max_memory_allocated as _backend_max_memory_allocated
+
+ logger.warning(
+ "backend_max_memory_allocated has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_max_memory_allocated(device)
+
+
+# These are callables which return boolean behaviour flags and can be used to specify some
+# device agnostic alternative where the feature is unsupported.
+def backend_supports_training(device: str):
+ from .torch_utils import backend_supports_training as _backend_supports_training
+
+ logger.warning(
+ "backend_supports_training has been moved to diffusers.utils.torch_utils. "
+ "diffusers.utils.testing_utils is deprecated and will be removed in a future version."
+ )
+ return _backend_supports_training(device)
+
+
+# Guard for when Torch is not available
+if is_torch_available():
+ # Update device function dict mapping
+ def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str):
+ try:
+ # Try to import the function directly
+ spec_fn = getattr(device_spec_module, attribute_name)
+ device_fn_dict[torch_device] = spec_fn
+ except AttributeError as e:
+ # If the function doesn't exist, and there is no default, throw an error
+ if "default" not in device_fn_dict:
+ raise AttributeError(
+ f"`{attribute_name}` not found in '{device_spec_path}' and no default fallback function found."
+ ) from e
+
+ if "DIFFUSERS_TEST_DEVICE_SPEC" in os.environ:
+ device_spec_path = os.environ["DIFFUSERS_TEST_DEVICE_SPEC"]
+ if not Path(device_spec_path).is_file():
+ raise ValueError(f"Specified path to device specification file is not found. Received {device_spec_path}")
+
+ try:
+ import_name = device_spec_path[: device_spec_path.index(".py")]
+ except ValueError as e:
+ raise ValueError(f"Provided device spec file is not a Python file! Received {device_spec_path}") from e
+
+ device_spec_module = importlib.import_module(import_name)
+
+ try:
+ device_name = device_spec_module.DEVICE_NAME
+ except AttributeError:
+ raise AttributeError("Device spec file did not contain `DEVICE_NAME`")
+
+ if "DIFFUSERS_TEST_DEVICE" in os.environ and torch_device != device_name:
+ msg = f"Mismatch between environment variable `DIFFUSERS_TEST_DEVICE` '{torch_device}' and device found in spec '{device_name}'\n"
+ msg += "Either unset `DIFFUSERS_TEST_DEVICE` or ensure it matches device spec name."
+ raise ValueError(msg)
+
+ torch_device = device_name
+
+ # Add one entry here for each `BACKEND_*` dictionary.
+ update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN")
+ update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN")
+ update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN")
+ update_mapping_from_spec(BACKEND_SUPPORTS_TRAINING, "SUPPORTS_TRAINING")
+ update_mapping_from_spec(BACKEND_RESET_PEAK_MEMORY_STATS, "RESET_PEAK_MEMORY_STATS_FN")
+ update_mapping_from_spec(BACKEND_RESET_MAX_MEMORY_ALLOCATED, "RESET_MAX_MEMORY_ALLOCATED_FN")
+ update_mapping_from_spec(BACKEND_MAX_MEMORY_ALLOCATED, "MAX_MEMORY_ALLOCATED_FN")
+
+
+# Modified from https://github.com/huggingface/transformers/blob/cdfb018d0300fef3b07d9220f3efe9c2a9974662/src/transformers/testing_utils.py#L3090
+
+# Type definition of key used in `Expectations` class.
+DeviceProperties = tuple[str | None, int | None]
+
+
+@functools.lru_cache
+def get_device_properties() -> DeviceProperties:
+ """
+ Get environment device properties.
+ """
+ if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM:
+ import torch
+
+ major, _ = torch.cuda.get_device_capability()
+ if IS_ROCM_SYSTEM:
+ return ("rocm", major)
+ else:
+ return ("cuda", major)
+ elif IS_XPU_SYSTEM:
+ import torch
+
+ # To get more info of the architecture meaning and bit allocation, refer to https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/device_architecture.def
+ arch = torch.xpu.get_device_capability()["architecture"]
+ gen_mask = 0x000000FF00000000
+ gen = (arch & gen_mask) >> 32
+ return ("xpu", gen)
+ else:
+ return (torch_device, None)
+
+
+if TYPE_CHECKING:
+ DevicePropertiesUserDict = UserDict[DeviceProperties, Any]
+else:
+ DevicePropertiesUserDict = UserDict
+
+if is_torch_available():
+ from diffusers.hooks._common import _GO_LC_SUPPORTED_PYTORCH_LAYERS
+ from diffusers.hooks.group_offloading import (
+ _GROUP_ID_LAZY_LEAF,
+ _compute_group_hash,
+ _find_parent_module_in_module_dict,
+ _gather_buffers_with_no_group_offloading_parent,
+ _gather_parameters_with_no_group_offloading_parent,
+ )
+
+ def _get_expected_safetensors_files(
+ module: torch.nn.Module,
+ offload_to_disk_path: str,
+ offload_type: str,
+ num_blocks_per_group: int | None = None,
+ ) -> Set[str]:
+ expected_files = set()
+
+ def get_hashed_filename(group_id: str) -> str:
+ short_hash = _compute_group_hash(group_id)
+ return os.path.join(offload_to_disk_path, f"group_{short_hash}.safetensors")
+
+ if offload_type == "block_level":
+ if num_blocks_per_group is None:
+ raise ValueError("num_blocks_per_group must be provided for 'block_level' offloading.")
+
+ # Handle groups of ModuleList and Sequential blocks
+ unmatched_modules = []
+ for name, submodule in module.named_children():
+ if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
+ unmatched_modules.append(module)
+ continue
+
+ for i in range(0, len(submodule), num_blocks_per_group):
+ current_modules = submodule[i : i + num_blocks_per_group]
+ if not current_modules:
+ continue
+ group_id = f"{name}_{i}_{i + len(current_modules) - 1}"
+ expected_files.add(get_hashed_filename(group_id))
+
+ # Handle the group for unmatched top-level modules and parameters
+ for module in unmatched_modules:
+ expected_files.add(get_hashed_filename(f"{module.__class__.__name__}_unmatched_group"))
+
+ elif offload_type == "leaf_level":
+ # Handle leaf-level module groups
+ for name, submodule in module.named_modules():
+ if isinstance(submodule, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
+ # These groups will always have parameters, so a file is expected
+ expected_files.add(get_hashed_filename(name))
+
+ # Handle groups for non-leaf parameters/buffers
+ modules_with_group_offloading = {
+ name for name, sm in module.named_modules() if isinstance(sm, _GO_LC_SUPPORTED_PYTORCH_LAYERS)
+ }
+ parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
+ buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
+
+ all_orphans = parameters + buffers
+ if all_orphans:
+ parent_to_tensors = {}
+ module_dict = dict(module.named_modules())
+ for tensor_name, _ in all_orphans:
+ parent_name = _find_parent_module_in_module_dict(tensor_name, module_dict)
+ if parent_name not in parent_to_tensors:
+ parent_to_tensors[parent_name] = []
+ parent_to_tensors[parent_name].append(tensor_name)
+
+ for parent_name in parent_to_tensors:
+ # A file is expected for each parent that gathers orphaned tensors
+ expected_files.add(get_hashed_filename(parent_name))
+ expected_files.add(get_hashed_filename(_GROUP_ID_LAZY_LEAF))
+
+ else:
+ raise ValueError(f"Unsupported offload_type: {offload_type}")
+
+ return expected_files
+
+ def _check_safetensors_serialization(
+ module: torch.nn.Module,
+ offload_to_disk_path: str,
+ offload_type: str,
+ num_blocks_per_group: int | None = None,
+ ) -> bool:
+ if not os.path.isdir(offload_to_disk_path):
+ return False, None, None
+
+ expected_files = _get_expected_safetensors_files(
+ module, offload_to_disk_path, offload_type, num_blocks_per_group
+ )
+ actual_files = set(glob.glob(os.path.join(offload_to_disk_path, "*.safetensors")))
+ missing_files = expected_files - actual_files
+ extra_files = actual_files - expected_files
+
+ is_correct = not missing_files and not extra_files
+ return is_correct, extra_files, missing_files
+
+
+class Expectations(DevicePropertiesUserDict):
+ def get_expectation(self) -> Any:
+ """
+ Find best matching expectation based on environment device properties.
+ """
+ return self.find_expectation(get_device_properties())
+
+ @staticmethod
+ def is_default(key: DeviceProperties) -> bool:
+ return all(p is None for p in key)
+
+ @staticmethod
+ def score(key: DeviceProperties, other: DeviceProperties) -> int:
+ """
+ Returns score indicating how similar two instances of the `Properties` tuple are. Points are calculated using
+ bits, but documented as int. Rules are as follows:
+ * Matching `type` gives 8 points.
+ * Semi-matching `type`, for example cuda and rocm, gives 4 points.
+ * Matching `major` (compute capability major version) gives 2 points.
+ * Default expectation (if present) gives 1 points.
+ """
+ (device_type, major) = key
+ (other_device_type, other_major) = other
+
+ score = 0b0
+ if device_type == other_device_type:
+ score |= 0b1000
+ elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]:
+ score |= 0b100
+
+ if major == other_major and other_major is not None:
+ score |= 0b10
+
+ if Expectations.is_default(other):
+ score |= 0b1
+
+ return int(score)
+
+ def find_expectation(self, key: DeviceProperties = (None, None)) -> Any:
+ """
+ Find best matching expectation based on provided device properties.
+ """
+ (result_key, result) = max(self.data.items(), key=lambda x: Expectations.score(key, x[0]))
+
+ if Expectations.score(key, result_key) == 0:
+ raise ValueError(f"No matching expectation found for {key}")
+
+ return result
+
+ def __repr__(self):
+ return f"{self.data}"
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/torch_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/torch_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..165770d48b128433a468e8e669f84cb4cd82f679
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/torch_utils.py
@@ -0,0 +1,409 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+PyTorch utilities: Utilities related to PyTorch
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+from typing import Callable, ParamSpec, TypeVar
+
+from . import logging
+from .import_utils import (
+ is_torch_available,
+ is_torch_mlu_available,
+ is_torch_neuronx_available,
+ is_torch_npu_available,
+ is_torch_version,
+)
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+if is_torch_available():
+ import torch
+ from torch.fft import fftn, fftshift, ifftn, ifftshift
+
+ BACKEND_SUPPORTS_TRAINING = {
+ "cuda": True,
+ "xpu": True,
+ "cpu": True,
+ "mps": False,
+ "neuron": False,
+ "default": True,
+ }
+ BACKEND_EMPTY_CACHE = {
+ "cuda": torch.cuda.empty_cache,
+ "xpu": torch.xpu.empty_cache,
+ "cpu": None,
+ "mps": torch.mps.empty_cache,
+ "neuron": None,
+ "default": None,
+ }
+ BACKEND_DEVICE_COUNT = {
+ "cuda": torch.cuda.device_count,
+ "xpu": torch.xpu.device_count,
+ "cpu": lambda: 0,
+ "mps": lambda: 0,
+ "neuron": lambda: getattr(getattr(torch, "neuron", None), "device_count", lambda: 0)(),
+ "default": 0,
+ }
+ BACKEND_MANUAL_SEED = {
+ "cuda": torch.cuda.manual_seed,
+ "xpu": torch.xpu.manual_seed,
+ "cpu": torch.manual_seed,
+ "mps": torch.mps.manual_seed,
+ "neuron": torch.manual_seed,
+ "default": torch.manual_seed,
+ }
+ BACKEND_RESET_PEAK_MEMORY_STATS = {
+ "cuda": torch.cuda.reset_peak_memory_stats,
+ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
+ "cpu": None,
+ "mps": None,
+ "neuron": None,
+ "default": None,
+ }
+ BACKEND_RESET_MAX_MEMORY_ALLOCATED = {
+ "cuda": torch.cuda.reset_max_memory_allocated,
+ "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
+ "cpu": None,
+ "mps": None,
+ "neuron": None,
+ "default": None,
+ }
+ BACKEND_MAX_MEMORY_ALLOCATED = {
+ "cuda": torch.cuda.max_memory_allocated,
+ "xpu": getattr(torch.xpu, "max_memory_allocated", None),
+ "cpu": 0,
+ "mps": 0,
+ "neuron": 0,
+ "default": 0,
+ }
+ BACKEND_SYNCHRONIZE = {
+ "cuda": torch.cuda.synchronize,
+ "xpu": getattr(torch.xpu, "synchronize", None),
+ "cpu": None,
+ "mps": None,
+ "neuron": getattr(getattr(torch, "neuron", None), "synchronize", None),
+ "default": None,
+ }
+
+ _FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"})
+ _INT64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"})
+ _DTYPE_DOWNCAST = {torch.float64: torch.float32, torch.int64: torch.int32}
+ _DTYPE_UNSUPPORTED_DEVICES = {torch.float64: _FP64_UNSUPPORTED_DEVICES, torch.int64: _INT64_UNSUPPORTED_DEVICES}
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+try:
+ from torch._dynamo import allow_in_graph as maybe_allow_in_graph
+except (ImportError, ModuleNotFoundError):
+
+ def maybe_allow_in_graph(cls):
+ return cls
+
+
+# This dispatches a defined function according to the accelerator from the function definitions.
+def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, callable], *args, **kwargs):
+ if device not in dispatch_table:
+ return dispatch_table["default"](*args, **kwargs)
+
+ fn = dispatch_table[device]
+
+ # Some device agnostic functions return values. Need to guard against 'None' instead at
+ # user level
+ if not callable(fn):
+ return fn
+
+ return fn(*args, **kwargs)
+
+
+# These are callables which automatically dispatch the function specific to the accelerator
+def backend_manual_seed(device: str, seed: int):
+ return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed)
+
+
+def backend_synchronize(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE)
+
+
+def backend_empty_cache(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE)
+
+
+def backend_device_count(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT)
+
+
+def backend_reset_peak_memory_stats(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS)
+
+
+def backend_reset_max_memory_allocated(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED)
+
+
+def backend_max_memory_allocated(device: str):
+ return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED)
+
+
+# These are callables which return boolean behaviour flags and can be used to specify some
+# device agnostic alternative where the feature is unsupported.
+def backend_supports_training(device: str):
+ if not is_torch_available():
+ return False
+
+ if device not in BACKEND_SUPPORTS_TRAINING:
+ device = "default"
+
+ return BACKEND_SUPPORTS_TRAINING[device]
+
+
+def maybe_adjust_dtype_for_device(dtype: "torch.dtype", device: "torch.device") -> "torch.dtype":
+ unsupported = _DTYPE_UNSUPPORTED_DEVICES.get(dtype)
+ return _DTYPE_DOWNCAST[dtype] if unsupported and device.type in unsupported else dtype
+
+
+def randn_tensor(
+ shape: tuple | list,
+ generator: list["torch.Generator"] | "torch.Generator" | None = None,
+ device: str | "torch.device" | None = None,
+ dtype: "torch.dtype" | None = None,
+ layout: "torch.layout" | None = None,
+):
+ """A helper function to create random tensors on the desired `device` with the desired `dtype`. When
+ passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor
+ is always created on the CPU.
+ """
+ # device on which tensor is created defaults to device
+ if isinstance(device, str):
+ device = torch.device(device)
+ rand_device = device
+ batch_size = shape[0]
+
+ layout = layout or torch.strided
+ device = device or torch.device("cpu")
+
+ # Neuron (XLA) does not support creating random tensors directly on device; always use CPU
+ if device.type == "neuron":
+ rand_device = torch.device("cpu")
+
+ if generator is not None:
+ gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type
+ if gen_device_type != device.type and gen_device_type == "cpu":
+ rand_device = "cpu"
+ if device.type not in ("mps", "neuron"):
+ logger.info(
+ f"The passed generator was created on 'cpu' even though a tensor on {device} was expected."
+ f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"
+ f" slightly speed up this function by passing a generator that was created on the {device} device."
+ )
+ elif gen_device_type != device.type and gen_device_type == "cuda":
+ raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.")
+
+ # make sure generator list of length 1 is treated like a non-list
+ if isinstance(generator, list) and len(generator) == 1:
+ generator = generator[0]
+
+ if isinstance(generator, list):
+ shape = (1,) + shape[1:]
+ latents = [
+ torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout)
+ for i in range(batch_size)
+ ]
+ latents = torch.cat(latents, dim=0).to(device)
+ else:
+ latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)
+
+ return latents
+
+
+def is_compiled_module(module) -> bool:
+ """Check whether the module was compiled with torch.compile()"""
+ if is_torch_version("<", "2.0.0") or not hasattr(torch, "_dynamo"):
+ return False
+ return isinstance(module, torch._dynamo.eval_frame.OptimizedModule)
+
+
+def unwrap_module(module):
+ """Unwraps a module if it was compiled with torch.compile()"""
+ return module._orig_mod if is_compiled_module(module) else module
+
+
+def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.Tensor":
+ """Fourier filter as introduced in FreeU (https://huggingface.co/papers/2309.11497).
+
+ This version of the method comes from here:
+ https://github.com/huggingface/diffusers/pull/5164#issuecomment-1732638706
+ """
+ x = x_in
+ B, C, H, W = x.shape
+
+ # Non-power of 2 images must be float32
+ if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
+ x = x.to(dtype=torch.float32)
+ # fftn does not support bfloat16, and produces the experimental ComplexHalf
+ # dtype (torch.complex32) when given float16, which is numerically unstable
+ # and triggers a UserWarning. Upcast any non-float32 dtype to float32.
+ elif x.dtype != torch.float32:
+ x = x.to(dtype=torch.float32)
+
+ # FFT
+ x_freq = fftn(x, dim=(-2, -1))
+ x_freq = fftshift(x_freq, dim=(-2, -1))
+
+ B, C, H, W = x_freq.shape
+ mask = torch.ones((B, C, H, W), device=x.device)
+
+ crow, ccol = H // 2, W // 2
+ mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scale
+ x_freq = x_freq * mask
+
+ # IFFT
+ x_freq = ifftshift(x_freq, dim=(-2, -1))
+ x_filtered = ifftn(x_freq, dim=(-2, -1)).real
+
+ return x_filtered.to(dtype=x_in.dtype)
+
+
+def apply_freeu(
+ resolution_idx: int, hidden_states: "torch.Tensor", res_hidden_states: "torch.Tensor", **freeu_kwargs
+) -> tuple["torch.Tensor", "torch.Tensor"]:
+ """Applies the FreeU mechanism as introduced in https:
+ //arxiv.org/abs/2309.11497. Adapted from the official code repository: https://github.com/ChenyangSi/FreeU.
+
+ Args:
+ resolution_idx (`int`): Integer denoting the UNet block where FreeU is being applied.
+ hidden_states (`torch.Tensor`): Inputs to the underlying block.
+ res_hidden_states (`torch.Tensor`): Features from the skip block corresponding to the underlying block.
+ s1 (`float`): Scaling factor for stage 1 to attenuate the contributions of the skip features.
+ s2 (`float`): Scaling factor for stage 2 to attenuate the contributions of the skip features.
+ b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
+ b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
+ """
+ if resolution_idx == 0:
+ num_half_channels = hidden_states.shape[1] // 2
+ hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b1"]
+ res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s1"])
+ if resolution_idx == 1:
+ num_half_channels = hidden_states.shape[1] // 2
+ hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b2"]
+ res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s2"])
+
+ return hidden_states, res_hidden_states
+
+
+def get_torch_cuda_device_capability():
+ if torch.cuda.is_available():
+ device = torch.device("cuda")
+ compute_capability = torch.cuda.get_device_capability(device)
+ compute_capability = f"{compute_capability[0]}.{compute_capability[1]}"
+ return float(compute_capability)
+ else:
+ return None
+
+
+@functools.lru_cache
+def get_device():
+ if torch.cuda.is_available():
+ return "cuda"
+ elif is_torch_npu_available():
+ return "npu"
+ elif hasattr(torch, "xpu") and torch.xpu.is_available():
+ return "xpu"
+ elif torch.backends.mps.is_available():
+ return "mps"
+ elif is_torch_mlu_available():
+ return "mlu"
+ elif is_torch_neuronx_available() and hasattr(torch, "neuron") and torch.neuron.is_available():
+ return "neuron"
+ else:
+ return "cpu"
+
+
+def empty_device_cache(device_type: str | None = None):
+ if device_type is None:
+ device_type = get_device()
+ if device_type in ["cpu"]:
+ return
+ device_mod = getattr(torch, device_type, torch.cuda)
+ device_mod.empty_cache()
+
+
+def device_synchronize(device_type: str | None = None):
+ if device_type is None:
+ device_type = get_device()
+ device_mod = getattr(torch, device_type, torch.cuda)
+ device_mod.synchronize()
+
+
+def enable_full_determinism():
+ """
+ Helper function for reproducible behavior during distributed training. See
+ - https://pytorch.org/docs/stable/notes/randomness.html for pytorch
+ """
+ # Enable PyTorch deterministic mode. This potentially requires either the environment
+ # variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,
+ # depending on the CUDA version, so we set them both here
+ os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
+ torch.use_deterministic_algorithms(True)
+
+ # Enable CUDNN deterministic mode
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+ torch.backends.cuda.matmul.allow_tf32 = False
+
+
+def disable_full_determinism():
+ os.environ["CUDA_LAUNCH_BLOCKING"] = "0"
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ""
+ torch.use_deterministic_algorithms(False)
+
+
+@functools.wraps(functools.lru_cache)
+def lru_cache_unless_export(maxsize=128, typed=False):
+ def outer_wrapper(fn: Callable[P, T]):
+ cached = functools.lru_cache(maxsize=maxsize, typed=typed)(fn)
+ if is_torch_version("<", "2.7.0"):
+ return cached
+
+ @functools.wraps(fn)
+ def inner_wrapper(*args: P.args, **kwargs: P.kwargs):
+ compiler = getattr(torch, "compiler", None)
+ is_exporting = bool(compiler and hasattr(compiler, "is_exporting") and compiler.is_exporting())
+ is_compiling = bool(compiler and hasattr(compiler, "is_compiling") and compiler.is_compiling())
+
+ # Fallback for older builds where compiler.is_compiling is unavailable.
+ if not is_compiling:
+ dynamo = getattr(torch, "_dynamo", None)
+ if dynamo is not None and hasattr(dynamo, "is_compiling"):
+ is_compiling = dynamo.is_compiling()
+
+ if is_exporting or is_compiling:
+ return fn(*args, **kwargs)
+ return cached(*args, **kwargs)
+
+ return inner_wrapper
+
+ return outer_wrapper
+
+
+if is_torch_available():
+ torch_device = get_device()
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/typing_utils.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/typing_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d106fe635b3a3f559facfaea8e6aeb86ac970665
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/typing_utils.py
@@ -0,0 +1,91 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Typing utilities: Utilities related to type checking and validation
+"""
+
+from typing import Any, Dict, List, Set, Tuple, Type, Union, get_args, get_origin
+
+
+def _is_valid_type(obj: Any, class_or_tuple: Type | tuple[Type, ...]) -> bool:
+ """
+ Checks if an object is an instance of any of the provided types. For collections, it checks if every element is of
+ the correct type as well.
+ """
+ if not isinstance(class_or_tuple, tuple):
+ class_or_tuple = (class_or_tuple,)
+
+ # Unpack unions
+ unpacked_class_or_tuple = []
+ for t in class_or_tuple:
+ if get_origin(t) is Union:
+ unpacked_class_or_tuple.extend(get_args(t))
+ else:
+ unpacked_class_or_tuple.append(t)
+ class_or_tuple = tuple(unpacked_class_or_tuple)
+
+ if Any in class_or_tuple:
+ return True
+
+ obj_type = type(obj)
+ # Classes with obj's type
+ class_or_tuple = {t for t in class_or_tuple if isinstance(obj, get_origin(t) or t)}
+
+ # Singular types (e.g. int, ControlNet, ...)
+ # Untyped collections (e.g. List, but not List[int])
+ elem_class_or_tuple = {get_args(t) for t in class_or_tuple}
+ if () in elem_class_or_tuple:
+ return True
+ # Typed lists or sets
+ elif obj_type in (list, set):
+ return any(all(_is_valid_type(x, t) for x in obj) for t in elem_class_or_tuple)
+ # Typed tuples
+ elif obj_type is tuple:
+ return any(
+ # Tuples with any length and single type (e.g. Tuple[int, ...])
+ (len(t) == 2 and t[-1] is Ellipsis and all(_is_valid_type(x, t[0]) for x in obj))
+ or
+ # Tuples with fixed length and any types (e.g. Tuple[int, str])
+ (len(obj) == len(t) and all(_is_valid_type(x, tt) for x, tt in zip(obj, t)))
+ for t in elem_class_or_tuple
+ )
+ # Typed dicts
+ elif obj_type is dict:
+ return any(
+ all(_is_valid_type(k, kt) and _is_valid_type(v, vt) for k, v in obj.items())
+ for kt, vt in elem_class_or_tuple
+ )
+
+ else:
+ return False
+
+
+def _get_detailed_type(obj: Any) -> Type:
+ """
+ Gets a detailed type for an object, including nested types for collections.
+ """
+ obj_type = type(obj)
+
+ if obj_type in (list, set):
+ obj_origin_type = List if obj_type is list else Set
+ elems_type = Union[tuple({_get_detailed_type(x) for x in obj})]
+ return obj_origin_type[elems_type]
+ elif obj_type is tuple:
+ return Tuple[tuple(_get_detailed_type(x) for x in obj)]
+ elif obj_type is dict:
+ keys_type = Union[tuple({_get_detailed_type(k) for k in obj.keys()})]
+ values_type = Union[tuple({_get_detailed_type(k) for k in obj.values()})]
+ return Dict[keys_type, values_type]
+ else:
+ return obj_type
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/versions.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/versions.py
new file mode 100644
index 0000000000000000000000000000000000000000..466fc4019a28d5513d75f9c068bd4af4d807e08a
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/utils/versions.py
@@ -0,0 +1,116 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Utilities for working with package versions
+"""
+
+import importlib.metadata
+import operator
+import re
+import sys
+
+from packaging import version
+
+
+ops = {
+ "<": operator.lt,
+ "<=": operator.le,
+ "==": operator.eq,
+ "!=": operator.ne,
+ ">=": operator.ge,
+ ">": operator.gt,
+}
+
+
+def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint):
+ if got_ver is None or want_ver is None:
+ raise ValueError(
+ f"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider"
+ f" reinstalling {pkg}."
+ )
+ if not ops[op](version.parse(got_ver), version.parse(want_ver)):
+ raise ImportError(
+ f"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}"
+ )
+
+
+def require_version(requirement: str, hint: str | None = None) -> None:
+ """
+ Perform a runtime check of the dependency versions, using the exact same syntax used by pip.
+
+ The installed module version comes from the *site-packages* dir via *importlib.metadata*.
+
+ Args:
+ requirement (`str`): pip style definition, e.g., "tokenizers==0.9.4", "tqdm>=4.27", "numpy"
+ hint (`str`, *optional*): what suggestion to print in case of requirements not being met
+
+ Example:
+
+ ```python
+ require_version("pandas>1.1.2")
+ require_version("numpy>1.18.5", "this is important to have for whatever reason")
+ ```"""
+
+ hint = f"\n{hint}" if hint is not None else ""
+
+ # non-versioned check
+ if re.match(r"^[\w_\-\d]+$", requirement):
+ pkg, op, want_ver = requirement, None, None
+ else:
+ match = re.findall(r"^([^!=<>\s]+)([\s!=<>]{1,2}.+)", requirement)
+ if not match:
+ raise ValueError(
+ "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"
+ f" got {requirement}"
+ )
+ pkg, want_full = match[0]
+ want_range = want_full.split(",") # there could be multiple requirements
+ wanted = {}
+ for w in want_range:
+ match = re.findall(r"^([\s!=<>]{1,2})(.+)", w)
+ if not match:
+ raise ValueError(
+ "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"
+ f" but got {requirement}"
+ )
+ op, want_ver = match[0]
+ wanted[op] = want_ver
+ if op not in ops:
+ raise ValueError(f"{requirement}: need one of {list(ops.keys())}, but got {op}")
+
+ # special case
+ if pkg == "python":
+ got_ver = ".".join([str(x) for x in sys.version_info[:3]])
+ for op, want_ver in wanted.items():
+ _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)
+ return
+
+ # check if any version is installed
+ try:
+ got_ver = importlib.metadata.version(pkg)
+ except importlib.metadata.PackageNotFoundError:
+ raise importlib.metadata.PackageNotFoundError(
+ f"The '{requirement}' distribution was not found and is required by this application. {hint}"
+ )
+
+ # check that the right version is installed if version number or a range was provided
+ if want_ver is not None:
+ for op, want_ver in wanted.items():
+ _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)
+
+
+def require_version_core(requirement):
+ """require_version wrapper which emits a core-specific hint on failure"""
+ hint = "Try: pip install transformers -U or pip install -e '.[dev]' if you're working with git main"
+ return require_version(requirement, hint)
diff --git a/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/video_processor.py b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/video_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f6b442ee2fe9fa7de3f59fd2a2c1bc59dd1ee58
--- /dev/null
+++ b/.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/video_processor.py
@@ -0,0 +1,180 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+import numpy as np
+import PIL
+import torch
+import torch.nn.functional as F
+
+from .image_processor import VaeImageProcessor, is_valid_image, is_valid_image_imagelist
+
+
+class VideoProcessor(VaeImageProcessor):
+ r"""Simple video processor."""
+
+ def preprocess_video(self, video, height: int | None = None, width: int | None = None, **kwargs) -> torch.Tensor:
+ r"""
+ Preprocesses input video(s). Keyword arguments will be forwarded to `VaeImageProcessor.preprocess`.
+
+ Args:
+ video (`list[PIL.Image]`, `list[list[PIL.Image]]`, `torch.Tensor`, `np.array`, `list[torch.Tensor]`, `list[np.array]`):
+ The input video. It can be one of the following:
+ * list of the PIL images.
+ * list of list of PIL images.
+ * 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height, width)`).
+ * 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`).
+ * list of 4D Torch tensors (expected shape for each tensor `(num_frames, num_channels, height,
+ width)`).
+ * list of 4D NumPy arrays (expected shape for each array `(num_frames, height, width, num_channels)`).
+ * 5D NumPy arrays: expected shape for each array `(batch_size, num_frames, height, width,
+ num_channels)`.
+ * 5D Torch tensors: expected shape for each array `(batch_size, num_frames, num_channels, height,
+ width)`.
+ height (`int`, *optional*, defaults to `None`):
+ The height in preprocessed frames of the video. If `None`, will use the `get_default_height_width()` to
+ get default height.
+ width (`int`, *optional*`, defaults to `None`):
+ The width in preprocessed frames of the video. If `None`, will use get_default_height_width()` to get
+ the default width.
+
+ Returns:
+ `torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`:
+ A 5D tensor holding the batched channels-first video(s).
+ """
+ if isinstance(video, list) and isinstance(video[0], np.ndarray) and video[0].ndim == 5:
+ warnings.warn(
+ "Passing `video` as a list of 5d np.ndarray is deprecated."
+ "Please concatenate the list along the batch dimension and pass it as a single 5d np.ndarray",
+ FutureWarning,
+ )
+ video = np.concatenate(video, axis=0)
+ if isinstance(video, list) and isinstance(video[0], torch.Tensor) and video[0].ndim == 5:
+ warnings.warn(
+ "Passing `video` as a list of 5d torch.Tensor is deprecated."
+ "Please concatenate the list along the batch dimension and pass it as a single 5d torch.Tensor",
+ FutureWarning,
+ )
+ video = torch.cat(video, axis=0)
+
+ # ensure the input is a list of videos:
+ # - if it is a batch of videos (5d torch.Tensor or np.ndarray), it is converted to a list of videos (a list of 4d torch.Tensor or np.ndarray)
+ # - if it is a single video, it is converted to a list of one video.
+ if isinstance(video, (np.ndarray, torch.Tensor)) and video.ndim == 5:
+ video = list(video)
+ elif isinstance(video, list) and is_valid_image(video[0]) or is_valid_image_imagelist(video):
+ video = [video]
+ elif isinstance(video, list) and is_valid_image_imagelist(video[0]):
+ video = video
+ else:
+ raise ValueError(
+ "Input is in incorrect format. Currently, we only support numpy.ndarray, torch.Tensor, PIL.Image.Image"
+ )
+
+ video = torch.stack([self.preprocess(img, height=height, width=width, **kwargs) for img in video], dim=0)
+
+ # move the number of channels before the number of frames.
+ video = video.permute(0, 2, 1, 3, 4)
+
+ return video
+
+ def postprocess_video(
+ self, video: torch.Tensor, output_type: str = "np", **kwargs
+ ) -> np.ndarray | torch.Tensor | list[PIL.Image.Image]:
+ r"""
+ Converts a video tensor to a list of frames for export. Keyword arguments will be forwarded to
+ `VaeImageProcessor.postprocess`.
+
+ Args:
+ video (`torch.Tensor`): The video as a tensor.
+ output_type (`str`, defaults to `"np"`): Output type of the postprocessed `video` tensor.
+ """
+ batch_size = video.shape[0]
+ outputs = []
+ for batch_idx in range(batch_size):
+ batch_vid = video[batch_idx].permute(1, 0, 2, 3)
+ batch_output = self.postprocess(batch_vid, output_type, **kwargs)
+ outputs.append(batch_output)
+
+ if output_type == "np":
+ outputs = np.stack(outputs)
+ elif output_type == "pt":
+ outputs = torch.stack(outputs)
+ elif not output_type == "pil":
+ raise ValueError(f"{output_type} does not exist. Please choose one of ['np', 'pt', 'pil']")
+
+ return outputs
+
+ @staticmethod
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
+ r"""
+ Returns the binned height and width based on the aspect ratio.
+
+ Args:
+ height (`int`): The height of the image.
+ width (`int`): The width of the image.
+ ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
+
+ Returns:
+ `tuple[int, int]`: The closest binned height and width.
+ """
+ ar = float(height / width)
+ closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
+ default_hw = ratios[closest_ratio]
+ return int(default_hw[0]), int(default_hw[1])
+
+ @staticmethod
+ def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
+ r"""
+ Resizes and crops a tensor of videos to the specified dimensions.
+
+ Args:
+ samples (`torch.Tensor`):
+ A tensor of shape (N, C, T, H, W) where N is the batch size, C is the number of channels, T is the
+ number of frames, H is the height, and W is the width.
+ new_width (`int`): The desired width of the output videos.
+ new_height (`int`): The desired height of the output videos.
+
+ Returns:
+ `torch.Tensor`: A tensor containing the resized and cropped videos.
+ """
+ orig_height, orig_width = samples.shape[3], samples.shape[4]
+
+ # Check if resizing is needed
+ if orig_height != new_height or orig_width != new_width:
+ ratio = max(new_height / orig_height, new_width / orig_width)
+ resized_width = int(orig_width * ratio)
+ resized_height = int(orig_height * ratio)
+
+ # Reshape to (N*T, C, H, W) for interpolation
+ n, c, t, h, w = samples.shape
+ samples = samples.permute(0, 2, 1, 3, 4).reshape(n * t, c, h, w)
+
+ # Resize
+ samples = F.interpolate(
+ samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
+ )
+
+ # Center Crop
+ start_x = (resized_width - new_width) // 2
+ end_x = start_x + new_width
+ start_y = (resized_height - new_height) // 2
+ end_y = start_y + new_height
+ samples = samples[:, :, start_y:end_y, start_x:end_x]
+
+ # Reshape back to (N, C, T, H, W)
+ samples = samples.reshape(n, t, c, new_height, new_width).permute(0, 2, 1, 3, 4)
+
+ return samples
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/.gitattributes b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..dab9a4e17afd2ef39d90ccb0b40ef2786fe77422
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/.gitattributes
@@ -0,0 +1,35 @@
+*.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
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/README.md b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..947d3ce5bcf29f412de0f9af4e5b2d0b1307ae12
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/README.md
@@ -0,0 +1,36 @@
+---
+title: Qwen3-VL conditioner
+emoji: 🧠
+colorFrom: indigo
+colorTo: blue
+sdk: gradio
+sdk_version: 6.20.0
+app_file: app.py
+pinned: false
+short_description: Qwen3-VL 33B prompt conditioning as a service
+suggested_hardware: zero-a10g
+---
+
+# Qwen3-VL conditioner
+
+Encodes prompts (optionally with keyframes or ordered media references) into conditioning tensors over
+the gradio API, running the 33B Qwen3-VL conditioner unquantized in bfloat16 on ZeroGPU. Built as the
+text encoding half of a split deployment for large video generation models.
+
+Both endpoints take an optional trailing `rewrite_prompt` (default `False`). Turned on, the conditioner
+first writes the structured prompt format the video model was trained on — a deterministic task
+instruction line, then a shot-by-shot timeline and the two sound fields, written by the conditioner's own
+language-model head while it looks at the keyframes or references — and encodes that instead. The rewrite
+comes back under the plan's `refined_prompt`, and a rewrite that fails falls back to encoding the raw
+prompt rather than failing the call.
+
+## Where the code comes from
+
+The weights are the public [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) diffusers
+checkpoint — only `text_encoder/`, `tokenizer/` and `processor/` are ever fetched here.
+
+`diffusers` is installed from the canonical MiniMax-H3 pull request,
+[huggingface/diffusers#14371](https://github.com/huggingface/diffusers/pull/14371), pinned to the commit `665f5782`
+(`refs/pull/14371/head`) rather than to the moving `minimax-h3-refactor` branch. That PR is a WIP, so it needs
+re-pinning whenever it updates, and `h3_split_blocks.py` — which subclasses its block classes to cut the pipeline in
+two — has to be re-checked against the new head at the same time.
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/__pycache__/app.cpython-310.pyc b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/__pycache__/app.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8135539bc6a060438b4dba2ac111fc6a9633e58c
Binary files /dev/null and b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/__pycache__/app.cpython-310.pyc differ
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/app.py b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..99a91cf7d0d8b3ccff17fa6dd1edef6ee8f89135
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/app.py
@@ -0,0 +1,527 @@
+"""Qwen3-VL conditioner service.
+
+Encodes a text prompt (optionally with keyframe images or ordered media references) into the embedding
+tensors a downstream video generation Space consumes: `prompt_embeds` and per row token tags, plus the
+resolved generation plan. Runs the 33B Qwen3-VL conditioner unquantized in bfloat16 on ZeroGPU.
+
+Optionally rewrites the prompt first. The conditioner is the full `Qwen3VLForConditionalGeneration`, language-model
+head and chat template included, so the structured prompt format the video model was actually trained on can be
+written here, by the same weights that are about to encode it — see `prompt_rewrite`. It is off by default and
+strictly additive: `rewrite_prompt=False` is the path that was here before it, unchanged.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+import time
+import traceback
+
+# First, and at module level. On ZeroGPU `import spaces` patches `torch.cuda` so that `is_available()` is True and
+# `get_device_capability()` answers the pool's own (12, 0) before any GPU is attached — which is what lets the whole
+# 62.14 GiB load happen at **startup**, off GPU time. It also has to precede anything that initializes CUDA.
+import spaces
+import gradio as gr
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "MiniMaxAI/MiniMax-H3") # public; overridable as a Space variable
+# A rewrite and an encode are booked separately, and both bookings are small. An `xlarge` booking costs twice the
+# seconds it reserves, and a caller ZeroGPU cannot attribute may book at most 120 credits at a time — 60 s of
+# `xlarge` — so a single booking covering both would be refused before any work happened. One forward through 50
+# decoder layers is 2 to 8 seconds.
+GPU_DURATION = int(os.environ.get("GPU_DURATION", "45"))
+GPU_SIZE = os.environ.get("GPU_SIZE", "xlarge")
+ON_SPACES = bool(os.environ.get("SPACE_ID"))
+# The decode budget of a rewrite. A base rewrite lands around 250 tokens and a six-section `ref2va` one around 600;
+# a decode that runs into this budget is trimmed back to its last finished sentence rather than fed in half.
+REWRITE_MAX_NEW_TOKENS = int(os.environ.get("REWRITE_MAX_NEW_TOKENS", "1000"))
+# The rewrite's own booking: one ~10k token prefill plus the decode, 15 to 50 seconds on the pool's card.
+REWRITE_GPU_DURATION = int(os.environ.get("REWRITE_GPU_DURATION", "60"))
+
+# The canvas table shared with the caller; labels are the wire contract.
+# evaluated for the six released aspect ratios. Hardcoded so the UI renders before `diffusers` is importable.
+# Must stay identical to the generator's table: the generator forwards the *label* to `/encode`, so a canvas this
+# half does not know is rejected here and surfaces as a failure over there.
+CANVASES = {
+ # 16:9
+ "960x544 · 16:9 fast": (544, 960),
+ "1024x576 · 16:9 fast": (576, 1024),
+ "1152x640 · 16:9": (640, 1152),
+ "1280x704 · 16:9": (704, 1280),
+ "1344x768 · 16:9 full": (768, 1344),
+ # 9:16
+ "544x960 · 9:16 fast": (960, 544),
+ "640x1152 · 9:16": (1152, 640),
+ "768x1344 · 9:16 full": (1344, 768),
+ # 1:1
+ "544x544 · 1:1 fast": (544, 544),
+ "768x768 · 1:1 full": (768, 768),
+ # 4:3 / 3:4
+ "768x576 · 4:3 fast": (576, 768),
+ "1024x768 · 4:3 full": (768, 1024),
+ "576x768 · 3:4 fast": (768, 576),
+ "768x1024 · 3:4 full": (1024, 768),
+ # 21:9
+ "1152x512 · 21:9 fast": (512, 1152),
+ "1536x672 · 21:9 full": (672, 1536),
+}
+DEFAULT_CANVAS = "960x544 · 16:9 fast"
+# The demo Spaces offer durations from 2 s, and both halves validate the frame count against the same pipeline floor.
+MIN_DURATION = 2
+
+PIPE = None
+REF2VA_PIPE = None
+LOAD_ERROR: str | None = None
+LOADED_IN: float | None = None
+NUM_LAYERS: int | None = None
+
+
+def align_num_frames(num_frames: int) -> int:
+ """The `17 * n + 5` frame count MiniMax-H3's video VAE can decode, and the duration window it has to land in.
+
+ The `t2va` / `fl2va` geometry is resolved by the layout step, which lives on the denoising side of the split, so
+ the frame count the plan reports is resolved here off the same library arithmetic. The VAE's numbers come from the
+ pipeline, which falls back to the released checkpoint's `17` and `5` when no VAE is loaded, as here.
+ """
+ from diffusers.modular_pipelines.minimax_h3.modular_pipeline import align_num_frames as align
+
+ aligned = align(int(num_frames), PIPE.vae_frames_per_chunk, PIPE.vae_latents_per_chunk)
+ # The duration the request generates is the *aligned* one, so that is what the window holds for: 346 frames
+ # would otherwise pass and then be rounded up to 362, i.e. 15.083 seconds.
+ duration = aligned / PIPE.fps
+ if not PIPE.min_duration <= duration <= PIPE.max_duration:
+ raise gr.Error(
+ f"MiniMax-H3 generates between {PIPE.min_duration:g} and {PIPE.max_duration:g} seconds at {PIPE.fps} fps, "
+ f"so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can encode, must be between "
+ f"{int(PIPE.min_duration * PIPE.fps)} and {int(PIPE.max_duration * PIPE.fps)}, got {int(num_frames)} "
+ f"(rounded up to {aligned})."
+ )
+ return aligned
+
+
+def lower_duration_floor(seconds: float = MIN_DURATION) -> None:
+ """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the released checkpoint."""
+ from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+
+ MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
+
+
+def status() -> str:
+ if LOAD_ERROR:
+ return LOAD_ERROR
+ if PIPE is None:
+ return f"Loading `{MODEL_REPO}:text_encoder/` (62.14 GiB). Watch the Space logs."
+ return (
+ f"Ready · Qwen3-VL conditioner, **bfloat16, unquantized**, {NUM_LAYERS} decoder layers, read at layer 50 · "
+ f"loaded in {LOADED_IN:.0f}s · repo `{MODEL_REPO}`"
+ )
+
+
+def load_models() -> str | None:
+ """Load the conditioner. At **startup**, and onto the card.
+
+ Only `text_encoder`, `tokenizer` and `processor` are fetched: those are the pretrained components
+ `MiniMaxH3ConditionerBlocks` declares — its `image_processor` is built from config and downloads nothing — and
+ `load_components` resolves each against `modular_model_index.json`, so neither transformer partition nor either
+ VAE is ever downloaded here. That is what keeps this Space at 66.7 GB of the 150 GB quota.
+
+ The weights are moved onto the card here as well. These are plain bfloat16 tensors, so ZeroGPU's startup packing
+ (`aten.empty_like(..., pin_memory=True)`) handles them — it is a torchao `Float8Tensor` that does not implement it,
+ and there is none here.
+ """
+ global PIPE, REF2VA_PIPE, LOAD_ERROR, LOADED_IN, NUM_LAYERS
+
+ if PIPE is not None or LOAD_ERROR is not None:
+ return LOAD_ERROR
+
+ started = time.time()
+ try:
+ import torch
+
+ from h3_split_blocks import MiniMaxH3ConditionerBlocks, MiniMaxH3Ref2VAConditionerBlocks
+
+ lower_duration_floor()
+ blocks = MiniMaxH3ConditionerBlocks()
+ print(f"[cond] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
+ pipe = blocks.init_pipeline(MODEL_REPO)
+ pipe.load_components(dtype=torch.bfloat16)
+ pipe.to("cuda")
+
+ # The `ref2va` conditioner half declares the same three components, so it is handed the ones that are
+ # already resident: `update_components` fills them in, `load_components` would then find nothing left to
+ # fetch, and there is never a second 62.14 GiB copy of the Qwen3-VL on the card.
+ ref2va_pipe = MiniMaxH3Ref2VAConditionerBlocks().init_pipeline(MODEL_REPO)
+ ref2va_pipe.update_components(
+ text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer, processor=pipe.processor
+ )
+
+ PIPE, REF2VA_PIPE = pipe, ref2va_pipe
+ NUM_LAYERS = pipe.text_encoder.config.text_config.num_hidden_layers
+ LOADED_IN = time.time() - started
+ print(f"[cond] ready in {LOADED_IN:.0f}s", flush=True)
+ except Exception as error:
+ traceback.print_exc()
+ LOAD_ERROR = f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: `{type(error).__name__}: {error}`"
+ return LOAD_ERROR
+
+
+def _refine(refine, *arguments):
+ """Run one prompt refinement, or give up on it.
+
+ A rewriter is an optional accessory to an encoding service, so it is not allowed to take the service down with it:
+ anything that goes wrong is logged and answered with `None`, which every caller reads as "encode the raw prompt".
+ """
+ try:
+ import prompt_rewrite
+
+ refined = getattr(prompt_rewrite, refine)(*arguments, REWRITE_MAX_NEW_TOKENS)
+ print(f"[cond] refined prompt:\n{refined}", flush=True)
+ return refined
+ except Exception as error:
+ traceback.print_exc()
+ print(f"[cond] prompt refinement failed ({type(error).__name__}: {error}); encoding the raw prompt", flush=True)
+ return None
+
+
+@spaces.GPU(duration=REWRITE_GPU_DURATION, size=GPU_SIZE)
+def _rewrite_keyframe_prompt(prompt, image, last_image, num_frames):
+ """The rewrite, on a booking of its own — see the note on `GPU_DURATION`.
+
+ `PIPE` is read as a global rather than passed in: a `@spaces.GPU` argument crosses a process boundary by pickling.
+ """
+ return _refine("refine_keyframe_prompt", PIPE, prompt, image, last_image, num_frames)
+
+
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def _encode(prompt, image, last_image, height, width):
+ """The only thing on GPU time here: one forward through 50 Qwen3-VL decoder layers.
+
+ `num_frames` is resolved by the caller and never reaches the blocks: this half is `[resize, text_encoder]`, and
+ neither step has an opinion about the duration.
+ """
+ state = PIPE(
+ prompt=prompt,
+ image=image,
+ last_image=last_image,
+ height=int(height),
+ width=int(width),
+ )
+ return (
+ state.get("prompt_embeds").cpu().contiguous(),
+ state.get("text_token_tags").cpu().contiguous(),
+ int(state.get("height")),
+ int(state.get("width")),
+ )
+
+
+def caller(request) -> str:
+ """Which ZeroGPU identity this request's booking is charged to, for the log."""
+ headers = getattr(request, "headers", None)
+ return "the caller's own quota" if headers is not None and headers.get("x-ip-token") else "this Space's IP quota"
+
+
+def encode(prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=False, request: gr.Request = None):
+ """Encode a request and hand back one safetensors file plus the plan it resolved.
+
+ The file is the whole wire format of the split:
+
+ prompt_embeds (1, num_text_tokens, 5120) bfloat16 the layer-50 hidden state
+ text_token_tags (num_text_tokens,) int64 per-row modality tag, 1 text / 0 vision
+
+ with the resolved `height`, `width` and `num_frames` in the safetensors metadata header, so the caller can pin the
+ same canvas on the generator half rather than re-deriving it.
+
+ `rewrite_prompt` is optional and last, so a client that does not pass it gets exactly the call it always made.
+ With it, the prompt is first rewritten into MiniMax-H3's trained format and the rewrite comes back under the
+ plan's `refined_prompt` — `None` whenever no rewrite happened, whether because none was asked for or because one
+ failed.
+ """
+ if LOAD_ERROR:
+ raise gr.Error(LOAD_ERROR)
+ if PIPE is None:
+ raise gr.Error("The conditioner is still loading.")
+ if not prompt or not prompt.strip():
+ raise gr.Error("A prompt is required.")
+ print(f"[cond] /encode on {caller(request)}", flush=True)
+
+ from PIL import Image, ImageOps
+ from safetensors.torch import save_file
+
+ height, width = CANVASES[canvas]
+ # EXIF-transposed and in RGB before the blocks see them: they take a keyframe as given, and a phone photo would
+ # otherwise be conditioned on sideways.
+ image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB") if image_path else None
+ last_image = ImageOps.exif_transpose(Image.open(last_image_path)).convert("RGB") if last_image_path else None
+ num_frames = align_num_frames(num_frames)
+
+ started = time.time()
+ refined = _rewrite_keyframe_prompt(prompt, image, last_image, num_frames) if rewrite_prompt else None
+ embeds, tags, height, width = _encode(refined or prompt, image, last_image, height, width)
+ elapsed = time.time() - started
+
+ directory = os.path.join(tempfile.gettempdir(), "h3-wire")
+ os.makedirs(directory, exist_ok=True)
+ path = os.path.join(directory, f"h3-cond-{int(time.time() * 1000)}.safetensors")
+ # `prompt` is the prompt these embeddings encode, which is the rewrite when there was one; the request's own
+ # wording is then kept next to it rather than lost.
+ metadata = {
+ "height": str(height),
+ "width": str(width),
+ "num_frames": str(num_frames),
+ "prompt": refined or prompt,
+ **({"raw_prompt": prompt} if refined else {}),
+ }
+ save_file({"prompt_embeds": embeds, "text_token_tags": tags}, path, metadata)
+
+ plan = {
+ "height": height,
+ "width": width,
+ "num_frames": num_frames,
+ "num_text_tokens": int(embeds.shape[1]),
+ "hidden_size": int(embeds.shape[2]),
+ "dtype": str(embeds.dtype).removeprefix("torch."),
+ "bytes": os.path.getsize(path),
+ "seconds": round(elapsed, 2),
+ "refined_prompt": refined,
+ }
+ print(f"[cond] {plan}", flush=True)
+ return path, plan
+
+
+def build_references(references):
+ """The `(kind, path)` pairs of a request as decoded reference dataclasses, in packed order.
+
+ One class per modality, each decoding its own file through `from_file`, which brings the rates along: a video its
+ own frame rate and its soundtrack, a clip its sample rate. The blocks themselves never open a media file.
+ """
+ from diffusers.modular_pipelines.minimax_h3 import (
+ MiniMaxH3AudioReference,
+ MiniMaxH3ImageReference,
+ MiniMaxH3VideoReference,
+ )
+
+ classes = {"image": MiniMaxH3ImageReference, "video": MiniMaxH3VideoReference, "audio": MiniMaxH3AudioReference}
+ return [classes[kind].from_file(path) for kind, path in references]
+
+
+def resolve_reference_num_frames(built, num_frames):
+ """The frame count of a `ref2va` request, including the case where a single soundtrack sets it.
+
+ `0` over the wire means "leave it to the references", and the setup step requires a count, so it is resolved here
+ off the decoded waveform: `round(samples / sample_rate * 24)`, snapped up to the next `17 * n + 5` and refused
+ when the result leaves the duration window.
+ """
+ from diffusers.modular_pipelines.minimax_h3.modular_pipeline import align_num_frames as align
+
+ frames_per_chunk = REF2VA_PIPE.vae_frames_per_chunk
+ latents_per_chunk = REF2VA_PIPE.vae_latents_per_chunk
+ if num_frames:
+ return align(int(num_frames), frames_per_chunk, latents_per_chunk)
+
+ audio_bearing = [reference for reference in built if reference.has_audio]
+ if len(audio_bearing) != 1:
+ raise ValueError(
+ "`num_frames` may only be left to the references when exactly one of them carries audio, got "
+ f"{len(audio_bearing)}."
+ )
+ reference = audio_bearing[0]
+ sample_rate = reference.sample_rate or REF2VA_PIPE.audio_sampling_rate
+ duration = reference.audio.shape[-1] / sample_rate
+ if not REF2VA_PIPE.min_duration <= duration <= REF2VA_PIPE.max_duration:
+ raise ValueError(
+ f"The reference soundtrack is {duration:g} seconds long, outside the {REF2VA_PIPE.min_duration:g} to "
+ f"{REF2VA_PIPE.max_duration:g} seconds MiniMax-H3 generates."
+ )
+ resolved = align(round(duration * REF2VA_PIPE.fps), frames_per_chunk, latents_per_chunk)
+ # The duration the request generates is the one of the *aligned* frame count: a 14.99 second soundtrack rounds up
+ # to 362 frames, i.e. 15.083 seconds.
+ if resolved / REF2VA_PIPE.fps > REF2VA_PIPE.max_duration:
+ raise ValueError(
+ f"The reference soundtrack is {duration:g} seconds long, which rounds up to {resolved} frames "
+ f"(`17 * n + 5`), i.e. {resolved / REF2VA_PIPE.fps:g} seconds — past the "
+ f"{REF2VA_PIPE.max_duration:g} seconds MiniMax-H3 generates. Pass `num_frames` to generate a shorter "
+ "video from this soundtrack."
+ )
+ return resolved
+
+
+@spaces.GPU(duration=REWRITE_GPU_DURATION, size=GPU_SIZE)
+def _rewrite_reference_prompt(prompt, references, num_frames):
+ """The `ref2va` rewrite, on a booking of its own — see the note on `GPU_DURATION`.
+
+ The references are handed over as paths and decoded here, as they are for the encode: a `@spaces.GPU` argument
+ crosses a process boundary by pickling, and a 5 s 1344x768 clip is 370 MB of frames once PyAV has expanded it.
+ """
+ built = build_references(references)
+ # The rewriter is shown the references it has to write about, so it can name what each one contributes.
+ return _refine(
+ "refine_reference_prompt", REF2VA_PIPE, prompt, built, resolve_reference_num_frames(built, num_frames)
+ )
+
+
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def _encode_ref2va(prompt, references, height, width, num_frames):
+ """The only thing on GPU time here: one forward through 50 decoder layers, over the `ref2va` presentation."""
+ # A path is decoded as the reference is built, which is where a video picks up its own soundtrack.
+ built = build_references(references)
+ num_frames = resolve_reference_num_frames(built, num_frames)
+ state = REF2VA_PIPE(
+ prompt=prompt,
+ references=built,
+ height=int(height),
+ width=int(width),
+ num_frames=int(num_frames),
+ )
+ return (
+ state.get("prompt_embeds").cpu().contiguous(),
+ state.get("text_token_tags").cpu().contiguous(),
+ int(state.get("height")),
+ int(state.get("width")),
+ int(state.get("num_frames")),
+ )
+
+
+def encode_ref2va(prompt, media, kinds, canvas, num_frames, rewrite_prompt=False, request: gr.Request = None):
+ """Encode a `ref2va` request. Same wire format as `/encode`, a different presentation behind it.
+
+ `ref2va` puts a label in front of every reference, numbered per modality, plus a vision block per image and per
+ merged video frame pair — so the references reach this half as files and are decoded here. `media` and `kinds`
+ are parallel and **ordered**: the order numbers the labels and advances the shared rotary clock, so it is part
+ of the request rather than a detail of the call.
+
+ `num_frames` may be `0`, which is "leave it to the references"; accepted when exactly one
+ reference carries a soundtrack, and the duration is then that soundtrack's; the resolved count comes back in the
+ plan either way, so the denoising half pins it rather than re-deriving it.
+
+ `rewrite_prompt` is optional and last, so a client that does not pass it gets exactly the call it always made.
+ With it, the prompt is first rewritten into MiniMax-H3's full-reference format — the rewriter is shown the
+ references, so it can name what each one contributes — and the rewrite comes back under the plan's
+ `refined_prompt`, `None` whenever no rewrite happened.
+ """
+ if LOAD_ERROR:
+ raise gr.Error(LOAD_ERROR)
+ if REF2VA_PIPE is None:
+ raise gr.Error("The conditioner is still loading.")
+ if not prompt or not prompt.strip():
+ raise gr.Error("A prompt is required.")
+ print(f"[cond] /encode_ref2va on {caller(request)}", flush=True)
+
+ from safetensors.torch import save_file
+
+ media = [media] if isinstance(media, str) else list(media or [])
+ kinds = [kind.strip() for kind in (kinds or "").split(",") if kind.strip()]
+ if len(media) != len(kinds):
+ raise gr.Error(f"{len(media)} files against {len(kinds)} kinds; `media` and `kinds` are parallel and ordered.")
+ if unknown := sorted(set(kinds) - {"image", "video", "audio"}):
+ raise gr.Error(f"A reference is an `image`, a `video` or an `audio`, got {unknown}.")
+ references = list(zip(kinds, media))
+
+ height, width = CANVASES[canvas]
+ requested_num_frames = int(num_frames) or None
+ started = time.time()
+ refined = _rewrite_reference_prompt(prompt, references, requested_num_frames) if rewrite_prompt else None
+ embeds, tags, height, width, num_frames = _encode_ref2va(
+ refined or prompt, references, height, width, requested_num_frames
+ )
+ elapsed = time.time() - started
+
+ directory = os.path.join(tempfile.gettempdir(), "h3-wire")
+ os.makedirs(directory, exist_ok=True)
+ path = os.path.join(directory, f"h3-cond-ref2va-{int(time.time() * 1000)}.safetensors")
+ metadata = {
+ "height": str(height),
+ "width": str(width),
+ "num_frames": str(num_frames),
+ "prompt": refined or prompt,
+ **({"raw_prompt": prompt} if refined else {}),
+ }
+ save_file({"prompt_embeds": embeds, "text_token_tags": tags}, path, metadata)
+
+ plan = {
+ "height": height,
+ "width": width,
+ "num_frames": num_frames,
+ "references": kinds,
+ "num_text_tokens": int(embeds.shape[1]),
+ "hidden_size": int(embeds.shape[2]),
+ "dtype": str(embeds.dtype).removeprefix("torch."),
+ "bytes": os.path.getsize(path),
+ "seconds": round(elapsed, 2),
+ "refined_prompt": refined,
+ }
+ print(f"[cond] ref2va {plan}", flush=True)
+ return path, plan
+
+
+load_models()
+
+INTRO = """# Qwen3-VL conditioner
+
+An embedding service: sends back the conditioning tensors (`prompt_embeds`, token tags and the resolved
+plan) that a downstream video generation Space consumes. Call it over the gradio API; the UI below is
+for inspection only.
+
+`rewrite_prompt` is optional and off by default. Turned on, the conditioner first *writes* the structured
+prompt format the video model was trained on — task instruction, shot-by-shot timeline, soundscape — and
+encodes that instead; it comes back under the plan's `refined_prompt`.
+"""
+
+REWRITE_LABEL = "Rewrite the prompt into the trained format first (adds ~15-50s)"
+
+
+with gr.Blocks(title="Qwen3-VL conditioner") as demo:
+ gr.Markdown(INTRO)
+ banner = gr.Markdown(status())
+
+ with gr.Tab("Keyframes (t2va / fl2va)"):
+ with gr.Row():
+ with gr.Column():
+ prompt = gr.Textbox(
+ label="Prompt",
+ lines=3,
+ value="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
+ )
+ with gr.Row():
+ image = gr.Image(label="First keyframe (optional)", type="filepath")
+ last_image = gr.Image(label="Last keyframe (optional)", type="filepath")
+ canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
+ num_frames = gr.Number(label="num_frames (17n + 5)", value=124, precision=0)
+ rewrite = gr.Checkbox(label=REWRITE_LABEL, value=False)
+ run = gr.Button("Encode", variant="primary")
+ with gr.Column():
+ wire = gr.File(label="prompt_embeds + text_token_tags (safetensors)")
+ plan = gr.JSON(label="Resolved plan")
+
+ with gr.Tab("References (ref2va)"):
+ with gr.Row():
+ with gr.Column():
+ ref_prompt = gr.Textbox(
+ label="Prompt",
+ lines=3,
+ value="The character walks through a neon-lit street in the rain, humming to themselves",
+ )
+ ref_media = gr.File(label="References, in the order the model reads them", file_count="multiple")
+ ref_kinds = gr.Textbox(label="Kinds, parallel to the files", value="image", placeholder="video,image")
+ ref_canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
+ ref_num_frames = gr.Number(
+ label="num_frames (17n + 5, or 0 to take it from a single soundtrack)", value=0, precision=0
+ )
+ ref_rewrite = gr.Checkbox(label=REWRITE_LABEL, value=False)
+ ref_run = gr.Button("Encode", variant="primary")
+ with gr.Column():
+ ref_wire = gr.File(label="prompt_embeds + text_token_tags (safetensors)")
+ ref_plan = gr.JSON(label="Resolved plan")
+
+ # `rewrite_prompt` is appended *after* every input that was already here, and every existing input keeps its
+ # position, so a positional client that predates it keeps working untouched — it simply takes the default.
+ run.click(encode, [prompt, image, last_image, canvas, num_frames, rewrite], [wire, plan], api_name="encode")
+ ref_run.click(
+ encode_ref2va,
+ [ref_prompt, ref_media, ref_kinds, ref_canvas, ref_num_frames, ref_rewrite],
+ [ref_wire, ref_plan],
+ api_name="encode_ref2va",
+ )
+ demo.load(status, None, banner, api_name="status")
+
+
+if __name__ == "__main__":
+ demo.queue(max_size=8).launch(show_error=True)
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/h3_split_blocks.py b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/h3_split_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..36e52def9248247c788fca65d9f44681eb1ed161
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/h3_split_blocks.py
@@ -0,0 +1,147 @@
+"""The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.
+
+MiniMax-H3 is 195.9 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage, so `MiniMaxH3Blocks` is cut
+at its `text_encoder` step: the 62.14 GiB Qwen3-VL runs in the conditioner Space, everything else in a generator
+Space, and `prompt_embeds` + `text_token_tags` is the whole wire format between them.
+
+`resize` / `setup` run on **both** sides: they own no pretrained component, and each half needs the canvas and the
+prepared keyframes or normalized references. Both conditioner halves also return the resolved `height` / `width` /
+`num_frames`, which the generating half pins rather than re-deriving.
+
+Two things the blocks leave to the caller: a keyframe reaches them EXIF-transposed and in RGB, and the `t2va` / `fl2va`
+frame count is aligned to `17 * n + 5` before the call, since that arithmetic lives on the denoising side of the cut.
+"""
+
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.encoders import (
+ MiniMaxH3Ref2VAReferenceEncoderStep,
+ MiniMaxH3Ref2VATextEncoderStep,
+ MiniMaxH3TextEncoderStep,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
+ MiniMaxH3AutoKeyframeVaeEncoderStep,
+ MiniMaxH3AutoResizeStep,
+ MiniMaxH3CoreDenoiseStep,
+ MiniMaxH3DecodeStep,
+ MiniMaxH3Ref2VACoreDenoiseStep,
+ _generation_outputs,
+)
+from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
+
+
+def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
+ """The wire format of the split. `num_frames` is declared by the `ref2va` half alone, whose setup resolves one."""
+ return [
+ OutputParam.template("prompt_embeds"),
+ OutputParam("text_token_tags", description="The per-row modality tag of every row of `prompt_embeds`."),
+ OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+ OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+ *(
+ [OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5.")]
+ if num_frames
+ else []
+ ),
+ ]
+
+
+class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
+ """The conditioner half of a split MiniMax-H3: the keyframes on the canvas plus the Qwen3-VL read at layer 50."""
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3AutoResizeStep, MiniMaxH3TextEncoderStep]
+ block_names = ["resize", "text_encoder"]
+
+ @property
+ def description(self):
+ return (
+ "The conditioner half of a split MiniMax-H3 deployment: puts the keyframes onto the target canvas and "
+ "encodes MiniMax-H3's presentation of the request into the `prompt_embeds` / `text_token_tags` pair the "
+ "denoising half consumes. The frame count is the caller's to align."
+ )
+
+ @property
+ def outputs(self):
+ return _wire_outputs(num_frames=False)
+
+
+class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
+ """The denoising half of a split MiniMax-H3: `MiniMaxH3Blocks` with its `text_encoder` step removed."""
+
+ model_name = "minimax-h3"
+ block_classes = [
+ MiniMaxH3AutoResizeStep,
+ MiniMaxH3AutoKeyframeVaeEncoderStep,
+ MiniMaxH3CoreDenoiseStep,
+ MiniMaxH3AfterDenoiseStep,
+ MiniMaxH3DecodeStep,
+ ]
+ block_names = ["resize", "vae_encoder", "denoise", "after_denoise", "decode"]
+
+ @property
+ def description(self):
+ return (
+ "The denoising half of a split MiniMax-H3 deployment: the `t2va` / `fl2va` branch of `MiniMaxH3Blocks` "
+ "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
+ "62.14 GiB Qwen3-VL conditioner is never loaded here."
+ )
+
+ @property
+ def outputs(self):
+ return _generation_outputs()
+
+
+class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
+ """The conditioner half of a split `ref2va`: the resolved plan plus the Qwen3-VL read at its 50th layer.
+
+ Component for component this is `MiniMaxH3ConditionerBlocks`, so one conditioner Space serves both partitions.
+ What differs is the presentation: `ref2va` prepends a label per reference and a vision block per image and per
+ merged video frame pair, so the references themselves have to reach this half.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3Ref2VATextEncoderStep]
+ block_names = ["setup", "text_encoder"]
+
+ @property
+ def description(self):
+ return (
+ "The conditioner half of a split MiniMax-H3 `ref2va` deployment: resolves the request plan (canvas, frame "
+ "count, references normalized onto MiniMax-H3's own rates and resolutions) and encodes MiniMax-H3's "
+ "presentation of it into the `prompt_embeds` / `text_token_tags` pair the denoising half consumes."
+ )
+
+ @property
+ def outputs(self):
+ return _wire_outputs()
+
+
+class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
+ """The denoising half of a split `ref2va`: the `ref2va` branch with its `text_encoder` step removed.
+
+ `reference_encoder` stays here, next to the two autoencoders it runs: its output shapes are where every reference
+ block's geometry in the packed layout comes from.
+ """
+
+ model_name = "minimax-h3"
+ block_classes = [
+ MiniMaxH3Ref2VASetupStep,
+ MiniMaxH3Ref2VAReferenceEncoderStep,
+ MiniMaxH3Ref2VACoreDenoiseStep,
+ MiniMaxH3AfterDenoiseStep,
+ MiniMaxH3DecodeStep,
+ ]
+ block_names = ["setup", "reference_encoder", "denoise", "after_denoise", "decode"]
+
+ @property
+ def description(self):
+ return (
+ "The denoising half of a split MiniMax-H3 `ref2va` deployment: the `ref2va` branch of `MiniMaxH3Blocks` "
+ "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
+ "62.14 GiB Qwen3-VL conditioner is never loaded here. The transformer is the `transformer_ref` partition."
+ )
+
+ @property
+ def outputs(self):
+ return _generation_outputs()
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/packages.txt b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/packages.txt
new file mode 100644
index 0000000000000000000000000000000000000000..2a2134b01ca333d811948a16c5a6c27166e9d297
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/packages.txt
@@ -0,0 +1 @@
+ffmpeg
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/prompt_rewrite.py b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/prompt_rewrite.py
new file mode 100644
index 0000000000000000000000000000000000000000..0077b721c30daad814814065db3ffa86fea725d8
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/prompt_rewrite.py
@@ -0,0 +1,878 @@
+"""Context-IR style prompt refinement, on the conditioner's own Qwen3-VL.
+
+MiniMax-H3 is trained on a *rewritten* prompt: a structured Context Intermediate Representation, not free text.
+MiniMax's own rewriter (`H3-Context-IR`) is a hosted service that is not part of the open release, but the two
+official prompting guides specify the target format exactly, and the conditioner this Space already holds is the
+full 64-layer `Qwen3VLForConditionalGeneration` — language-model head, chat template and all. So the rewrite runs
+here, on the same weights that are about to encode its output, and costs one extra prefill plus a short decode.
+
+The split of responsibilities is deliberate:
+
+ * **The instruction line is code.** Every keyframe task opens with a fixed alignment sentence whose only variables
+ are the task, the effective duration to two decimals and the index of the final shot. A language model asked for
+ a template with numbers in it will paraphrase the template and round the numbers, so [`base_instruction_line`]
+ formats it from the request instead. `t2va` has no instruction line at all.
+ * **The descriptive core is the model's.** `integrated_multimodal_description` (or, in full-reference mode, the six
+ sections from `subject_definitions` to `non_diegetic_music`) is what actually needs a vision-language model:
+ it has to look at the keyframes or the references and write the timeline they imply.
+
+Nothing here is on the encoding path. The refined prompt is fed to the unchanged encoder as `prompt`, and any failure
+falls back to the raw prompt.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import time
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+# The conditioner reads a keyframe at the full canvas and a reference image at a 2048 pixel short edge. The *rewriter*
+# does not need that: it needs to see what is in the frame. A 1024 pixel long edge is ~700 vision tokens per image
+# through the 16 pixel patches and the 2x2 merge, which keeps a nine-image request inside a sane prefill.
+REWRITE_IMAGE_LONG_EDGE = 1024
+# A reference video reaches the conditioner at 2 fps, so that is what the rewriter is shown as well — capped, so that
+# one long clip cannot own the prefill. 32 frames is 16 seconds, past the 15 MiniMax-H3 generates.
+REWRITE_VIDEO_FPS = 2.0
+REWRITE_VIDEO_MAX_FRAMES = 32
+
+MINIMAX_H3_FPS = 24
+# The video VAE's own chunking, which the frame count has to land on: `17 * n + 5`. Arguments to the library's
+# `align_num_frames` since the blocks were refactored, where they used to be module constants of a `packing` module
+# that no longer exists.
+FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 17, 5
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# The deterministic half: task resolution, duration and the instruction line.
+# ---------------------------------------------------------------------------------------------------------------
+
+
+def base_task(image, last_image) -> str:
+ """Which of the four base tasks a `/encode` request is, from the keyframes it carries."""
+ if image is not None and last_image is not None:
+ return "fl2va"
+ if image is not None:
+ return "i2va"
+ if last_image is not None:
+ return "l2va"
+ return "t2va"
+
+
+def aligned_num_frames(num_frames: int) -> int:
+ """The `17 * n + 5` frame count the video VAE decodes, which is what the request actually generates."""
+ from diffusers.modular_pipelines.minimax_h3.modular_pipeline import align_num_frames
+
+ return int(align_num_frames(int(num_frames), FRAMES_PER_CHUNK, LATENTS_PER_CHUNK))
+
+
+def duration_seconds(num_frames: int) -> float:
+ """The effective duration of the request, at MiniMax-H3's fixed 24 fps."""
+ return aligned_num_frames(num_frames) / MINIMAX_H3_FPS
+
+
+def final_shot_index(core: str) -> int:
+ """The index of the last `[Shot N]` the rewrite actually wrote, which the instruction line points at."""
+ indices = [int(match) for match in re.findall(r"\[Shot (\d+)\]", core)]
+ return max(indices) if indices else 1
+
+
+def base_instruction_line(task: str, duration: float, final_shot: int) -> str | None:
+ """The alignment instruction a keyframe task opens with, verbatim from the guide with the numbers filled in.
+
+ The bracket styles differ between the three templates — `fl2va` writes `Picture 2 (from Shot 1)` while `l2va`
+ writes ` (from [Shot 1])`. That is how the guide specifies them, so that is what is reproduced.
+
+ Args:
+ task (`str`): `"t2va"`, `"i2va"`, `"fl2va"` or `"l2va"`.
+ duration (`float`): The effective duration in seconds, rendered to exactly two decimals.
+ final_shot (`int`): The index of the final shot of the rewrite, `N` in the templates.
+
+ Returns:
+ `str` or `None`: the instruction line, or `None` for `t2va`, which has none.
+ """
+ if task == "t2va":
+ return None
+ if task == "i2va":
+ return (
+ "For the target video, at 0.00 seconds into the target video, (from [Shot 1]) is fully "
+ "referenced."
+ )
+ if task == "fl2va":
+ return (
+ "How the reference pictures align with the target video — Picture 1 (from Shot 1) aligns with the "
+ f"0.00-second mark of the target video; Picture 2 (from Shot {final_shot}) aligns with the "
+ f"{duration:.2f}-second mark of the target video."
+ )
+ if task == "l2va":
+ return (
+ "How the reference pictures align with the target video — (from [Shot "
+ f"{final_shot}]) aligns with the {duration:.2f}-second mark of the target video."
+ )
+ raise ValueError(f"Unknown base task {task!r}.")
+
+
+def reference_labels(references) -> list[tuple[str, str]]:
+ """The label MiniMax-H3's `ref2va` presentation gives every reference, in packed order.
+
+ Mirrors `tokenize_ref2va_presentation`: labels are numbered *per modality*, and a video that carries a soundtrack
+ is labelled `` before ``, so a rewrite that cites the wrong number cites the wrong asset.
+
+ Args:
+ references (`list[MiniMaxH3Reference]`): The references of the request, in packed order.
+
+ Returns:
+ `list[tuple[str, str]]`: `(kind, label)` per reference, where a sound-bearing video's label names both.
+ """
+ counts = {"image": 0, "video": 0, "audio": 0}
+ labels = []
+ for reference in references:
+ audio_label = None
+ if reference.has_audio:
+ counts["audio"] += 1
+ audio_label = f""
+ if reference.kind == "image":
+ counts["image"] += 1
+ labels.append(("image", f""))
+ elif reference.kind == "video":
+ counts["video"] += 1
+ label = f""
+ labels.append(("video", f"{audio_label} + {label}" if audio_label else label))
+ else:
+ labels.append(("audio", audio_label))
+ return labels
+
+
+def reference_duration_seconds(references, num_frames) -> float | None:
+ """The duration a `ref2va` request generates, including the case where the soundtrack sets it.
+
+ `num_frames` may be left open, and the `ref2va` setup step then takes the duration off the single audio-bearing
+ reference. The rewrite needs that number to place its cut timestamps, so the same resolution is repeated here —
+ on the already-decoded waveform, so it costs nothing — and simply gives up (`None`) if the request is one setup
+ would reject anyway.
+ """
+ if num_frames:
+ return duration_seconds(num_frames)
+ audio_bearing = [reference for reference in references if reference.has_audio]
+ if len(audio_bearing) != 1:
+ return None
+ reference = audio_bearing[0]
+ # A decoded reference always carries the rate its container reported; one built from an in-memory waveform without
+ # a rate is resolved against the audio VAE's own, which is the pipeline's to know, so the hint is left out.
+ if not reference.sample_rate:
+ return None
+ seconds = reference.audio.shape[-1] / reference.sample_rate
+ return duration_seconds(round(seconds * MINIMAX_H3_FPS))
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# The system prompts: the two official guides, condensed to what a rewrite has to obey.
+# ---------------------------------------------------------------------------------------------------------------
+
+_SHARED_RULES = """\
+## Shots and cuts
+
+Begin the body with `[Shot 1]`, which never carries a timestamp. Later shots open with a strictly increasing cut time
+that falls inside the video duration, written `[Shot 2] At 00:03.500, the camera cuts to ...` (`MM:SS.mmm`). Use `the
+camera cuts to`, `the shot cuts to`, `the shot transitions to`, `the shot changes to` or `the shot switches to` for an
+ordinary cut; cross-dissolve, fade or wipe only when the user asked for one. A cut must introduce new information
+about the subject, space, state, viewpoint or time — if only the distance or the angle changes, move the camera
+instead. Never place a cut at or past the end of the video.
+
+A cut is expressed *only* by opening the next `[Shot N]` with its timestamp. Never describe a transition, a cut or a
+jump to another scene in the middle of a shot: if the picture changes, the shot has ended and the next `[Shot N]`
+begins. Everything inside one `[Shot N]` is one continuous take.
+
+## Camera motion
+
+Write camera motion as natural English inside the shot, never as labels stacked at the end of a sentence. Motion type
+is one of `Zoom In / Zoom Out`, `Push In / Pull Out`, `Pan Left / Pan Right`, `Truck Left / Truck Right`, `Tilt Up /
+Tilt Down`, `Pedestal Up / Pedestal Down`, `Arc Shot`, `Tracking Shot`, `Static Shot`, `Shake Slightly / Shake
+Strongly`, `POV`, `Roll Clockwise / Roll Counterclockwise`. Add amplitude (`with small amplitude`, `with large
+amplitude`) and speed (`at slow speed`, `at fast speed`) only when they mean something. Medium amplitude and normal
+speed are the default and are left unsaid — never write `with medium amplitude`, `at normal speed` or any other
+amplitude or speed phrase than the four above.
+
+ The camera pushes in with small amplitude at slow speed toward the folded letter in her hands.
+ The camera holds a static shot as the runner exits the frame.
+
+## Speakers, dialogue and singing
+
+Anyone who speaks, sings or produces an off-screen human voice gets a stable ID: `(S1)`, `(S2)`, and `(S1,S2)` when
+already-numbered speakers vocalize together. An ID belongs to the same person across every shot; characters who never
+vocalize get none. At a speaker's first appearance, establish a stable identity from what is visible and audible —
+character type, age, gender, on-screen or not, pitch, timbre, speaking rate, accent.
+
+The identifying phrase, the ID, the action and the delivery go *outside* ``. Inside `` there is only a language
+tag and the spoken content itself, every original word and punctuation mark preserved verbatim, never translated and
+never rewritten. The wrapper is literally the two characters-plus-one letter `` and its closing ` ` — never any
+other tag name, and never a tag out of your own chat vocabulary:
+
+ The young woman with a quiet, breathy voice (S1) says: [English] I get off at the next station.
+ The two children (S1,S2) shout together, [English] Wait for us!
+
+For voiceover use the exact phrase `says in an off-screen voiceover`, and immediately after that `` block state
+that the on-screen character's lips remain closed. When one line crosses a cut, write `` at the
+connecting point in both parts and say the audio continues across the cut (`continues seamlessly across the cut`,
+`continues uninterrupted into the next shot`, `carries over from the previous shot`, `remains audible across the
+transition`). Use `` when speech is truncated by the end of the video.
+
+Never invent dialogue the user did not ask for. If the user gave no spoken lines, write no `` blocks at all.
+
+## On-screen text
+
+Any banner, sign, label, subtitle or neon text actually visible on screen goes in English double quotation marks,
+verbatim and untranslated: `A red neon sign reading "营业中" glows above the doorway.`
+
+## overall_soundscape
+
+One continuous paragraph, 1-4 English sentences, summarizing ambience, physical action sounds and non-verbal human
+sounds across the whole video: wind, rain, traffic, footsteps, fabric, impacts, breathing, laughter, panting.
+Dialogue, singing and diegetic music belong in the description body and are not repeated here. `N/A` only when the
+user explicitly asked for complete silence.
+
+## non_diegetic_music
+
+1-3 English sentences on the background music only the audience hears. Instrumentation, tempo, rhythm and dynamic
+change — no abstract mood words and no explanation of what the score does emotionally. Music the characters can hear
+(singing, instruments, a radio, a television, a phone) is diegetic and belongs in the body instead. `N/A` when there
+is no non-diegetic music.\
+"""
+
+BASE_SYSTEM_PROMPT = f"""\
+You are H3-Context-IR, the prompt rewriter of the MiniMax-H3 audiovisual generation model. You turn a user's raw
+request into the structured representation H3 was trained on. Everything you write is consumed directly by the
+generator, so you never address the user and never explain yourself.
+
+# Output contract
+
+Emit exactly these three fields, in this order, each separated from the next by one blank line, and nothing else:
+
+ integrated_multimodal_description: [Shot 1] ...
+
+ overall_soundscape: ...
+
+ non_diegetic_music: ...
+
+Hard rules:
+
+- Plain text only. No markdown, no code fences, no headings, no bullet lists, no preamble, no closing remark.
+- Start your very first token with `integrated_multimodal_description:`.
+- Never write the picture-alignment instruction line (`For the target video, ...` / `How the reference pictures align
+ ...`). The system prepends it. Writing one yourself is an error.
+- Write in English. Only dialogue and lyrics inside ``, and text visibly present in the scene, keep their original
+ language.
+- Stay faithful to the user's request: keep every subject, action, line of dialogue and sound they specified. Add
+ scene, character, action and sound detail only where it is consistent with their intent.
+
+# integrated_multimodal_description
+
+This is the body of the rewrite, developed along the timeline. Every detail must correspond to something that can be
+seen or heard: visual style, opening composition, subject appearance and position, scene and key props, actions and
+reactions, cuts, spoken language, and sound synchronized to what happens.
+
+Open `[Shot 1]` with the overall style and the initial composition — `Cinematic`, `live-action`, `2D-animated`,
+`3D CG`, `claymation`, `watercolor`, `vintage film` and the like. With keyframes, read the style off the reference
+image; without them, take it from the user's text.
+
+ [Shot 1] Live-action, cinematic, a medium-wide shot frames...
+
+{_SHARED_RULES}
+
+# Working from keyframes
+
+**i2va (first frame given).** `` *is* the frame at 0.00 seconds and belongs to `[Shot 1]`. Establish the
+style, subjects, composition and scene anchors of the image first, then the action that follows. Character identity,
+clothing, colors, key objects and spatial relationships stay consistent. Shape: first-frame anchor, action onset,
+continuous development, result or reaction.
+
+**fl2va (first and last frame given).** Picture 1 opens and Picture 2 ends. Do not describe two static images — supply
+the motion path between them: how the subject moves, how the pose changes, how objects are handled, how the
+composition, scene and lighting evolve. Prefer a single shot so the model can interpolate continuously, unless the
+user explicitly asked for more, and reach Picture 2's state at the very end of the final shot. Shape: first-frame
+state, observable intermediate changes, progressively narrowing differences, last-frame state.
+
+**l2va (last frame given).** `` is the *final* frame and belongs to the last shot, not to Shot 1. Infer a
+plausible earlier state from the user's intent and the image, then describe how characters, objects, camera and scene
+gradually approach it. Shape: plausible preceding state, explicit action and transition path, gradual convergence,
+last-frame landing.
+
+# Example (text-only request, 8.00 seconds)
+
+integrated_multimodal_description: [Shot 1] Live-action, cinematic, a medium-wide shot frames a baker opening the
+shutters of a small street bakery before sunrise. The camera pushes in with small amplitude at slow speed as the
+middle-aged baker with a calm, slightly raspy voice (S1) places a fresh loaf on the wooden counter and says:
+[English] First batch of the morning. [Shot 2] At 00:05.000, the camera cuts to a close-up of steam rising
+from the sliced bread while the baker's final words carry over from the previous shot.
+
+overall_soundscape: Wooden shutters scrape open over a quiet street as trays clink softly inside the bakery. The
+doorbell rings once, followed by light footsteps and the crisp sound of bread being sliced.
+
+non_diegetic_music: A soft acoustic-guitar pattern at a moderate tempo, joined by sparse upright-bass notes and a
+gentle fade at the end.
+
+# Example (first frame given, 8.00 seconds)
+
+integrated_multimodal_description: [Shot 1] Live-action, cinematic, the young woman shown in remains
+beside the rain-covered train window, preserving her appearance, clothing, seat position, and the carriage layout.
+The camera trucks right with small amplitude at slow speed as she lifts her gaze from the folded letter toward the
+passing city lights. Her reflection moves across the glass while the quiet, breathy young woman (S1) says:
+[English] I get off at the next station. She folds the letter along its existing crease.
+
+overall_soundscape: The train wheels produce a steady metallic rhythm beneath a low ventilation hum. Rain ticks
+against the window while paper rustles softly in her hands.
+
+non_diegetic_music: Sustained cello notes at a slow tempo with widely spaced piano tones, gradually decreasing in
+volume.\
+"""
+
+REF_SYSTEM_PROMPT = f"""\
+You are H3-Context-IR, the prompt rewriter of the MiniMax-H3 audiovisual generation model, working in
+**full-reference mode**: the request comes with reference images, videos and audio, and your rewrite has to state
+what each of them contributes to the target video. Everything you write is consumed directly by the generator, so you
+never address the user and never explain yourself.
+
+# Output contract
+
+Emit exactly these six sections, in this order, each section name on its own line followed by its content, and one
+blank line between sections. Nothing else.
+
+ subject_definitions:
+
+
+ summary:
+ [task type] one short paragraph
+
+ retention_analysis:
+
+
+ detailed_description:
+
+
+ overall_soundscape:
+ ...
+
+ non_diegetic_music:
+ ...
+
+Hard rules:
+
+- Plain text only. No markdown, no code fences, no bullet lists, no preamble, no closing remark.
+- Start your very first token with `subject_definitions:`.
+- Write in English. Only dialogue and lyrics inside ``, and text visibly present in the scene, keep their original
+ language.
+- Use only the reference labels the request lists, exactly as spelled there. Never invent a label, never renumber one,
+ and never introduce a new label after `subject_definitions`.
+- A label keeps one meaning across all six sections.
+
+# subject_definitions
+
+One line per piece of referenced content that is tracked separately later. Four label types:
+
+- `` — visible content reused or modified in the target video: people, animals, objects, scenes,
+ backgrounds, environments, clothing, props, interfaces, effects, styles, actions, expressions, poses. It is a
+ *content unit*, not a file. One subject may come from several assets, and one asset may provide several subjects.
+- `` — a reference image used as a concrete frame anchor (first frame, keyframe, last frame, edited
+ keyframe) or as a storyboard / composition reference.
+- `` — a reference video in a whole-video relationship: the source being edited, the clip being continued,
+ or a reference for camera movement, cuts, rhythm or temporal structure.
+- `` — an audio asset, or a reference video's synchronized track, that is copied or referenced.
+
+Say what the label denotes, what role it plays and which features to follow; name the source asset when provenance
+would otherwise be ambiguous.
+
+ is the young woman in , with long dark hair, a blue cardigan, and a thin silver necklace.
+ is the woman whose appearance comes from and whose walking motion comes from .
+ is the first frame of [Shot 1], showing a woman seated beside a café window.
+ is the source video for the target video edit.
+ is the voice-timbre reference for (S1).
+
+A label is only ever for content that is actually *in* one of the reference assets listed for this request. Anything
+that exists only in the user's text — a character, a place, a weather condition nobody sent you a picture of — is
+written into `detailed_description` as ordinary description and gets no `` of its own. Never define a
+subject you cannot see in a reference.
+
+An image used *only* to define a character, scene, costume or style gets no standalone `` line — cite it
+inside the `` definition it supports. A person, object, scene, action or effect taken out of a reference
+video is a ``, not a ``. A reference video does not produce an `` merely because the file
+has sound. When an `` corresponds to a target speaker, reuse that speaker's global ID: ` (Sx)`,
+or a stable voice description followed by `(Sx)`.
+
+# summary
+
+One short English paragraph, opening with a square-bracketed task-type prefix. Task types:
+
+| `keyframe completion` | an image is a concrete frame anchor of the target video |
+| `reference generation` | an image, video or audio guides generation of a character, scene, style, action, camera
+ movement or storyboard, without being a frame anchor or the video being edited or continued |
+| `video editing` | an existing source video is directly modified |
+| `video continuation` | new content continues, extends, resumes or transitions from a source video |
+| `audio reuse` | the same audio signal is reused in full or in part |
+| `audio reference` | only the music style, timbre, dialogue or lyric content, sound texture, beat or continuity of
+ an audio signal is referenced, not the signal itself |
+
+Combine several with ` + `, never repeating one: `[video continuation + keyframe completion]`, `[video editing +
+audio reuse]`. The mere presence of a video or an audio file does not create a task type — a video that only lends
+camera movement, cuts or rhythm is `reference generation`. For a video-editing task, the paragraph begins right after
+the prefix with `The target video is an edited version of .`
+
+# retention_analysis
+
+One line per reference label, in the meaning `subject_definitions` gave it. Visible content
+(``, ``, ``) uses `fully_preserved`, `partially_preserved`, `attribute_transfer` or
+`weak_reference`. Audio (``) uses `fully_copy`, `partially_copy`, `reference` or `weak_reference`. These
+markers are fixed values — write them exactly.
+
+ (appears in [Shot 1], [Shot 3]): fully_preserved - ...
+ ([Shot 1] first frame): fully_preserved - ...
+ (cut and pacing structure): weak_reference - ...
+ : fully_copy - is reused 1:1 as the target video's complete final audio track.
+
+Stay inside the role the label was already given. New actions, backgrounds or plot events in the target video are not
+losses of fidelity. Never write `(Sx)` in this section.
+
+# detailed_description
+
+The main body, shot by shot in playback order, with reference labels inserted where they apply. Establish the style
+in one or two English sentences *before* `[Shot 1]` — unlike text-only mode, where the style opens Shot 1.
+
+Make this as detailed and explicit as you can. For every shot: the current composition, subject appearance and
+position, environment and lighting, actions and state changes, camera movement, the sound at that moment, and the
+point where referenced content actually appears or takes effect. Never reduce it to a plot summary or a list of
+reference relationships. Aim for 350-500 English words for a generation task; dialogue-heavy content prioritizes
+fitting the whole spoken timeline over any word count. A single shot is not a licence to write less.
+
+At an important ``'s first clear appearance, describe its referenced characteristics, where it sits in the
+frame and what it is doing — then keep using the label without redefining it. Frame anchors read naturally: `the shot
+begins from `, `the shot's keyframe corresponds to `, `the shot ends on `. Cite
+`` where its source state, structure or continuation relationship applies, and `` in the shot or
+phase where its audio relationship is active.
+
+When a referenced subject speaks, keep both labels: ` (S1) turns toward the woman and says, [English]
+Last summer, I went to my grandfather's house. `. `` is the referenced subject, `(Sx)` the actual
+speaker; off-screen keeps the same form and is marked `off-screen`. A speaker with no defined subject gets a stable
+voice description plus `(Sx)`. Verbal content that exists only inside a directly reused soundtrack uses ``
+as its source and gets no `(Sx)`. Assign `(Sx)` once, in the order vocal events happen in the target video.
+
+When dialogue, narration or lyrics are reused from reference audio, or the user asked for them to be reperformed,
+keep the exact source words and original language inside ``, write `[unclear]` for spans you cannot make out, and
+standardize punctuation to `,` `.` `?` `!` — no tildes, emoji, bullets or decorative repeats. When only timbre,
+rhythm, emotion or delivery is referenced, do not carry the original words over.
+
+{_SHARED_RULES}
+
+State an audio copy-or-reference relationship in the section that matches the audible layer: ambience and effects in
+`overall_soundscape`, audience-only score in `non_diegetic_music`. Complete dialogue and lyrics appear only inside
+`` in `detailed_description`.
+
+# Example
+
+subject_definitions:
+ is the coffee-shop environment in , featuring an exposed brick wall, an orange tufted sofa with
+patterned pillows, a neon sign, and a wooden coffee table.
+ is the fluffy white Samoyed in , with thick white fur, pointed ears, a dark nose, and a curved
+tail.
+ is the young blonde woman in , with long blonde hair and a light-pink button-down shirt with
+rolled-up sleeves.
+ is the voice-timbre reference for (S1), containing a spoken English vocal layer.
+
+summary:
+[reference generation + audio reference] The target video shows eating a cookie in . A young man
+enters with , which lunges toward the cookie. The two-shot exchange uses as the voice-timbre
+reference for and ends with a canned audience laugh.
+
+retention_analysis:
+ (appears in [Shot 1], [Shot 2]): fully_preserved - the exposed brick wall, orange tufted sofa, patterned
+pillows, neon sign, and wooden coffee table are retained.
+ (appears in [Shot 1], [Shot 2]): fully_preserved - the Samoyed's thick white fur, pointed ears, dark nose,
+and curved tail are retained.
+ (appears in [Shot 1], [Shot 2]): fully_preserved - the blonde woman's identity, long hair, and light-pink
+shirt are retained.
+: reference - its vocal timbre guides the dialogue delivery of without copying the original
+signal.
+
+detailed_description:
+The target video uses a realistic multi-camera sitcom style with warm indoor lighting.
+[Shot 1] A medium shot establishes , the coffee shop with its exposed brick wall, orange tufted sofa,
+patterned pillows, neon sign, and wooden coffee table. (S1), the young woman with long blonde hair and a
+light-pink button-down shirt with rolled-up sleeves, sits on the sofa holding a chocolate-chip cookie. From the left, a
+young man in a dark-grey hoodie enters holding the leash of , the thick-furred white Samoyed with pointed
+ears, a dark nose, and a curved tail. The dog lunges toward the cookie and pulls the leash taut. (S1) jerks
+her hand back and, using the clear youthful voice timbre referenced from , exclaims with light annoyance,
+[English] Hey! Watch your dog! She closes her lips and guards the cookie while the man pulls the dog back.
+[Shot 2] At 00:03.000, the shot cuts to a close-up of (S1), the blonde woman in the light-pink shirt from
+Shot 1. Her annoyance softens as she looks toward the Samoyed. The camera pushes in with small amplitude at slow speed
+while a classic canned audience laugh begins and continues through the final frame.
+
+overall_soundscape:
+Soft indoor coffee-shop room tone continues throughout the scene, with the leash snapping taut and the dog's claws
+scrabbling on the wooden floor.
+
+non_diegetic_music:
+N/A\
+"""
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# The request the rewriter is shown.
+# ---------------------------------------------------------------------------------------------------------------
+
+_BASE_TASK_BRIEF = {
+ "t2va": "t2va — text only, no reference image. Build the whole timeline from the request.",
+ "i2va": (
+ "i2va — is the first frame of the video, at 0.00 seconds, and belongs to [Shot 1]. Cite "
+ " explicitly in [Shot 1] and anchor its style, subjects and composition before any action."
+ ),
+ "fl2va": (
+ "fl2va — Picture 1 is the first frame, at 0.00 seconds, and Picture 2 is the last frame, at the end of the "
+ "video. Write the motion path between them, and prefer a single shot."
+ ),
+ "l2va": "l2va — is the LAST frame of the video and belongs to the final shot, not to [Shot 1].",
+}
+
+
+def base_user_message(prompt: str, task: str, duration: float, keyframes: list[Image.Image]) -> list[dict]:
+ """The chat turn a base-task rewrite is asked for: the keyframes under their labels, then the request."""
+ content: list[dict] = []
+ labels = ["", ""]
+ for index, keyframe in enumerate(keyframes):
+ content.append({"type": "text", "text": f"{labels[index]}:"})
+ content.append({"type": "image", "image": _thumbnail(keyframe)})
+ content.append(
+ {
+ "type": "text",
+ "text": (
+ f"Task: {_BASE_TASK_BRIEF[task]}\n"
+ f"Video duration: {duration:.2f} seconds at 24 fps. Every cut timestamp must fall strictly inside it.\n"
+ f"\nUser request:\n{prompt.strip()}\n"
+ "\nAll three fields are mandatory, in order — write `non_diegetic_music: N/A` if the video carries no "
+ "audience-only score. Write the rewrite now, starting with `integrated_multimodal_description:`."
+ ),
+ }
+ )
+ return [{"role": "user", "content": content}]
+
+
+def reference_user_message(prompt: str, labels, references, duration: float | None) -> list[dict]:
+ """The chat turn a `ref2va` rewrite is asked for: every reference under its own label, then the request.
+
+ An audio reference never reaches the rewriter — Qwen3-VL has no audio tower, and it never reaches the conditioner
+ either — so it is announced by label and role only, and the rewrite has to place its relationship without having
+ heard it. A video is shown the way the conditioner sees it: sampled down to 2 fps, from the frames the reference
+ already decoded.
+ """
+ content: list[dict] = []
+ inventory = []
+ for (kind, label), reference in zip(labels, references):
+ inventory.append(f"{label} — {kind} reference")
+ if kind == "image":
+ content.append({"type": "text", "text": f"{label}:"})
+ content.append({"type": "image", "image": _thumbnail(reference.image)})
+ elif kind == "video":
+ content.append({"type": "text", "text": f"{label}:"})
+ content.append({"type": "video", "video": _video_frames(reference)})
+ else:
+ content.append(
+ {"type": "text", "text": f"{label}: an audio reference. You cannot hear it; describe its role only."}
+ )
+
+ duration_line = (
+ f"Video duration: {duration:.2f} seconds at 24 fps. Every cut timestamp must fall strictly inside it.\n"
+ if duration is not None
+ else "Video duration: taken from the reference soundtrack. Keep cut timestamps early and conservative.\n"
+ )
+ content.append(
+ {
+ "type": "text",
+ "text": (
+ "Task: full-reference generation (ref2va).\n"
+ "References, in the order the generator reads them, under the labels you must use:\n"
+ + "\n".join(f" {entry}" for entry in inventory)
+ + "\n"
+ + duration_line
+ + f"\nUser request:\n{prompt.strip()}\n"
+ "\nAll six sections are mandatory, in order. Write the rewrite now, starting with "
+ "`subject_definitions:`."
+ ),
+ }
+ )
+ return [{"role": "user", "content": content}]
+
+
+def _as_uint8_rgb(media) -> np.ndarray:
+ """An in-memory reference frame (or frame stack) as channels-last `uint8` RGB.
+
+ The library used to expose this as `reference_media_to_uint8`; the reference dataclasses replaced it, and the one
+ normalization left to do here is for media a caller built in memory rather than through `from_file` — which is
+ already `uint8` `(num_frames, height, width, 3)`. A `torch.Tensor` is channels-first, as everywhere else in
+ diffusers, a `np.ndarray` channels-last, and floating point values are read over `[0, 1]`.
+ """
+ if isinstance(media, list):
+ return np.stack([_as_uint8_rgb(item) for item in media])
+ if isinstance(media, Image.Image):
+ return np.asarray(media.convert("RGB"))
+ if isinstance(media, torch.Tensor):
+ media = media.movedim(-3, -1).cpu().numpy()
+ media = np.asarray(media)
+ if media.dtype != np.uint8:
+ media = (media * 255.0).round().clip(0, 255).astype(np.uint8)
+ return media
+
+
+def _thumbnail(image) -> Image.Image:
+ """A copy of an image small enough to look at cheaply. The encoding path never sees it."""
+ if not isinstance(image, Image.Image):
+ image = Image.fromarray(_as_uint8_rgb(image))
+ image = image.convert("RGB")
+ if max(image.size) > REWRITE_IMAGE_LONG_EDGE:
+ scale = REWRITE_IMAGE_LONG_EDGE / max(image.size)
+ image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))), Image.LANCZOS)
+ return image
+
+
+def _video_frames(reference) -> list[Image.Image]:
+ """A reference video at the 2 fps the conditioner reads it, capped so one clip cannot own the prefill.
+
+ Sampled off the frames the reference already decoded rather than off the file, so the rewriter looks at the same
+ footage the conditioner will and no container is opened twice.
+ """
+ frames = _as_uint8_rgb(reference.frames)
+ stride = max(1, round((reference.fps or MINIMAX_H3_FPS) / REWRITE_VIDEO_FPS))
+ sampled = list(range(0, len(frames), stride))[:REWRITE_VIDEO_MAX_FRAMES]
+ return [_thumbnail(Image.fromarray(frames[index])) for index in sampled]
+
+
+# ---------------------------------------------------------------------------------------------------------------
+# Generation and cleanup.
+# ---------------------------------------------------------------------------------------------------------------
+
+BASE_FIELDS = ("integrated_multimodal_description", "overall_soundscape", "non_diegetic_music")
+REF_FIELDS = (
+ "subject_definitions",
+ "summary",
+ "retention_analysis",
+ "detailed_description",
+ "overall_soundscape",
+ "non_diegetic_music",
+)
+
+_FENCE = re.compile(r"^\s*```[a-zA-Z]*\s*|\s*```\s*$")
+# The shapes a chat model signs off with. Only ever matched against *trailing* lines, so a rewrite that happens to
+# contain one of these phrases inside a shot description keeps it.
+_CHATTER = re.compile(
+ r"^\s*(?:```.*|(?:let me know|i hope|hope (?:this|that)|would you like|note|if you (?:want|need)|feel free|"
+ r"this (?:rewrite|prompt|follows))\b.*)$",
+ re.IGNORECASE,
+)
+_INSTRUCTION_LINE = re.compile(
+ r"^\s*(For the target video,.*?referenced\.|How the reference pictures align with the target video[^\n]*)\s*",
+ re.IGNORECASE,
+)
+
+
+def _generate(components, messages: list[dict], system: str, max_new_tokens: int) -> tuple[str, bool]:
+ """One greedy decode off the conditioner's own language-model head.
+
+ The conditioner is the full `Qwen3VLForConditionalGeneration`, so `generate` is available as it stands: the
+ encoding path just never calls it, reading `hidden_states[50]` off `text_encoder.model` instead. Greedy, because a
+ rewrite is a format-following task and sampling only adds ways to leave the format.
+ """
+ import torch
+
+ model, processor = components.text_encoder, components.processor
+ device = getattr(components, "_execution_device", None) or next(model.parameters()).device
+
+ started = time.time()
+ inputs = processor.apply_chat_template(
+ [{"role": "system", "content": system}, *messages],
+ tokenize=True,
+ add_generation_prompt=True,
+ return_dict=True,
+ return_tensors="pt",
+ # A frame list is already at the rate it should be read at, so no resampling in the processor.
+ processor_kwargs={"do_sample_frames": False},
+ )
+ inputs = {name: value.to(device) if hasattr(value, "to") else value for name, value in inputs.items()}
+ prefill = int(inputs["input_ids"].shape[1])
+
+ with torch.inference_mode():
+ generated = model.generate(
+ **inputs,
+ max_new_tokens=max_new_tokens,
+ do_sample=False,
+ num_beams=1,
+ repetition_penalty=1.0,
+ )
+ completion = generated[0, prefill:]
+ text = processor.batch_decode([completion], skip_special_tokens=True)[0]
+ # A decode that used its whole budget stopped where the budget ran out, not where the rewrite ended, and is the
+ # only case whose tail needs trimming. One that stopped on its own end token is complete by construction.
+ truncated = int(completion.shape[0]) >= max_new_tokens
+ print(
+ f"[rewrite] {prefill} prefill + {int(completion.shape[0])} decoded tokens in {time.time() - started:.1f}s"
+ f"{' (hit the budget)' if truncated else ''}",
+ flush=True,
+ )
+ if os.environ.get("REWRITE_DEBUG"):
+ print(f"[rewrite] raw:\n{text}\n[rewrite] end raw", flush=True)
+ return text, truncated
+
+
+def _drop_incomplete_tail(text: str, fields: tuple[str, ...]) -> str:
+ """Cut a decode that ran into the token budget back to something well formed.
+
+ A rewrite that stops mid-sentence is still fed to the encoder, so the tail is trimmed to the last sentence that
+ actually ended — and a whole trailing field that never got past its own name is dropped.
+ """
+ blocks = _split_fields(text, fields)
+ if not blocks:
+ return text
+ name, body = blocks[-1]
+ if not body.strip():
+ blocks = blocks[:-1]
+ elif body.rstrip()[-1] not in ".!?\"'>":
+ cut = max(body.rfind(mark) for mark in (". ", ".\n", "! ", "!\n", "? ", "?\n", " "))
+ if cut > 0:
+ end = cut + (4 if body[cut : cut + 4] == " " else 1)
+ blocks[-1] = (name, body[:end])
+ else:
+ blocks = blocks[:-1]
+ if not blocks:
+ return text
+ return "\n\n".join(f"{name}:{body}" if body.startswith("\n") else f"{name}: {body.strip()}" for name, body in blocks)
+
+
+def _split_fields(text: str, fields: tuple[str, ...]) -> list[tuple[str, str]]:
+ """The `name: body` blocks of a rewrite, in the order they appear."""
+ pattern = re.compile(rf"^({'|'.join(fields)})\s*:", re.MULTILINE)
+ matches = list(pattern.finditer(text))
+ return [
+ (match.group(1), text[match.end() : matches[index + 1].start() if index + 1 < len(matches) else len(text)])
+ for index, match in enumerate(matches)
+ ]
+
+
+def _strip_trailing_chatter(text: str) -> str:
+ """Drop the sign-off a chat model puts after the body: a closing fence, an offer to revise, a note."""
+ lines = text.rstrip().split("\n")
+ while lines and (not lines[-1].strip() or _CHATTER.match(lines[-1])):
+ lines.pop()
+ return "\n".join(lines)
+
+
+def _scrub_default_camera_qualifiers(text: str) -> str:
+ """Delete the amplitude and speed phrases the guide says are the unsaid defaults.
+
+ `with medium amplitude` and `at normal speed` are exactly what "medium amplitude and normal speed are usually
+ omitted" rules out, and the model reaches for them anyway. They carry no information, so removing them is a pure
+ deterministic normalization rather than a judgement about the rewrite.
+ """
+ text = re.sub(r"\s+(?:with|at)\s+(?:medium|moderate)\s+amplitude", "", text, flags=re.IGNORECASE)
+ return re.sub(r"\s+at\s+(?:normal|medium|moderate|steady)\s+speed", "", text, flags=re.IGNORECASE)
+
+
+def _repair_dialogue_tags(text: str) -> str:
+ """Put a spoken line back inside `` when the model reached for a tag out of its own chat vocabulary.
+
+ Qwen3-VL occasionally writes `[English] ... ` where the format wants
+ `[English] ... ` — the language tag and the line itself are right, only the wrapper is the model's own. The
+ wrapper is what MiniMax-H3 keys dialogue off, so it is renamed rather than left to be encoded as prose.
+ """
+ for tag in set(re.findall(r"<([a-zA-Z_][\w-]*)>\s*\[[A-Z]", text)) - {"d"}:
+ text = text.replace(f"<{tag}>", "").replace(f"{tag}>", " ")
+ return text
+
+
+def clean_core(text: str, fields: tuple[str, ...], truncated: bool = False) -> str:
+ """Strip everything a chat model adds around the body it was asked for.
+
+ Fences, an echoed instruction line, a leading `assistant`, a stray closing remark after the last field, a dialogue
+ wrapper out of the model's own chat vocabulary, and — only when the decode ran out of budget — a tail that stops
+ mid-sentence. What is left is the prompt body itself, which is what the encoder is handed.
+ """
+ text = text.strip()
+ text = _FENCE.sub("", text).strip()
+ text = _strip_trailing_chatter(text)
+ text = re.sub(r"^assistant\s*[:\n]", "", text, flags=re.IGNORECASE).strip()
+ # The instruction line is the caller's to assemble, so an echoed one is dropped rather than duplicated.
+ text = _INSTRUCTION_LINE.sub("", text).strip()
+ text = _repair_dialogue_tags(text)
+ text = _scrub_default_camera_qualifiers(text)
+
+ blocks = _split_fields(text, fields)
+ if blocks:
+ # Anything before the first field name is chatter; anything after the last field's body is a closing remark.
+ text = text[text.index(blocks[0][0]) :].strip()
+ elif "[Shot 1]" in text and fields is BASE_FIELDS:
+ # The body without its field name still is a body; name it rather than throwing the rewrite away.
+ text = f"{BASE_FIELDS[0]}: {text}"
+
+ if truncated:
+ text = _drop_incomplete_tail(text, fields)
+ # One blank line between fields, whatever the model emitted.
+ text = re.sub(r"\n{3,}", "\n\n", text)
+ for name in fields[1:]:
+ text = re.sub(rf"\n+({name}\s*:)", r"\n\n\1", text)
+ return text.strip()
+
+
+def refine_keyframe_prompt(components, prompt, image, last_image, num_frames, max_new_tokens) -> str:
+ """Rewrite a `/encode` request into MiniMax-H3's trained prompt format.
+
+ Args:
+ components: The conditioner pipeline, for its `text_encoder` and `processor`.
+ prompt (`str`): The user's raw prompt.
+ image (`PIL.Image.Image`, *optional*): The first keyframe, if any.
+ last_image (`PIL.Image.Image`, *optional*): The last keyframe, if any.
+ num_frames (`int`): The requested frame count, before alignment.
+ max_new_tokens (`int`): The decode budget.
+
+ Returns:
+ `str`: the final prompt — the code-assembled instruction line, a blank line, and the model's core fields.
+ """
+ task = base_task(image, last_image)
+ duration = duration_seconds(num_frames)
+ keyframes = [keyframe for keyframe in (image, last_image) if keyframe is not None]
+
+ raw, truncated = _generate(
+ components, base_user_message(prompt, task, duration, keyframes), BASE_SYSTEM_PROMPT, max_new_tokens
+ )
+ core = clean_core(raw, BASE_FIELDS, truncated)
+ if not core:
+ raise ValueError("The rewrite came back empty.")
+
+ instruction = base_instruction_line(task, duration, final_shot_index(core))
+ return core if instruction is None else f"{instruction}\n\n{core}"
+
+
+def refine_reference_prompt(components, prompt, references, num_frames, max_new_tokens) -> str:
+ """Rewrite an `/encode_ref2va` request into MiniMax-H3's full-reference format.
+
+ Full-reference mode has no alignment instruction line: the six sections are the whole prompt, and the reference
+ relationships that a keyframe task states in one templated sentence are what `subject_definitions`,
+ `summary` and `retention_analysis` carry instead.
+
+ Args:
+ components: The conditioner pipeline, for its `text_encoder` and `processor`.
+ prompt (`str`): The user's raw prompt.
+ references (`list[MiniMaxH3Reference]`): The built references, in packed order — their labels and their media.
+ num_frames (`int` or `None`): The requested frame count, `None` when the soundtrack sets it.
+ max_new_tokens (`int`): The decode budget.
+
+ Returns:
+ `str`: the final prompt, the six sections as the model wrote them.
+ """
+ labels = reference_labels(references)
+ duration = reference_duration_seconds(references, num_frames)
+
+ raw, truncated = _generate(
+ components, reference_user_message(prompt, labels, references, duration), REF_SYSTEM_PROMPT, max_new_tokens
+ )
+ core = clean_core(raw, REF_FIELDS, truncated)
+ if not core:
+ raise ValueError("The rewrite came back empty.")
+ return core
diff --git a/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/requirements.txt b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..6b86bbd7641f791e3fd5bd80b402e8e6f5ef2bf8
--- /dev/null
+++ b/.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/requirements.txt
@@ -0,0 +1,33 @@
+# `diffusers` is installed from the canonical MiniMax-H3 pull request,
+# https://github.com/huggingface/diffusers/pull/14371 ("Minimax h3 follow up (review & refactor)"), pinned to a
+# **commit** rather than to its `minimax-h3-refactor` branch: the PR is a WIP and its head moves, and this Space's
+# blocks subclass its block classes. Re-pin — and re-check `h3_split_blocks.py` against the block names of the new
+# head — whenever the PR updates.
+#
+# 665f578278365ea4a3318cb8c9b66ce6c01204b9 = refs/pull/14371/head at the time of this deploy
+#
+# Nothing is quantized in this deployment, so there is no `torchao`.
+--extra-index-url https://download.pytorch.org/whl/cu130
+diffusers @ git+https://github.com/huggingface/diffusers.git@665f578278365ea4a3318cb8c9b66ce6c01204b9
+torch==2.11.0
+torchvision==0.26.0
+# A reference soundtrack that is not already at the audio VAE's 32 kHz is resampled with `torchaudio`, which is the
+# only thing in the `ref2va` path that needs it — and it is easy to miss, because a 32 kHz reference skips the
+# resample entirely. Both halves need it: the conditioner's `setup` step normalizes the same waveforms this one does.
+torchaudio==2.11.0
+# Pinned to the version the MiniMax-H3 parity work was verified against: the Qwen3-VL processor decides the vision
+# patch count, so a different minor changes the conditioning.
+transformers==5.8.0
+accelerate==1.14.0
+# diffusers pins <2; 1.24.0 is the version the parity work ran on.
+huggingface-hub==1.24.0
+gradio==6.20.0
+spaces==0.51.1
+# No `kernels` pin on purpose: the Hub attention backends want `kernels>=0.12.3`, and that version breaks
+# transformers 5.8.0 at import.
+# PyAV: decoding the reference media, i.e. `MiniMaxH3VideoReference.from_file` / `MiniMaxH3AudioReference.from_file`.
+av
+pillow
+numpy
+requests
+safetensors>=0.8.0
diff --git a/.slim/codemap.json b/.slim/codemap.json
new file mode 100644
index 0000000000000000000000000000000000000000..e0e8fd8b3bacac3a0d915a2400814c898697852a
--- /dev/null
+++ b/.slim/codemap.json
@@ -0,0 +1,34 @@
+{
+ "metadata": {
+ "version": "1.0.0",
+ "last_run": "2026-08-08T14:28:47.030Z",
+ "root": "C:\\Users\\Alan\\Desktop\\portableapps\\minimax-h3-repo",
+ "include_patterns": [
+ "*.py",
+ "*.txt",
+ "*.md",
+ "*.gitattributes"
+ ],
+ "exclude_patterns": [
+ ".cache/**",
+ "__pycache__/**",
+ "examples/**"
+ ],
+ "exceptions": []
+ },
+ "file_hashes": {
+ ".gitattributes": "a859f8a89685747ffd4171b870540c41",
+ "AGENTS.md": "e6b229267954151e3788539ee543e1aa",
+ "README.md": "19d13b1ab44ef10a78299fde3b0b1f77",
+ "app.py": "826c182cd74ec65244f70484b9b70149",
+ "codemap.md": "6e3339b340ed7930d52910b6226540a7",
+ "h3_aoti.py": "97e2d71787f0c6690480747d84a8bf91",
+ "h3_split_blocks.py": "10998027d8aa8c3c4886d2e32eccbb45",
+ "packages.txt": "ea0c93931b23f89c79cf4f8822ec0c85",
+ "requirements.txt": "3081962f013df35bff515811fb574e92",
+ "spaces_constant_binding_patch.py": "9eb60bae9876a1a4f432ed1601482cb0"
+ },
+ "folder_hashes": {
+ ".": "f67bb5d11661914fe466998247be9611"
+ }
+}
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 0000000000000000000000000000000000000000..b579780a1b390689a9c9464bd5560f919124e183
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,76 @@
+# AGENTS.md
+
+## What this is
+
+A Hugging Face **ZeroGPU Space** — the denoising half of a two-Space MiniMax-H3 deployment (video + soundtrack, unquantized bf16). Gradio 6.20.0 app, entrypoint `app.py`. The other half, `multimodalart/qwen3vl-conditioner` (62 GiB Qwen3-VL), is called over the gradio API on every request.
+
+The pipeline is split in `h3_split_blocks.py`: `MiniMaxH3GeneratorBlocks` subclasses diffusers' `SequentialPipelineBlocks` with the `text_encoder` step removed. The entire wire format between the halves is two tensors — `prompt_embeds` `(1, num_text_tokens, 5120)` bf16 + `text_token_tags` int64 — carried as one safetensors file with `height`/`width`/`num_frames` in its metadata.
+
+## Critical rules
+
+- **`import spaces` must precede anything that initializes CUDA** (first import in `app.py`; it patches `torch.cuda` so the 77.3 GB load happens at startup on CPU). Importing torch first silently breaks placement.
+- **`diffusers` is not from PyPI.** It installs from `refs/pull/14371/head` of huggingface/diffusers, pinned to commit `665f5782` (see `requirements.txt`). The PR is WIP and moves: when re-pinning, re-check `h3_split_blocks.py` against the new head's block class names, and update the pin and the comment together.
+- **Canvas labels are a cross-Space contract.** `CANVASES` keys in `app.py` go over the wire by label and must stay identical to the conditioner's table — an unknown label is rejected there. Adding a canvas means editing both Spaces.
+- **This half cannot run standalone** (no conditioner → no request) and **nothing here can be tested end-to-end offline**: every generate calls the live conditioner Space.
+- Changing the wire format, canvas table, or frame-count arithmetic breaks the split — coordinate with the conditioner Space.
+
+## Non-obvious constraints
+
+- Frames are snapped to `17 * n + 5` at 24 fps (`snap_frames`). UI duration is 2–14 s because 15 s rounds to 362 frames, which the checkpoint refuses.
+- **Autoencoders must stay float32** (`_keep_in_fp32_modules`): a bf16 audio VAE decodes the soundtrack ~20 dB too quiet. `load_components(dtype=torch.bfloat16)` is for the transformer only.
+- `H3_PLACEMENT` default is **`pack`** in code (pack the transformer at startup; the ~10 GB fp32 VAEs move on the first GPU call). `lazy` moves all 72 GiB on the first GPU call; `offload` needs `_arm_decode_hooks` or the VAE decodes never offload. The README's env table says `lazy` — the code is authoritative.
+- `@spaces.GPU(duration=get_duration, ...)`: the duration estimate must keep covering the request (linear + quadratic in latent rows) or the call is evicted. If you raise steps or canvas sizes, re-check `_DUR_B`/`_DUR_C`.
+- `_generate` returns only the three outputs because a `@spaces.GPU` return crosses a process boundary by pickling; the pipeline state stays on the card.
+- `upsample` is deliberately the **last** positional arg of `generate` so positional API clients that predate it are unaffected — keep it last.
+- Attention backend default `_native_cudnn`; flash-attention 3 is sm90-only and this pool is sm120 — don't switch without measuring.
+
+## Dependency pins (all deliberate)
+
+- `torch==2.11.0` (cu130 wheel index); bumping breaks the AoTI key match (silent fallback to eager).
+- `transformers==5.8.0` — "a different minor changes the conditioning": the Qwen3-VL processor decides the vision patch count.
+- `huggingface-hub==1.24.0` because diffusers pins `<2`.
+- No `kernels` pin on purpose: `kernels>=0.12.3` breaks transformers 5.8.0 at import. Do not add one.
+- `gradio==6.20.0` must match `sdk_version` in the README frontmatter (the Space config).
+
+## AoTI (`h3_aoti.py`)
+
+- `maybe_load` is the runtime path: off unless `H3_AOTI=1`, and only for the exact key `bf16/torch2.11/sm120/dynamic`; any mismatch falls back to eager. **A mismatched compiled package segfaults rather than raising** — hence the strict key check. `H3_AOTI_TEMB_ROWS` (default 4) must match the published package.
+- `export_block` / `compile_and_save` / `upload` are the **build path** for the separate `multimodalart/minimax-h3-aoti` repo: requires a same-arch GPU and `HF_TOKEN`, runs on GPU time, and the Space itself never calls it.
+- `spaces_constant_binding_patch.py` patches `spaces`' AoTI loader so anonymously-lifted constants bind instead of segfaulting; keep it importable (missing is tolerated with a warning).
+
+## Layout
+
+- `app.py` — Gradio UI + request flow (`generate` → remote `/encode` → `@spaces.GPU _generate` → `encode_video` mux). `load_models()` runs at import; a syntax/import error kills the whole Space at startup.
+- `h3_split_blocks.py` — the split block classes for both partitions (`t2va` and `ref2va`) and both halves.
+- `h3_aoti.py` — AoTI runtime patch + export/compile/upload build path.
+- `spaces_constant_binding_patch.py` — AoTI constant-binding patch (compile side + load side).
+- `packages.txt` (`ffmpeg`) — ZeroGPU apt packages, Space-level, not pip. `examples/` holds keyframe images for `gr.Examples`.
+
+## Repository Map
+
+A full codemap is available at `codemap.md` in the project root.
+
+Before working on any task, read `codemap.md` to understand:
+- Project architecture and entry points
+- Directory responsibilities and design patterns
+- Data flow and integration points between modules
+
+For deep work on a specific folder, also read that folder's `codemap.md`.
+
+## Verification
+
+No tests, lint, or CI in this repo. Local sanity check:
+
+```
+python -m py_compile app.py h3_aoti.py h3_split_blocks.py spaces_constant_binding_patch.py
+```
+
+Full verification is running the Space itself — 77.3 GB startup download plus a real generation against the live conditioner — treat that as an expensive manual gate, not a routine step.
+
+## Cloned Dependency Source
+
+Read-only dependency source repositories are available under
+`.slim/clonedeps/repos/` for inspection. Do not edit these clones.
+
+- `.slim/clonedeps/repos/huggingface__diffusers/` — huggingface/diffusers at the exact pinned commit `665f5782` (sparse-checked to `src/diffusers/modular_pipelines/`); the `SequentialPipelineBlocks` + MiniMax-H3 block classes `h3_split_blocks.py` subclasses and `h3_aoti.py` exports.
+- `.slim/clonedeps/repos/multimodalart__qwen3vl-conditioner/` — the sibling conditioner Space at `main`; the `CANVASES`/wire-format contract this Space must match, plus the `h3_core` / `prompt_rewrite` modules this repo references.
diff --git a/__pycache__/app.cpython-310.pyc b/__pycache__/app.cpython-310.pyc
index dea4933d50005906e217da8c44d06fcc98bc0a74..32fd7b498931fc084d5bcdcdc9bb3ebc5049df1d 100644
Binary files a/__pycache__/app.cpython-310.pyc and b/__pycache__/app.cpython-310.pyc differ
diff --git a/__pycache__/h3_aoti.cpython-310.pyc b/__pycache__/h3_aoti.cpython-310.pyc
index 46df86e96aa991ce9ce3704efd13c171a4c889e5..30ca07d2d27f438aff6a9d2b14d98d27ca85efa6 100644
Binary files a/__pycache__/h3_aoti.cpython-310.pyc and b/__pycache__/h3_aoti.cpython-310.pyc differ
diff --git a/__pycache__/h3_split_blocks.cpython-310.pyc b/__pycache__/h3_split_blocks.cpython-310.pyc
index 8c440c370d5ddc02abb78bcfa4f2cccde4f8867b..3ccb21a9794be568bc7cf476ffeb0f557cca2531 100644
Binary files a/__pycache__/h3_split_blocks.cpython-310.pyc and b/__pycache__/h3_split_blocks.cpython-310.pyc differ
diff --git a/__pycache__/spaces_constant_binding_patch.cpython-310.pyc b/__pycache__/spaces_constant_binding_patch.cpython-310.pyc
index 175c57c4dbfaef6dc61b532c53cb52ec8b7a4834..36ea811563b1c302948810519fc61478dff5c182 100644
Binary files a/__pycache__/spaces_constant_binding_patch.cpython-310.pyc and b/__pycache__/spaces_constant_binding_patch.cpython-310.pyc differ
diff --git a/app.py b/app.py
index a832ddfbd71c06aaaf67fb60c4805645ef26db52..7d505cf35149405a889a3e80ef64c13a4a808b32 100644
--- a/app.py
+++ b/app.py
@@ -22,6 +22,11 @@ PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
+# lightx2v distillation LoRA: merges a 4-step delta into the transformer at startup.
+TURBO_LORA_REPO = os.environ.get("H3_TURBO_LORA", "lightx2v/Minimax-h3-Turbo")
+TURBO_LORA_FILE = "minimax_h3_fl2v_turbo_4step_v0.1.safetensors"
+TURBO_LORA_SCALE = 0.0625 # alpha(8) / rank(128)
+
# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
# is rejected there and surfaces as a failure here.
CANVASES = {
@@ -73,6 +78,7 @@ PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
+TURBO_AVAILABLE: bool = False
def status() -> str:
if LOAD_ERROR:
@@ -103,6 +109,7 @@ def load_models() -> str | None:
started = time.time()
try:
import torch
+ from huggingface_hub import hf_hub_download
from diffusers import ComponentsManager
from h3_split_blocks import MiniMaxH3GeneratorBlocks
@@ -115,6 +122,44 @@ def load_models() -> str | None:
pipe.load_components(dtype=torch.bfloat16)
pipe.transformer.set_attention_backend(ATTENTION)
+ # Merge the lightx2v 4-step distillation LoRA into the transformer at startup (CPU, no GPU time).
+ # The LoRA delta is small (~0.0625 scale) and bakes into the base weights permanently.
+ global TURBO_AVAILABLE
+ if TURBO_LORA_REPO:
+ from safetensors.torch import load_file as _load_lora
+
+ print(f"[gen] loading turbo LoRA from {TURBO_LORA_REPO} ...", flush=True)
+ _lora_path = hf_hub_download(TURBO_LORA_REPO, TURBO_LORA_FILE)
+ _lora_sd = _load_lora(_lora_path)
+
+ # Group LoRA A/B pairs by target module name.
+ _lora_pairs: dict[str, dict[str, torch.Tensor]] = {}
+ for _k, _v in _lora_sd.items():
+ if ".lora_A." in _k:
+ _base = _k.split(".lora_A.")[0]
+ _lora_pairs.setdefault(_base, {})["A"] = _v
+ elif ".lora_B." in _k:
+ _base = _k.split(".lora_B.")[0]
+ _lora_pairs.setdefault(_base, {})["B"] = _v
+
+ _merged = 0
+ for _base_name, _pair in _lora_pairs.items():
+ if "A" not in _pair or "B" not in _pair:
+ continue
+ _delta = (_pair["B"] @ _pair["A"]) * TURBO_LORA_SCALE
+ # Navigate from "transformer_blocks.0.attn.to_q" to the module.
+ _parts = _base_name.split(".")
+ _obj = pipe.transformer
+ for _p in _parts:
+ _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
+ _obj.weight.data += _delta.to(_obj.weight.dtype)
+ _merged += 1
+
+ del _lora_sd, _lora_pairs
+ import gc as _gc; _gc.collect()
+ TURBO_AVAILABLE = True
+ print(f"[gen] turbo LoRA merged ({_merged} modules)", flush=True)
+
# Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
# worker. Off unless `H3_AOTI=1`.
import h3_aoti
@@ -236,7 +281,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
-def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=25, seed=-1, upsample=False, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=25, seed=-1, upsample=False, turbo=False, progress=gr.Progress(track_tqdm=True)):
"""One request. `upsample` is last and defaults off, so a positional API client that predates it is unaffected."""
if LOAD_ERROR:
raise gr.Error(LOAD_ERROR)
@@ -251,6 +296,9 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
num_frames = snap_frames(duration)
+ if turbo and TURBO_AVAILABLE:
+ steps = 5 # 5 sigma points → 4 transformer evaluations
+
progress(0.0, desc=f"Upsampling the prompt on {CONDITIONER_SPACE} ..." if upsample else f"Conditioning on {CONDITIONER_SPACE} ...")
conditioned = time.time()
prompt_embeds, text_token_tags, metadata, plan = encode_remote(
@@ -370,6 +418,9 @@ with gr.Blocks(title="MiniMax-H3") as demo:
canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
duration = gr.Slider(label="Duration (s)", minimum=MIN_UI_DURATION, maximum=MAX_UI_DURATION, step=1, value=5)
steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=25)
+ turbo = gr.Checkbox(label="Turbo (4-step LoRA)", value=False,
+ visible=TURBO_AVAILABLE,
+ info="Merge distillation LoRA, force 4 steps")
seed = gr.Number(label="Seed (−1 = random)", value=-1, precision=0)
with gr.Column():
@@ -398,7 +449,7 @@ with gr.Blocks(title="MiniMax-H3") as demo:
run.click(
generate,
- [prompt, image, last_image, canvas, duration, steps, seed, upsample],
+ [prompt, image, last_image, canvas, duration, steps, seed, upsample, turbo],
[video, report, upsampled, upsampled_panel],
api_name="generate",
)
diff --git a/codemap.md b/codemap.md
new file mode 100644
index 0000000000000000000000000000000000000000..368ea5c998b8f02efd97145fb3218f78b3b262bf
--- /dev/null
+++ b/codemap.md
@@ -0,0 +1,40 @@
+# Repository Atlas: minimax-h3-repo
+
+## Project Responsibility
+
+The denoising half of a **two-Space MiniMax-H3 deployment** on Hugging Face ZeroGPU: joint video + synchronized soundtrack generation, unquantized bf16. The 195.9 GiB checkpoint cannot fit one ZeroGPU Space (150 GB storage cap), so `MiniMaxH3Blocks` is cut at its `text_encoder` step. This repo holds the transformer + two autoencoders; the 62 GiB Qwen3-VL conditioner runs in the sibling `multimodalart/qwen3vl-conditioner` Space and is called over the gradio API on every request.
+
+## System Entry Points
+
+- `app.py` — Gradio UI + request orchestration. **Runs `load_models()` at import** (77.3 GB startup load on CPU). Entrypoint per README frontmatter (`app_file: app.py`).
+- `requirements.txt` — Dependency manifest. **`diffusers` comes from a pinned WIP PR commit** (`refs/pull/14371/head` @ `665f5782`), not PyPI.
+- `README.md` — Space config (frontmatter `sdk_version` must match `gradio==6.20.0`) + design rationale and measured performance.
+
+## File Map
+
+| File | Responsibility | Detailed Notes |
+|------|---------------|----------------|
+| `app.py` | Request flow + Gradio UI | `generate` → remote `/encode` → `@spaces.GPU _generate` → `encode_video` mux. Owns `CANVASES` (cross-Space contract), `snap_frames` (17n+5), placement modes (`pack`/`lazy`/`offload`), `get_duration` GPU budget estimator. |
+| `h3_split_blocks.py` | Split pipeline definition | `SequentialPipelineBlocks` subclasses for both halves (`Conditioner`/`Generator`) and both partitions (`t2va`/`ref2va`). Defines the wire-format outputs (`prompt_embeds` + `text_token_tags` + resolved height/width/num_frames). |
+| `h3_aoti.py` | AoTI acceleration | `maybe_load`/`patch_blocks` = runtime path (strict key `bf16/torch2.11/sm120/dynamic`, else eager). `export_block`/`compile_and_save`/`upload` = build path for the separate `multimodalart/minimax-h3-aoti` repo (never called by the Space). |
+| `spaces_constant_binding_patch.py` | AoTI constant binding | Compile-side (`register_loose_tensors`, `write_constant_aliases`) + load-side (`apply_spaces_constant_binding_patch`) glue preventing SIGSEGV from anonymously-lifted export constants. |
+| `requirements.txt` | Pip dependencies | All pins deliberate — see notes in file. |
+| `packages.txt` | Space apt packages | `ffmpeg` (ZeroGPU system-level, not pip). |
+| `.gitattributes` | HF LFS rules | Standard Space LFS filters for model artifacts. |
+| `AGENTS.md` | Agent instructions | Critical rules, dependency pins, verification steps. |
+
+## Data & Control Flow
+
+1. Browser/client calls `generate(prompt, image?, last_image?, canvas, duration, steps, seed, upsample)`.
+2. `snap_frames` aligns the duration to `17*n + 5` frames at 24 fps; `encode_remote` calls the conditioner's `/encode` over `gradio_client` (forwarding the caller's own ZeroGPU token).
+3. Conditioner returns one safetensors file carrying `prompt_embeds` `(1, num_text_tokens, 5120)` bf16 + `text_token_tags` int64, with resolved `height`/`width`/`num_frames` in metadata + the plan.
+4. `_generate` (decorated `@spaces.GPU(duration=get_duration, size=xlarge)`) runs the packed-sequence denoise loop + two decoders on GPU. Returns only `(video frames, audio, sampling_rate)` — a `@spaces.GPU` return crosses a process boundary by pickling, so pipeline state stays on the card.
+5. `encode_video` (PyAV) muxes h264 frames + stereo AAC soundtrack into an MP4 in temp storage.
+
+## Integration Points
+
+- **Cross-Space wire contract**: `CANVASES` label strings in `app.py` go over the wire and must stay identical to the conditioner's table; the wire format (two tensors + metadata) must stay in lockstep. Changing either breaks the split.
+- **AoTI package repo**: `multimodalart/minimax-h3-aoti` — consumed at runtime via `snapshot_download` when `H3_AOTI=1`; produced by this repo's build path (requires `HF_TOKEN` + same-arch GPU).
+- **Model repo**: `MiniMaxAI/MiniMax-H3` (public diffusers checkpoint, subfolders `transformer/` + `vae/` + `audio_vae/`).
+- **Conditioner Space**: `multimodalart/qwen3vl-conditioner` — remote `/encode` endpoint, billed to the requesting user's token.
+- **diffusers upstream**: `huggingface/diffusers` PR #14371 — the pinned WIP commit whose block classes `h3_split_blocks.py` subclasses.