observantdistressed commited on
Commit
6891cba
·
verified ·
1 Parent(s): 7ec04b1

feat: add turbo LoRA option (lightx2v 4-step distillation)

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitignore +3 -0
  2. .ignore +9 -0
  3. .slim/clonedeps.json +23 -0
  4. .slim/clonedeps/repos/huggingface__diffusers/.gitignore +188 -0
  5. .slim/clonedeps/repos/huggingface__diffusers/AGENTS.md +1 -0
  6. .slim/clonedeps/repos/huggingface__diffusers/CITATION.cff +52 -0
  7. .slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md +1 -0
  8. .slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md +130 -0
  9. .slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md +1 -0
  10. .slim/clonedeps/repos/huggingface__diffusers/LICENSE +201 -0
  11. .slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in +2 -0
  12. .slim/clonedeps/repos/huggingface__diffusers/Makefile +120 -0
  13. .slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md +1 -0
  14. .slim/clonedeps/repos/huggingface__diffusers/README.md +232 -0
  15. .slim/clonedeps/repos/huggingface__diffusers/SECURITY.md +106 -0
  16. .slim/clonedeps/repos/huggingface__diffusers/_typos.toml +13 -0
  17. .slim/clonedeps/repos/huggingface__diffusers/pyproject.toml +32 -0
  18. .slim/clonedeps/repos/huggingface__diffusers/setup.py +305 -0
  19. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py +1746 -0
  20. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py +244 -0
  21. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py +752 -0
  22. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py +34 -0
  23. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py +57 -0
  24. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py +1468 -0
  25. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py +230 -0
  26. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py +47 -0
  27. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py +714 -0
  28. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py +120 -0
  29. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py +211 -0
  30. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py +404 -0
  31. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py +381 -0
  32. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py +52 -0
  33. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py +1109 -0
  34. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py +49 -0
  35. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py +113 -0
  36. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py +1331 -0
  37. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py +166 -0
  38. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py +259 -0
  39. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py +889 -0
  40. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py +1056 -0
  41. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +1301 -0
  42. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +240 -0
  43. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +130 -0
  44. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py +47 -0
  45. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py +270 -0
  46. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py +92 -0
  47. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py +236 -0
  48. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py +264 -0
  49. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py +200 -0
  50. .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_pipeline.py +110 -0
.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ # BEGIN oh-my-opencode-slim clonedeps
2
+ .slim/clonedeps/repos/
3
+ # END oh-my-opencode-slim clonedeps
.ignore ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ # BEGIN oh-my-opencode-slim clonedeps
2
+ !.slim/
3
+ !.slim/clonedeps.json
4
+ !.slim/clonedeps/
5
+ !.slim/clonedeps/repos/
6
+ !.slim/clonedeps/repos/**
7
+ .slim/clonedeps/repos/**/.git/
8
+ .slim/clonedeps/repos/**/.git/**
9
+ # END oh-my-opencode-slim clonedeps
.slim/clonedeps.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "version": "1.0.0",
3
+ "updatedAt": "2026-08-08T14:28:19.000Z",
4
+ "dependencies": [
5
+ {
6
+ "name": "diffusers",
7
+ "resolvedVersion": "665f578278365ea4a3318cb8c9b66ce6c01204b9 (refs/pull/14371/head pin)",
8
+ "repoUrl": "https://github.com/huggingface/diffusers.git",
9
+ "ref": "665f578278365ea4a3318cb8c9b66ce6c01204b9",
10
+ "path": ".slim/clonedeps/repos/huggingface__diffusers",
11
+ "packagePath": "src/diffusers/modular_pipelines",
12
+ "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"
13
+ },
14
+ {
15
+ "name": "multimodalart/qwen3vl-conditioner",
16
+ "resolvedVersion": "20d92357b94f1ee6a93b4fcaeaa06b9d151122d7 (main)",
17
+ "repoUrl": "https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner",
18
+ "ref": "main",
19
+ "path": ".slim/clonedeps/repos/multimodalart__qwen3vl-conditioner",
20
+ "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"
21
+ }
22
+ ]
23
+ }
.slim/clonedeps/repos/huggingface__diffusers/.gitignore ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Initially taken from GitHub's Python gitignore file
2
+
3
+ # Byte-compiled / optimized / DLL files
4
+ __pycache__/
5
+ *.py[cod]
6
+ *$py.class
7
+
8
+ # C extensions
9
+ *.so
10
+
11
+ # tests and logs
12
+ tests/fixtures/cached_*_text.txt
13
+ logs/
14
+ lightning_logs/
15
+ lang_code_data/
16
+
17
+ # Distribution / packaging
18
+ .Python
19
+ build/
20
+ develop-eggs/
21
+ dist/
22
+ downloads/
23
+ eggs/
24
+ .eggs/
25
+ lib/
26
+ lib64/
27
+ parts/
28
+ sdist/
29
+ var/
30
+ wheels/
31
+ *.egg-info/
32
+ .installed.cfg
33
+ *.egg
34
+ MANIFEST
35
+
36
+ # PyInstaller
37
+ # Usually these files are written by a Python script from a template
38
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
39
+ *.manifest
40
+ *.spec
41
+
42
+ # Installer logs
43
+ pip-log.txt
44
+ pip-delete-this-directory.txt
45
+
46
+ # Unit test / coverage reports
47
+ htmlcov/
48
+ .tox/
49
+ .nox/
50
+ .coverage
51
+ .coverage.*
52
+ .cache
53
+ nosetests.xml
54
+ coverage.xml
55
+ *.cover
56
+ .hypothesis/
57
+ .pytest_cache/
58
+
59
+ # Translations
60
+ *.mo
61
+ *.pot
62
+
63
+ # Django stuff:
64
+ *.log
65
+ local_settings.py
66
+ db.sqlite3
67
+
68
+ # Flask stuff:
69
+ instance/
70
+ .webassets-cache
71
+
72
+ # Scrapy stuff:
73
+ .scrapy
74
+
75
+ # Sphinx documentation
76
+ docs/_build/
77
+
78
+ # PyBuilder
79
+ target/
80
+
81
+ # Jupyter Notebook
82
+ .ipynb_checkpoints
83
+
84
+ # IPython
85
+ profile_default/
86
+ ipython_config.py
87
+
88
+ # pyenv
89
+ .python-version
90
+
91
+ # celery beat schedule file
92
+ celerybeat-schedule
93
+
94
+ # SageMath parsed files
95
+ *.sage.py
96
+
97
+ # Environments
98
+ .env
99
+ .venv
100
+ env/
101
+ venv/
102
+ ENV/
103
+ env.bak/
104
+ venv.bak/
105
+
106
+ # Spyder project settings
107
+ .spyderproject
108
+ .spyproject
109
+
110
+ # Rope project settings
111
+ .ropeproject
112
+
113
+ # mkdocs documentation
114
+ /site
115
+
116
+ # mypy
117
+ .mypy_cache/
118
+ .dmypy.json
119
+ dmypy.json
120
+
121
+ # Pyre type checker
122
+ .pyre/
123
+
124
+ # vscode
125
+ .vs
126
+ .vscode
127
+
128
+ # Cursor
129
+ .cursor
130
+
131
+ # Pycharm
132
+ .idea
133
+
134
+ # TF code
135
+ tensorflow_code
136
+
137
+ # Models
138
+ proc_data
139
+
140
+ # examples
141
+ runs
142
+ /runs_old
143
+ /wandb
144
+ /examples/runs
145
+ /examples/**/*.args
146
+ /examples/rag/sweep
147
+
148
+ # data
149
+ /data
150
+ serialization_dir
151
+
152
+ # emacs
153
+ *.*~
154
+ debug.env
155
+
156
+ # vim
157
+ .*.swp
158
+
159
+ # ctags
160
+ tags
161
+
162
+ # pre-commit
163
+ .pre-commit*
164
+
165
+ # .lock
166
+ *.lock
167
+
168
+ # DS_Store (MacOS)
169
+ .DS_Store
170
+
171
+ # RL pipelines may produce mp4 outputs
172
+ *.mp4
173
+ *.jpg
174
+ *.jepg
175
+ *.wav
176
+
177
+ # dependencies
178
+ /transformers
179
+
180
+ # ruff
181
+ .ruff_cache
182
+
183
+ # wandb
184
+ wandb
185
+
186
+ # AI agent generated symlinks
187
+ /.agents/skills
188
+ /.claude/skills
.slim/clonedeps/repos/huggingface__diffusers/AGENTS.md ADDED
@@ -0,0 +1 @@
 
 
1
+ .ai/AGENTS.md
.slim/clonedeps/repos/huggingface__diffusers/CITATION.cff ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ cff-version: 1.2.0
2
+ title: 'Diffusers: State-of-the-art diffusion models'
3
+ message: >-
4
+ If you use this software, please cite it using the
5
+ metadata from this file.
6
+ type: software
7
+ authors:
8
+ - given-names: Patrick
9
+ family-names: von Platen
10
+ - given-names: Suraj
11
+ family-names: Patil
12
+ - given-names: Anton
13
+ family-names: Lozhkov
14
+ - given-names: Pedro
15
+ family-names: Cuenca
16
+ - given-names: Nathan
17
+ family-names: Lambert
18
+ - given-names: Kashif
19
+ family-names: Rasul
20
+ - given-names: Mishig
21
+ family-names: Davaadorj
22
+ - given-names: Dhruv
23
+ family-names: Nair
24
+ - given-names: Sayak
25
+ family-names: Paul
26
+ - given-names: Steven
27
+ family-names: Liu
28
+ - given-names: William
29
+ family-names: Berman
30
+ - given-names: Yiyi
31
+ family-names: Xu
32
+ - given-names: Thomas
33
+ family-names: Wolf
34
+ repository-code: 'https://github.com/huggingface/diffusers'
35
+ abstract: >-
36
+ Diffusers provides pretrained diffusion models across
37
+ multiple modalities, such as vision and audio, and serves
38
+ as a modular toolbox for inference and training of
39
+ diffusion models.
40
+ keywords:
41
+ - deep-learning
42
+ - pytorch
43
+ - image-generation
44
+ - hacktoberfest
45
+ - diffusion
46
+ - text2image
47
+ - image2image
48
+ - score-based-generative-modeling
49
+ - stable-diffusion
50
+ - stable-diffusion-diffusers
51
+ license: Apache-2.0
52
+ version: 0.12.1
.slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md ADDED
@@ -0,0 +1 @@
 
 
1
+ .ai/AGENTS.md
.slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ # Contributor Covenant Code of Conduct
3
+
4
+ ## Our Pledge
5
+
6
+ We as members, contributors, and leaders pledge to make participation in our
7
+ community a harassment-free experience for everyone, regardless of age, body
8
+ size, visible or invisible disability, ethnicity, sex characteristics, gender
9
+ identity and expression, level of experience, education, socio-economic status,
10
+ nationality, personal appearance, race, caste, color, religion, or sexual identity
11
+ and orientation.
12
+
13
+ We pledge to act and interact in ways that contribute to an open, welcoming,
14
+ diverse, inclusive, and healthy community.
15
+
16
+ ## Our Standards
17
+
18
+ Examples of behavior that contributes to a positive environment for our
19
+ community include:
20
+
21
+ * Demonstrating empathy and kindness toward other people
22
+ * Being respectful of differing opinions, viewpoints, and experiences
23
+ * Giving and gracefully accepting constructive feedback
24
+ * Accepting responsibility and apologizing to those affected by our mistakes,
25
+ and learning from the experience
26
+ * Focusing on what is best not just for us as individuals, but for the
27
+ overall Diffusers community
28
+
29
+ Examples of unacceptable behavior include:
30
+
31
+ * The use of sexualized language or imagery, and sexual attention or
32
+ advances of any kind
33
+ * Trolling, insulting or derogatory comments, and personal or political attacks
34
+ * Public or private harassment
35
+ * Publishing others' private information, such as a physical or email
36
+ address, without their explicit permission
37
+ * Spamming issues or PRs with links to projects unrelated to this library
38
+ * Other conduct which could reasonably be considered inappropriate in a
39
+ professional setting
40
+
41
+ ## Enforcement Responsibilities
42
+
43
+ Community leaders are responsible for clarifying and enforcing our standards of
44
+ acceptable behavior and will take appropriate and fair corrective action in
45
+ response to any behavior that they deem inappropriate, threatening, offensive,
46
+ or harmful.
47
+
48
+ Community leaders have the right and responsibility to remove, edit, or reject
49
+ comments, commits, code, wiki edits, issues, and other contributions that are
50
+ not aligned to this Code of Conduct, and will communicate reasons for moderation
51
+ decisions when appropriate.
52
+
53
+ ## Scope
54
+
55
+ This Code of Conduct applies within all community spaces, and also applies when
56
+ an individual is officially representing the community in public spaces.
57
+ Examples of representing our community include using an official e-mail address,
58
+ posting via an official social media account, or acting as an appointed
59
+ representative at an online or offline event.
60
+
61
+ ## Enforcement
62
+
63
+ Instances of abusive, harassing, or otherwise unacceptable behavior may be
64
+ reported to the community leaders responsible for enforcement at
65
+ feedback@huggingface.co.
66
+ All complaints will be reviewed and investigated promptly and fairly.
67
+
68
+ All community leaders are obligated to respect the privacy and security of the
69
+ reporter of any incident.
70
+
71
+ ## Enforcement Guidelines
72
+
73
+ Community leaders will follow these Community Impact Guidelines in determining
74
+ the consequences for any action they deem in violation of this Code of Conduct:
75
+
76
+ ### 1. Correction
77
+
78
+ **Community Impact**: Use of inappropriate language or other behavior deemed
79
+ unprofessional or unwelcome in the community.
80
+
81
+ **Consequence**: A private, written warning from community leaders, providing
82
+ clarity around the nature of the violation and an explanation of why the
83
+ behavior was inappropriate. A public apology may be requested.
84
+
85
+ ### 2. Warning
86
+
87
+ **Community Impact**: A violation through a single incident or series
88
+ of actions.
89
+
90
+ **Consequence**: A warning with consequences for continued behavior. No
91
+ interaction with the people involved, including unsolicited interaction with
92
+ those enforcing the Code of Conduct, for a specified period of time. This
93
+ includes avoiding interactions in community spaces as well as external channels
94
+ like social media. Violating these terms may lead to a temporary or
95
+ permanent ban.
96
+
97
+ ### 3. Temporary Ban
98
+
99
+ **Community Impact**: A serious violation of community standards, including
100
+ sustained inappropriate behavior.
101
+
102
+ **Consequence**: A temporary ban from any sort of interaction or public
103
+ communication with the community for a specified period of time. No public or
104
+ private interaction with the people involved, including unsolicited interaction
105
+ with those enforcing the Code of Conduct, is allowed during this period.
106
+ Violating these terms may lead to a permanent ban.
107
+
108
+ ### 4. Permanent Ban
109
+
110
+ **Community Impact**: Demonstrating a pattern of violation of community
111
+ standards, including sustained inappropriate behavior, harassment of an
112
+ individual, or aggression toward or disparagement of classes of individuals.
113
+
114
+ **Consequence**: A permanent ban from any sort of public interaction within
115
+ the community.
116
+
117
+ ## Attribution
118
+
119
+ This Code of Conduct is adapted from the [Contributor Covenant][homepage],
120
+ version 2.1, available at
121
+ https://www.contributor-covenant.org/version/2/1/code_of_conduct.html.
122
+
123
+ Community Impact Guidelines were inspired by [Mozilla's code of conduct
124
+ enforcement ladder](https://github.com/mozilla/diversity).
125
+
126
+ [homepage]: https://www.contributor-covenant.org
127
+
128
+ For answers to common questions about this code of conduct, see the FAQ at
129
+ https://www.contributor-covenant.org/faq. Translations are available at
130
+ https://www.contributor-covenant.org/translations.
.slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md ADDED
@@ -0,0 +1 @@
 
 
1
+ docs/source/en/conceptual/contribution.md
.slim/clonedeps/repos/huggingface__diffusers/LICENSE ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, Any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright [yyyy] [name of copyright owner]
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
.slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ include LICENSE
2
+ include src/diffusers/utils/model_card_template.md
.slim/clonedeps/repos/huggingface__diffusers/Makefile ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ .PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples codex claude clean-ai
2
+
3
+ # make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!)
4
+ export PYTHONPATH = src
5
+
6
+ check_dirs := examples scripts src tests utils benchmarks
7
+
8
+ modified_only_fixup:
9
+ $(eval modified_py_files := $(shell python utils/get_modified_files.py $(check_dirs)))
10
+ @if test -n "$(modified_py_files)"; then \
11
+ echo "Checking/fixing $(modified_py_files)"; \
12
+ ruff check $(modified_py_files) --fix; \
13
+ ruff format $(modified_py_files);\
14
+ else \
15
+ echo "No library .py files were modified"; \
16
+ fi
17
+
18
+ # Update src/diffusers/dependency_versions_table.py
19
+
20
+ deps_table_update:
21
+ @python setup.py deps_table_update
22
+
23
+ deps_table_check_updated:
24
+ @md5sum src/diffusers/dependency_versions_table.py > md5sum.saved
25
+ @python setup.py deps_table_update
26
+ @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)
27
+ @rm md5sum.saved
28
+
29
+ # autogenerating code
30
+
31
+ autogenerate_code: deps_table_update
32
+
33
+ # Check that the repo is in a good state
34
+
35
+ repo-consistency:
36
+ python utils/check_dummies.py
37
+ python utils/check_repo.py
38
+ python utils/check_inits.py
39
+ python utils/check_forward_call_docstrings.py
40
+
41
+ # this target runs checks on all files
42
+
43
+ quality:
44
+ ruff check $(check_dirs) setup.py
45
+ ruff format --check $(check_dirs) setup.py
46
+ doc-builder style src/diffusers docs/source --max_len 119 --check_only
47
+ python utils/check_doc_toc.py
48
+
49
+ # Format source code automatically and check is there are any problems left that need manual fixing
50
+
51
+ extra_style_checks:
52
+ python utils/custom_init_isort.py
53
+ python utils/check_doc_toc.py --fix_and_overwrite
54
+
55
+ # this target runs checks on all files and potentially modifies some of them
56
+
57
+ style:
58
+ ruff check $(check_dirs) setup.py --fix
59
+ ruff format $(check_dirs) setup.py
60
+ doc-builder style src/diffusers docs/source --max_len 119
61
+ ${MAKE} autogenerate_code
62
+ ${MAKE} extra_style_checks
63
+
64
+ # Super fast fix and check target that only works on relevant modified files since the branch was made
65
+
66
+ fixup: modified_only_fixup extra_style_checks autogenerate_code repo-consistency
67
+
68
+ # Make marked copies of snippets of codes conform to the original
69
+
70
+ fix-copies:
71
+ python utils/check_copies.py --fix_and_overwrite
72
+ python utils/check_dummies.py --fix_and_overwrite
73
+
74
+ # Auto docstrings in modular blocks
75
+ modular-autodoctrings:
76
+ python utils/modular_auto_docstring.py
77
+
78
+ # Verify forward() / __call__() arguments are documented in their docstrings
79
+ check-forward-call-docstrings:
80
+ python utils/check_forward_call_docstrings.py
81
+
82
+ # Run tests for the library
83
+
84
+ test:
85
+ python -m pytest -n auto --dist=loadfile -s -v ./tests/
86
+
87
+ # Run tests for examples
88
+
89
+ test-examples:
90
+ python -m pytest -n auto --dist=loadfile -s -v ./examples/
91
+
92
+
93
+ # Release stuff
94
+
95
+ pre-release:
96
+ python utils/release.py
97
+
98
+ pre-patch:
99
+ python utils/release.py --patch
100
+
101
+ post-release:
102
+ python utils/release.py --post_release
103
+
104
+ post-patch:
105
+ python utils/release.py --post_release --patch
106
+
107
+ # AI agent symlinks
108
+
109
+ codex:
110
+ mkdir -p .agents
111
+ rm -rf .agents/skills
112
+ ln -snf ../.ai/skills .agents/skills
113
+
114
+ claude:
115
+ mkdir -p .claude
116
+ rm -rf .claude/skills
117
+ ln -snf ../.ai/skills .claude/skills
118
+
119
+ clean-ai:
120
+ rm -rf .agents/skills .claude/skills
.slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md ADDED
@@ -0,0 +1 @@
 
 
1
+ docs/source/en/conceptual/philosophy.md
.slim/clonedeps/repos/huggingface__diffusers/README.md ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!---
2
+ Copyright 2022 - The HuggingFace Team. All rights reserved.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ -->
16
+
17
+ <p align="center">
18
+ <br>
19
+ <img src="https://raw.githubusercontent.com/huggingface/diffusers/main/docs/source/en/imgs/diffusers_library.jpg" width="400"/>
20
+ <br>
21
+ <p>
22
+ <p align="center">
23
+ <a href="https://github.com/huggingface/diffusers/blob/main/LICENSE"><img alt="GitHub" src="https://img.shields.io/github/license/huggingface/datasets.svg?color=blue"></a>
24
+ <a href="https://github.com/huggingface/diffusers/releases"><img alt="GitHub release" src="https://img.shields.io/github/release/huggingface/diffusers.svg"></a>
25
+ <a href="https://pepy.tech/project/diffusers"><img alt="GitHub release" src="https://static.pepy.tech/badge/diffusers/month"></a>
26
+ <a href="CODE_OF_CONDUCT.md"><img alt="Contributor Covenant" src="https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg"></a>
27
+ <a href="https://twitter.com/diffuserslib"><img alt="X account" src="https://img.shields.io/twitter/url/https/twitter.com/diffuserslib.svg?style=social&label=Follow%20%40diffuserslib"></a>
28
+ </p>
29
+
30
+ 🤗 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).
31
+
32
+ 🤗 Diffusers offers three core components:
33
+
34
+ - 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.
35
+ - Interchangeable noise [schedulers](https://huggingface.co/docs/diffusers/api/schedulers/overview) for different diffusion speeds and output quality.
36
+ - 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.
37
+
38
+ ## Installation
39
+
40
+ 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.
41
+
42
+ ### PyTorch
43
+
44
+ With `pip` (official package):
45
+
46
+ ```bash
47
+ pip install --upgrade diffusers[torch]
48
+ ```
49
+
50
+ With `conda` (maintained by the community):
51
+
52
+ ```sh
53
+ conda install -c conda-forge diffusers
54
+ ```
55
+
56
+ ### Apple Silicon (M1/M2) support
57
+
58
+ Please refer to the [How to use Stable Diffusion in Apple Silicon](https://huggingface.co/docs/diffusers/optimization/mps) guide.
59
+
60
+ ## Quickstart
61
+
62
+ 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):
63
+
64
+ ```python
65
+ from diffusers import DiffusionPipeline
66
+ import torch
67
+
68
+ pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16)
69
+ pipeline.to("cuda")
70
+ pipeline("An image of a squirrel in Picasso style").images[0]
71
+ ```
72
+
73
+ You can also dig into the models and schedulers toolbox to build your own diffusion system:
74
+
75
+ ```python
76
+ from diffusers import DDPMScheduler, UNet2DModel
77
+ from PIL import Image
78
+ import torch
79
+
80
+ scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
81
+ model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
82
+ scheduler.set_timesteps(50)
83
+
84
+ sample_size = model.config.sample_size
85
+ noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
86
+ input = noise
87
+
88
+ for t in scheduler.timesteps:
89
+ with torch.no_grad():
90
+ noisy_residual = model(input, t).sample
91
+ prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
92
+ input = prev_noisy_sample
93
+
94
+ image = (input / 2 + 0.5).clamp(0, 1)
95
+ image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
96
+ image = Image.fromarray((image * 255).round().astype("uint8"))
97
+ image
98
+ ```
99
+
100
+ Check out the [Quickstart](https://huggingface.co/docs/diffusers/quicktour) to launch your diffusion journey today!
101
+
102
+ ## How to navigate the documentation
103
+
104
+ | **Documentation** | **What can I learn?** |
105
+ |---------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
106
+ | [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. |
107
+ | [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. |
108
+ | [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. |
109
+ | [Optimization](https://huggingface.co/docs/diffusers/optimization/fp16) | Guides for how to optimize your diffusion model to run faster and consume less memory. |
110
+ | [Training](https://huggingface.co/docs/diffusers/training/overview) | Guides for how to train a diffusion model for different tasks with different training techniques. |
111
+ ## Contribution
112
+
113
+ We ❤️ contributions from the open-source community!
114
+ If you want to contribute to this library, please check out our [Contribution guide](https://huggingface.co/docs/diffusers/main/en/conceptual/contribution).
115
+ 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).
116
+ You can look out for [issues](https://github.com/huggingface/diffusers/issues) you'd like to tackle to contribute to the library.
117
+ - 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
118
+ - 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
119
+ - See [New scheduler](https://github.com/huggingface/diffusers/issues?q=is%3Aopen+is%3Aissue+label%3A%22New+scheduler%22)
120
+
121
+ Also, say 👋 in our public Discord channel <a href="https://discord.gg/G7tWnz98XR"><img alt="Join us on Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a>. We discuss the hottest trends about diffusion models, help each other with contributions, personal projects or just hang out ☕.
122
+
123
+
124
+ ## Popular Tasks & Pipelines
125
+
126
+ <table>
127
+ <tr>
128
+ <th>Task</th>
129
+ <th>Pipeline</th>
130
+ <th>🤗 Hub</th>
131
+ </tr>
132
+ <tr style="border-top: 2px solid black">
133
+ <td>Unconditional Image Generation</td>
134
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/ddpm"> DDPM </a></td>
135
+ <td><a href="https://huggingface.co/google/ddpm-ema-church-256"> google/ddpm-ema-church-256 </a></td>
136
+ </tr>
137
+ <tr style="border-top: 2px solid black">
138
+ <td>Text-to-Image</td>
139
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img">Stable Diffusion Text-to-Image</a></td>
140
+ <td><a href="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"> stable-diffusion-v1-5/stable-diffusion-v1-5 </a></td>
141
+ </tr>
142
+ <tr>
143
+ <td>Text-to-Image</td>
144
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/unclip">unCLIP</a></td>
145
+ <td><a href="https://huggingface.co/kakaobrain/karlo-v1-alpha"> kakaobrain/karlo-v1-alpha </a></td>
146
+ </tr>
147
+ <tr>
148
+ <td>Text-to-Image</td>
149
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/deepfloyd_if">DeepFloyd IF</a></td>
150
+ <td><a href="https://huggingface.co/DeepFloyd/IF-I-XL-v1.0"> DeepFloyd/IF-I-XL-v1.0 </a></td>
151
+ </tr>
152
+ <tr>
153
+ <td>Text-to-Image</td>
154
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/kandinsky">Kandinsky</a></td>
155
+ <td><a href="https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder"> kandinsky-community/kandinsky-2-2-decoder </a></td>
156
+ </tr>
157
+ <tr style="border-top: 2px solid black">
158
+ <td>Text-guided Image-to-Image</td>
159
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/controlnet">ControlNet</a></td>
160
+ <td><a href="https://huggingface.co/lllyasviel/sd-controlnet-canny"> lllyasviel/sd-controlnet-canny </a></td>
161
+ </tr>
162
+ <tr>
163
+ <td>Text-guided Image-to-Image</td>
164
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/pix2pix">InstructPix2Pix</a></td>
165
+ <td><a href="https://huggingface.co/timbrooks/instruct-pix2pix"> timbrooks/instruct-pix2pix </a></td>
166
+ </tr>
167
+ <tr>
168
+ <td>Text-guided Image-to-Image</td>
169
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/img2img">Stable Diffusion Image-to-Image</a></td>
170
+ <td><a href="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"> stable-diffusion-v1-5/stable-diffusion-v1-5 </a></td>
171
+ </tr>
172
+ <tr style="border-top: 2px solid black">
173
+ <td>Text-guided Image Inpainting</td>
174
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/inpaint">Stable Diffusion Inpainting</a></td>
175
+ <td><a href="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-inpainting"> stable-diffusion-v1-5/stable-diffusion-inpainting </a></td>
176
+ </tr>
177
+ <tr style="border-top: 2px solid black">
178
+ <td>Image Variation</td>
179
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/image_variation">Stable Diffusion Image Variation</a></td>
180
+ <td><a href="https://huggingface.co/lambdalabs/sd-image-variations-diffusers"> lambdalabs/sd-image-variations-diffusers </a></td>
181
+ </tr>
182
+ <tr style="border-top: 2px solid black">
183
+ <td>Super Resolution</td>
184
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/upscale">Stable Diffusion Upscale</a></td>
185
+ <td><a href="https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler"> stabilityai/stable-diffusion-x4-upscaler </a></td>
186
+ </tr>
187
+ <tr>
188
+ <td>Super Resolution</td>
189
+ <td><a href="https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/latent_upscale">Stable Diffusion Latent Upscale</a></td>
190
+ <td><a href="https://huggingface.co/stabilityai/sd-x2-latent-upscaler"> stabilityai/sd-x2-latent-upscaler </a></td>
191
+ </tr>
192
+ </table>
193
+
194
+ ## Popular libraries using 🧨 Diffusers
195
+
196
+ - https://github.com/microsoft/TaskMatrix
197
+ - https://github.com/invoke-ai/InvokeAI
198
+ - https://github.com/InstantID/InstantID
199
+ - https://github.com/apple/ml-stable-diffusion
200
+ - https://github.com/Sanster/lama-cleaner
201
+ - https://github.com/IDEA-Research/Grounded-Segment-Anything
202
+ - https://github.com/ashawkey/stable-dreamfusion
203
+ - https://github.com/deep-floyd/IF
204
+ - https://github.com/bentoml/BentoML
205
+ - https://github.com/bmaltais/kohya_ss
206
+ - +14,000 other amazing GitHub repositories 💪
207
+
208
+ Thank you for using us ❤️.
209
+
210
+ ## Credits
211
+
212
+ 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:
213
+
214
+ - @CompVis' latent diffusion models library, available [here](https://github.com/CompVis/latent-diffusion)
215
+ - @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)
216
+ - @ermongroup's DDIM implementation, available [here](https://github.com/ermongroup/ddim)
217
+ - @yang-song's Score-VE and Score-VP implementations, available [here](https://github.com/yang-song/score_sde_pytorch)
218
+
219
+ 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.
220
+
221
+ ## Citation
222
+
223
+ ```bibtex
224
+ @misc{von-platen-etal-2022-diffusers,
225
+ 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},
226
+ title = {Diffusers: State-of-the-art diffusion models},
227
+ year = {2022},
228
+ publisher = {GitHub},
229
+ journal = {GitHub repository},
230
+ howpublished = {\url{https://github.com/huggingface/diffusers}}
231
+ }
232
+ ```
.slim/clonedeps/repos/huggingface__diffusers/SECURITY.md ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Security Policy
2
+
3
+ ## Supported Versions
4
+
5
+ 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.
6
+
7
+ ## Reporting a Vulnerability
8
+
9
+ Report privately — **do not open a public issue or PR for a suspected vulnerability.**
10
+
11
+ - **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.
12
+ - **Email:** security@huggingface.co
13
+
14
+ 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.
15
+
16
+ ## Recognition
17
+
18
+ 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).
19
+
20
+ ## What your report must include
21
+
22
+ 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.
23
+
24
+ ```markdown
25
+ ### Summary
26
+ One sentence: what the vulnerability is and where.
27
+
28
+ ### Affected version / commit
29
+ Exact released version or commit SHA you reproduced on (e.g. v4.57.0 / a1b2c3d).
30
+ Not "latest" or "main".
31
+
32
+ ### Affected component
33
+ The public API, module, or entry point involved (e.g. `DiffusionPipeline.from_pretrained`, `AutoModel.from_pretrained`).
34
+
35
+ ### Vulnerability class
36
+ Type and CWE if known (e.g. deserialization / CWE-502, path traversal / CWE-22).
37
+
38
+ ### Attack vector & preconditions
39
+ - How is the vulnerable code reached? (which API call / input / config)
40
+ - Who is the attacker and what do they control?
41
+ - What must be true for the attack to work? (auth, a user action, a non-default
42
+ setting, a malicious file being loaded, etc.)
43
+
44
+ ### Proof of concept
45
+ A minimal, self-contained script or step sequence that runs on a clean install
46
+ of the version above. Include:
47
+ - the exact commands / code to run,
48
+ - any input files needed (attach them, or give a script that generates them),
49
+ - the **expected** behavior vs. the **actual** behavior you observed.
50
+ A snippet showing that a function *exists* or *could* be misused is not a PoC.
51
+
52
+ ### Impact
53
+ What an attacker gains in a realistic deployment. "Could theoretically…"
54
+ without a working chain is not an impact.
55
+
56
+ ### Scope
57
+ Which trust boundary (see below) does this cross? If your finding touches
58
+ anything in the "Out of scope" list, name which item and explain why it is
59
+ nonetheless a violation of a guarantee we make.
60
+
61
+ ### Suggested severity (optional)
62
+ We assign the final severity. Include a CVSS v3.1 vector only if you have one.
63
+
64
+ ### Suggested fix (optional)
65
+ ```
66
+
67
+ > 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.
68
+
69
+ ## Threat model & trust boundaries
70
+
71
+ Understanding these saves everyone time — most of what we close as not-a-vulnerability falls inside one of them.
72
+
73
+ **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:
74
+ - preferring the `safetensors` format over pickle-based formats (`use_safetensors=True`),
75
+ - pinning a specific, reviewed `revision`,
76
+ - setting `trust_remote_code=True` only for repositories whose modeling files you have inspected.
77
+
78
+ 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.
79
+
80
+ **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.
81
+
82
+ ## In scope
83
+
84
+ 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:
85
+
86
+ - 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;
87
+ - a control we advertise being bypassed (e.g. code running despite `safetensors`-only loading, or a pinned revision being ignored);
88
+ - exposure or mishandling of credentials, tokens, or another user's data by the library;
89
+ - CI/CD or supply-chain issues in this repository.
90
+
91
+ ## Out of scope
92
+
93
+ 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.
94
+
95
+ - 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).
96
+ - 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.)
97
+ - Findings in documentation, tests, or other non-packaged reference material.
98
+ - 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.
99
+ - 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.
100
+ - Vulnerabilities in third-party dependencies we do not vendor — report upstream (we'll bump once fixed).
101
+ - Theoretical issues without a working proof of concept, and reports auto-generated from scanners or LLMs without a verified, reproducible chain.
102
+ - 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.
103
+
104
+ ## Safe harbor
105
+
106
+ 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.
.slim/clonedeps/repos/huggingface__diffusers/_typos.toml ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Files for typos
2
+ # Instruction: https://github.com/marketplace/actions/typos-action#getting-started
3
+
4
+ [default.extend-identifiers]
5
+
6
+ [default.extend-words]
7
+ NIN="NIN" # NIN is used in scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py
8
+ nd="np" # nd may be np (numpy)
9
+ parms="parms" # parms is used in scripts/convert_original_stable_diffusion_to_diffusers.py
10
+
11
+
12
+ [files]
13
+ extend-exclude = ["_typos.toml"]
.slim/clonedeps/repos/huggingface__diffusers/pyproject.toml ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [tool.ruff]
2
+ line-length = 119
3
+ extend-exclude = [
4
+ "src/diffusers/pipelines/flux2/system_messages.py",
5
+ ]
6
+
7
+ [tool.ruff.lint]
8
+ # Never enforce `E501` (line length violations).
9
+ ignore = ["C901", "E501", "E721", "E741", "F402", "F823"]
10
+ select = ["C", "E", "F", "I", "W"]
11
+
12
+ # Ignore import violations in all `__init__.py` files.
13
+ [tool.ruff.lint.per-file-ignores]
14
+ "__init__.py" = ["E402", "F401", "F403", "F811"]
15
+ "src/diffusers/utils/dummy_*.py" = ["F401"]
16
+
17
+ [tool.ruff.lint.isort]
18
+ lines-after-imports = 2
19
+ known-first-party = ["diffusers"]
20
+
21
+ [tool.ruff.format]
22
+ # Like Black, use double quotes for strings.
23
+ quote-style = "double"
24
+
25
+ # Like Black, indent with spaces, rather than tabs.
26
+ indent-style = "space"
27
+
28
+ # Like Black, respect magic trailing commas.
29
+ skip-magic-trailing-comma = false
30
+
31
+ # Like Black, automatically detect the appropriate line ending.
32
+ line-ending = "auto"
.slim/clonedeps/repos/huggingface__diffusers/setup.py ADDED
@@ -0,0 +1,305 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """
16
+ Simple check list from AllenNLP repo: https://github.com/allenai/allennlp/blob/main/setup.py
17
+
18
+ To create the package for PyPI.
19
+
20
+ 1. Run `make pre-release` (or `make pre-patch` for a patch release) then run `make fix-copies` to fix the index of the
21
+ documentation.
22
+
23
+ If releasing on a special branch, copy the updated README.md on the main branch for the commit you will make
24
+ for the post-release and run `make fix-copies` on the main branch as well.
25
+
26
+ 2. Unpin specific versions from setup.py that use a git install.
27
+
28
+ 3. Checkout the release branch (v<RELEASE>-release, for example v4.19-release), and commit these changes with the
29
+ message: "Release: <RELEASE>" and push.
30
+
31
+ 4. Manually trigger the "Nightly and release tests on main/release branch" workflow from the release branch. Wait for
32
+ the tests to complete. We can safely ignore the known test failures.
33
+
34
+ 5. Wait for the tests on main to be completed and be green (otherwise revert and fix bugs).
35
+
36
+ 6. Add a tag in git to mark the release: "git tag v<RELEASE> -m 'Adds tag v<RELEASE> for PyPI'"
37
+ Push the tag to git: git push --tags origin v<RELEASE>-release
38
+
39
+ 7. Build both the sources and the wheel. Do not change anything in setup.py between
40
+ creating the wheel and the source distribution (obviously).
41
+
42
+ For the wheel, run: "python setup.py bdist_wheel" in the top level directory
43
+ (This will build a wheel for the Python version you use to build it).
44
+
45
+ For the sources, run: "python setup.py sdist"
46
+ You should now have a /dist directory with both .whl and .tar.gz source versions.
47
+
48
+ Long story cut short, you need to run both before you can upload the distribution to the
49
+ test PyPI and the actual PyPI servers:
50
+
51
+ python setup.py bdist_wheel && python setup.py sdist
52
+
53
+ 8. Check that everything looks correct by uploading the package to the PyPI test server:
54
+
55
+ twine upload dist/* -r pypitest
56
+ (pypi suggests using twine as other methods upload files via plaintext.)
57
+ You may have to specify the repository url, use the following command then:
58
+ twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
59
+
60
+ Check that you can install it in a virtualenv by running:
61
+ pip install -i https://testpypi.python.org/pypi diffusers
62
+
63
+ If you are testing from a Colab Notebook, for instance, then do:
64
+ pip install diffusers && pip uninstall diffusers
65
+ pip install -i https://testpypi.python.org/pypi diffusers
66
+
67
+ Check you can run the following commands:
68
+ python -c "from diffusers import __version__; print(__version__)"
69
+ python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('fusing/unet-ldm-dummy-update'); pipe()"
70
+ python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe', safety_checker=None); pipe('ah suh du')"
71
+ python -c "from diffusers import *"
72
+
73
+ 9. Upload the final version to the actual PyPI:
74
+ twine upload dist/* -r pypi
75
+
76
+ 10. Prepare the release notes and publish them on GitHub once everything is looking hunky-dory. You can use the following
77
+ Space to fetch all the commits applicable for the release: https://huggingface.co/spacmes/sayakpaul/auto-release-notes-diffusers.
78
+ It automatically fetches the correct tag and branch but also provides the option to configure them.
79
+ `tag` should be the previous release tag (v0.26.1, for example), and `branch` should be
80
+ the latest release branch (v0.27.0-release, for example). It denotes all commits that have happened on branch
81
+ v0.27.0-release after the tag v0.26.1 was created.
82
+
83
+ 11. Run `make post-release` (or, for a patch release, `make post-patch`). If you were on a branch for the release,
84
+ you need to go back to main before executing this.
85
+ """
86
+
87
+ import re
88
+ import sys
89
+
90
+ from setuptools import Command, find_packages, setup
91
+
92
+
93
+ # IMPORTANT:
94
+ # 1. all dependencies should be listed here with their version requirements if any
95
+ # 2. once modified, run: `make deps_table_update` to update src/diffusers/dependency_versions_table.py
96
+ _deps = [
97
+ "Pillow", # keep the PIL.Image.Resampling deprecation away
98
+ "accelerate>=0.31.0",
99
+ "datasets",
100
+ "filelock",
101
+ "ftfy",
102
+ "hf-doc-builder>=0.3.0",
103
+ "httpx<1.0.0",
104
+ "huggingface-hub>=1.23.0,<2.0",
105
+ "requests-mock==1.10.0",
106
+ "importlib_metadata",
107
+ "invisible-watermark>=0.2.0",
108
+ "isort>=5.5.4",
109
+ "Jinja2",
110
+ "torchsde",
111
+ "note_seq",
112
+ "librosa",
113
+ # librosa pulls in numba/llvmlite; pin to versions that support Python 3.10+
114
+ "llvmlite>=0.40.0",
115
+ "numba>=0.57.0",
116
+ "numpy",
117
+ "parameterized",
118
+ "peft>=0.17.0",
119
+ "protobuf>=3.20.3,<4",
120
+ "pytest",
121
+ "pytest-timeout",
122
+ "pytest-xdist",
123
+ "python>=3.10.0",
124
+ "ruff==0.9.10",
125
+ "safetensors>=0.8.0",
126
+ "sentencepiece>=0.1.91,!=0.1.92",
127
+ "GitPython<3.1.19",
128
+ "scipy",
129
+ "onnx",
130
+ "optimum_quanto>=0.2.6",
131
+ "gguf>=0.10.0",
132
+ "auto-round>=0.13.0",
133
+ "torchao>=0.7.0",
134
+ "bitsandbytes>=0.43.3",
135
+ "nvidia_modelopt[hf]>=0.33.1",
136
+ "sdnq>=0.2.2",
137
+ "regex!=2019.12.17",
138
+ "requests",
139
+ "tensorboard",
140
+ "tiktoken>=0.7.0",
141
+ "torch>=2.6",
142
+ "torchvision",
143
+ "transformers>=4.41.2",
144
+ "urllib3<=2.0.0",
145
+ "black",
146
+ "phonemizer",
147
+ "opencv-python",
148
+ "timm",
149
+ "flashpack",
150
+ ]
151
+
152
+ # this is a lookup table with items like:
153
+ #
154
+ # tokenizers: "huggingface-hub==0.8.0"
155
+ # packaging: "packaging"
156
+ #
157
+ # some of the values are versioned whereas others aren't.
158
+ deps = {b: a for a, b in (re.findall(r"^(([^!=<>~]+)(?:[!=<>~].*)?$)", x)[0] for x in _deps)}
159
+
160
+ # since we save this data in src/diffusers/dependency_versions_table.py it can be easily accessed from
161
+ # anywhere. If you need to quickly access the data from this table in a shell, you can do so easily with:
162
+ #
163
+ # python -c 'import sys; from diffusers.dependency_versions_table import deps; \
164
+ # print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets
165
+ #
166
+ # Just pass the desired package names to that script as it's shown with 2 packages above.
167
+ #
168
+ # If diffusers is not yet installed and the work is done from the cloned repo remember to add `PYTHONPATH=src` to the script above
169
+ #
170
+ # You can then feed this for example to `pip`:
171
+ #
172
+ # pip install -U $(python -c 'import sys; from diffusers.dependency_versions_table import deps; \
173
+ # print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets)
174
+ #
175
+
176
+
177
+ def deps_list(*pkgs):
178
+ return [deps[pkg] for pkg in pkgs]
179
+
180
+
181
+ class DepsTableUpdateCommand(Command):
182
+ """
183
+ A custom command that updates the dependency table.
184
+ usage: python setup.py deps_table_update
185
+ """
186
+
187
+ description = "build runtime dependency table"
188
+ user_options = [
189
+ # format: (long option, short option, description).
190
+ (
191
+ "dep-table-update",
192
+ None,
193
+ "updates src/diffusers/dependency_versions_table.py",
194
+ ),
195
+ ]
196
+
197
+ def initialize_options(self):
198
+ pass
199
+
200
+ def finalize_options(self):
201
+ pass
202
+
203
+ def run(self):
204
+ entries = "\n".join([f' "{k}": "{v}",' for k, v in deps.items()])
205
+ content = [
206
+ "# THIS FILE HAS BEEN AUTOGENERATED. To update:",
207
+ "# 1. modify the `_deps` dict in setup.py",
208
+ "# 2. run `make deps_table_update`",
209
+ "deps = {",
210
+ entries,
211
+ "}",
212
+ "",
213
+ ]
214
+ target = "src/diffusers/dependency_versions_table.py"
215
+ print(f"updating {target}")
216
+ with open(target, "w", encoding="utf-8", newline="\n") as f:
217
+ f.write("\n".join(content))
218
+
219
+
220
+ extras = {}
221
+ extras["quality"] = deps_list("urllib3", "isort", "ruff", "hf-doc-builder")
222
+ extras["docs"] = deps_list("hf-doc-builder")
223
+ extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft", "timm")
224
+ extras["test"] = deps_list(
225
+ "ftfy",
226
+ "GitPython",
227
+ "datasets",
228
+ "Jinja2",
229
+ "invisible-watermark",
230
+ "librosa",
231
+ "llvmlite",
232
+ "numba",
233
+ "parameterized",
234
+ "protobuf",
235
+ "pytest",
236
+ "pytest-timeout",
237
+ "pytest-xdist",
238
+ "requests-mock",
239
+ "safetensors",
240
+ "sentencepiece",
241
+ "scipy",
242
+ "tiktoken",
243
+ "torchsde",
244
+ "torchvision",
245
+ "transformers",
246
+ "phonemizer",
247
+ )
248
+ extras["torch"] = deps_list("torch", "accelerate")
249
+
250
+ extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate")
251
+ extras["gguf"] = deps_list("gguf", "accelerate")
252
+ extras["optimum_quanto"] = deps_list("optimum_quanto", "accelerate")
253
+ extras["torchao"] = deps_list("torchao", "accelerate")
254
+ extras["nvidia_modelopt"] = deps_list("nvidia_modelopt[hf]")
255
+ extras["sdnq"] = deps_list("sdnq")
256
+ extras["flashpack"] = deps_list("flashpack")
257
+
258
+ extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"]
259
+
260
+ install_requires = [
261
+ deps["importlib_metadata"],
262
+ deps["filelock"],
263
+ deps["httpx"],
264
+ deps["huggingface-hub"],
265
+ deps["numpy"],
266
+ deps["regex"],
267
+ deps["requests"],
268
+ deps["safetensors"],
269
+ deps["Pillow"],
270
+ ]
271
+
272
+ version_range_max = max(sys.version_info[1], 10) + 1
273
+
274
+ setup(
275
+ name="diffusers",
276
+ 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)
277
+ description="State-of-the-art diffusion in PyTorch.",
278
+ long_description=open("README.md", "r", encoding="utf-8").read(),
279
+ long_description_content_type="text/markdown",
280
+ keywords="deep learning diffusion pytorch stable diffusion audioldm",
281
+ license="Apache 2.0 License",
282
+ author="The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)",
283
+ author_email="diffusers@huggingface.co",
284
+ url="https://github.com/huggingface/diffusers",
285
+ package_dir={"": "src"},
286
+ packages=find_packages("src"),
287
+ package_data={"diffusers": ["py.typed"]},
288
+ include_package_data=True,
289
+ python_requires=">=3.10.0",
290
+ install_requires=list(install_requires),
291
+ extras_require=extras,
292
+ entry_points={"console_scripts": ["diffusers-cli=diffusers.commands.diffusers_cli:main"]},
293
+ classifiers=[
294
+ "Development Status :: 5 - Production/Stable",
295
+ "Intended Audience :: Developers",
296
+ "Intended Audience :: Education",
297
+ "Intended Audience :: Science/Research",
298
+ "License :: OSI Approved :: Apache Software License",
299
+ "Operating System :: OS Independent",
300
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
301
+ "Programming Language :: Python :: 3",
302
+ ]
303
+ + [f"Programming Language :: Python :: 3.{i}" for i in range(8, version_range_max)],
304
+ cmdclass={"deps_table_update": DepsTableUpdateCommand},
305
+ )
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py ADDED
@@ -0,0 +1,1746 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ __version__ = "0.40.0.dev0"
2
+
3
+ from typing import TYPE_CHECKING
4
+
5
+ from .utils import (
6
+ DIFFUSERS_SLOW_IMPORT,
7
+ OptionalDependencyNotAvailable,
8
+ _LazyModule,
9
+ is_accelerate_available,
10
+ is_auto_round_available,
11
+ is_bitsandbytes_available,
12
+ is_gguf_available,
13
+ is_librosa_available,
14
+ is_note_seq_available,
15
+ is_nvidia_modelopt_available,
16
+ is_onnx_available,
17
+ is_opencv_available,
18
+ is_optimum_quanto_available,
19
+ is_scipy_available,
20
+ is_sdnq_available,
21
+ is_sentencepiece_available,
22
+ is_torch_available,
23
+ is_torchao_available,
24
+ is_torchsde_available,
25
+ is_transformers_available,
26
+ is_transformers_version,
27
+ )
28
+
29
+
30
+ # Lazy Import based on
31
+ # https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
32
+
33
+ # 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,
34
+ # and is used to defer the actual importing for when the objects are requested.
35
+ # This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends).
36
+
37
+ _import_structure = {
38
+ "configuration_utils": ["ConfigMixin"],
39
+ "guiders": [],
40
+ "hooks": [],
41
+ "loaders": ["FromOriginalModelMixin"],
42
+ "models": [],
43
+ "modular_pipelines": [],
44
+ "pipelines": [],
45
+ "quantizers.pipe_quant_config": ["PipelineQuantizationConfig"],
46
+ "quantizers.quantization_config": [],
47
+ "schedulers": [],
48
+ "utils": [
49
+ "OptionalDependencyNotAvailable",
50
+ "is_inflect_available",
51
+ "is_invisible_watermark_available",
52
+ "is_librosa_available",
53
+ "is_note_seq_available",
54
+ "is_onnx_available",
55
+ "is_scipy_available",
56
+ "is_torch_available",
57
+ "is_torchsde_available",
58
+ "is_transformers_available",
59
+ "is_transformers_version",
60
+ "is_unidecode_available",
61
+ "logging",
62
+ ],
63
+ }
64
+
65
+ try:
66
+ if not is_torch_available() and not is_accelerate_available() and not is_bitsandbytes_available():
67
+ raise OptionalDependencyNotAvailable()
68
+ except OptionalDependencyNotAvailable:
69
+ from .utils import dummy_bitsandbytes_objects
70
+
71
+ _import_structure["utils.dummy_bitsandbytes_objects"] = [
72
+ name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_")
73
+ ]
74
+ else:
75
+ _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig")
76
+
77
+ try:
78
+ if not is_torch_available() and not is_accelerate_available() and not is_gguf_available():
79
+ raise OptionalDependencyNotAvailable()
80
+ except OptionalDependencyNotAvailable:
81
+ from .utils import dummy_gguf_objects
82
+
83
+ _import_structure["utils.dummy_gguf_objects"] = [
84
+ name for name in dir(dummy_gguf_objects) if not name.startswith("_")
85
+ ]
86
+ else:
87
+ _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
88
+
89
+ try:
90
+ if not is_torch_available() and not is_accelerate_available() and not is_torchao_available():
91
+ raise OptionalDependencyNotAvailable()
92
+ except OptionalDependencyNotAvailable:
93
+ from .utils import dummy_torchao_objects
94
+
95
+ _import_structure["utils.dummy_torchao_objects"] = [
96
+ name for name in dir(dummy_torchao_objects) if not name.startswith("_")
97
+ ]
98
+ else:
99
+ _import_structure["quantizers.quantization_config"].append("TorchAoConfig")
100
+
101
+ try:
102
+ if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available():
103
+ raise OptionalDependencyNotAvailable()
104
+ except OptionalDependencyNotAvailable:
105
+ from .utils import dummy_optimum_quanto_objects
106
+
107
+ _import_structure["utils.dummy_optimum_quanto_objects"] = [
108
+ name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_")
109
+ ]
110
+ else:
111
+ _import_structure["quantizers.quantization_config"].append("QuantoConfig")
112
+
113
+ try:
114
+ if not is_torch_available() and not is_accelerate_available() and not is_nvidia_modelopt_available():
115
+ raise OptionalDependencyNotAvailable()
116
+ except OptionalDependencyNotAvailable:
117
+ from .utils import dummy_nvidia_modelopt_objects
118
+
119
+ _import_structure["utils.dummy_nvidia_modelopt_objects"] = [
120
+ name for name in dir(dummy_nvidia_modelopt_objects) if not name.startswith("_")
121
+ ]
122
+ else:
123
+ _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig")
124
+
125
+ try:
126
+ if not is_torch_available():
127
+ raise OptionalDependencyNotAvailable()
128
+ except OptionalDependencyNotAvailable:
129
+ from .utils import dummy_nunchaku_lite_objects
130
+
131
+ _import_structure["utils.dummy_nunchaku_lite_objects"] = [
132
+ name for name in dir(dummy_nunchaku_lite_objects) if not name.startswith("_")
133
+ ]
134
+ else:
135
+ _import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig")
136
+
137
+ try:
138
+ if not is_auto_round_available():
139
+ raise OptionalDependencyNotAvailable()
140
+ except OptionalDependencyNotAvailable:
141
+ from .utils import dummy_auto_round_objects
142
+
143
+ _import_structure["utils.dummy_auto_round_objects"] = [
144
+ name for name in dir(dummy_auto_round_objects) if not name.startswith("_")
145
+ ]
146
+ else:
147
+ _import_structure["quantizers.quantization_config"].append("AutoRoundConfig")
148
+
149
+ try:
150
+ if not is_torch_available() and not is_accelerate_available() and not is_sdnq_available():
151
+ raise OptionalDependencyNotAvailable()
152
+ except OptionalDependencyNotAvailable:
153
+ from .utils import dummy_sdnq_objects
154
+
155
+ _import_structure["utils.dummy_sdnq_objects"] = [
156
+ name for name in dir(dummy_sdnq_objects) if not name.startswith("_")
157
+ ]
158
+ else:
159
+ _import_structure["quantizers.quantization_config"].append("SDNQConfig")
160
+
161
+ try:
162
+ if not is_onnx_available():
163
+ raise OptionalDependencyNotAvailable()
164
+ except OptionalDependencyNotAvailable:
165
+ from .utils import dummy_onnx_objects # noqa F403
166
+
167
+ _import_structure["utils.dummy_onnx_objects"] = [
168
+ name for name in dir(dummy_onnx_objects) if not name.startswith("_")
169
+ ]
170
+
171
+ else:
172
+ _import_structure["pipelines"].extend(["OnnxRuntimeModel"])
173
+
174
+ try:
175
+ if not is_torch_available():
176
+ raise OptionalDependencyNotAvailable()
177
+ except OptionalDependencyNotAvailable:
178
+ from .utils import dummy_pt_objects # noqa F403
179
+
180
+ _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
181
+
182
+ else:
183
+ _import_structure["guiders"].extend(
184
+ [
185
+ "AdaptiveProjectedGuidance",
186
+ "AdaptiveProjectedMixGuidance",
187
+ "AutoGuidance",
188
+ "BaseGuidance",
189
+ "ClassifierFreeGuidance",
190
+ "ClassifierFreeZeroStarGuidance",
191
+ "FrequencyDecoupledGuidance",
192
+ "PerturbedAttentionGuidance",
193
+ "SkipLayerGuidance",
194
+ "SmoothedEnergyGuidance",
195
+ "TangentialClassifierFreeGuidance",
196
+ ]
197
+ )
198
+ _import_structure["hooks"].extend(
199
+ [
200
+ "FasterCacheConfig",
201
+ "FirstBlockCacheConfig",
202
+ "HookRegistry",
203
+ "LayerSkipConfig",
204
+ "MagCacheConfig",
205
+ "PyramidAttentionBroadcastConfig",
206
+ "SmoothedEnergyGuidanceConfig",
207
+ "TaylorSeerCacheConfig",
208
+ "TextKVCacheConfig",
209
+ "apply_faster_cache",
210
+ "apply_first_block_cache",
211
+ "apply_layer_skip",
212
+ "apply_mag_cache",
213
+ "apply_pyramid_attention_broadcast",
214
+ "apply_taylorseer_cache",
215
+ "apply_text_kv_cache",
216
+ ]
217
+ )
218
+ _import_structure["image_processor"] = [
219
+ "InpaintProcessor",
220
+ "IPAdapterMaskProcessor",
221
+ "PixArtImageProcessor",
222
+ "VaeImageProcessor",
223
+ "VaeImageProcessorLDM3D",
224
+ ]
225
+ _import_structure["models"].extend(
226
+ [
227
+ "AceStepTransformer1DModel",
228
+ "AllegroTransformer3DModel",
229
+ "AnimaTextConditioner",
230
+ "AnyFlowFARTransformer3DModel",
231
+ "AnyFlowTransformer3DModel",
232
+ "AsymmetricAutoencoderKL",
233
+ "AttentionBackendName",
234
+ "AuraFlowTransformer2DModel",
235
+ "AutoencoderDC",
236
+ "AutoencoderKL",
237
+ "AutoencoderKLAllegro",
238
+ "AutoencoderKLCogVideoX",
239
+ "AutoencoderKLCosmos",
240
+ "AutoencoderKLFlux2",
241
+ "AutoencoderKLHunyuanImage",
242
+ "AutoencoderKLHunyuanImageRefiner",
243
+ "AutoencoderKLHunyuanVideo",
244
+ "AutoencoderKLHunyuanVideo15",
245
+ "AutoencoderKLKVAE",
246
+ "AutoencoderKLKVAEVideo",
247
+ "AutoencoderKLLTX2Audio",
248
+ "AutoencoderKLLTX2Video",
249
+ "AutoencoderKLLTXVideo",
250
+ "AutoencoderKLMagvit",
251
+ "AutoencoderKLMiniMaxH3",
252
+ "AutoencoderKLMiniMaxH3Audio",
253
+ "AutoencoderKLMochi",
254
+ "AutoencoderKLQwenImage",
255
+ "AutoencoderKLTemporalDecoder",
256
+ "AutoencoderKLWan",
257
+ "AutoencoderOobleck",
258
+ "AutoencoderRAE",
259
+ "AutoencoderTiny",
260
+ "AutoencoderVidTok",
261
+ "AutoModel",
262
+ "BriaFiboTransformer2DModel",
263
+ "BriaTransformer2DModel",
264
+ "CacheMixin",
265
+ "ChromaTransformer2DModel",
266
+ "ChronoEditTransformer3DModel",
267
+ "CogVideoXTransformer3DModel",
268
+ "CogView3PlusTransformer2DModel",
269
+ "CogView4Transformer2DModel",
270
+ "ConsisIDTransformer3DModel",
271
+ "ConsistencyDecoderVAE",
272
+ "ContextParallelConfig",
273
+ "ControlNetModel",
274
+ "ControlNetUnionModel",
275
+ "ControlNetXSAdapter",
276
+ "Cosmos3AVAEAudioTokenizer",
277
+ "Cosmos3OmniTransformer",
278
+ "CosmosControlNetModel",
279
+ "CosmosTransformer3DModel",
280
+ "DiTTransformer2DModel",
281
+ "DreamLiteTransformer2DModel",
282
+ "DreamLiteUNetModel",
283
+ "EasyAnimateTransformer3DModel",
284
+ "ErnieImageTransformer2DModel",
285
+ "Flux2Transformer2DModel",
286
+ "FluxControlNetModel",
287
+ "FluxMultiControlNetModel",
288
+ "FluxTransformer2DModel",
289
+ "GlmImageTransformer2DModel",
290
+ "HeliosTransformer3DModel",
291
+ "HiDreamImageTransformer2DModel",
292
+ "HunyuanDiT2DControlNetModel",
293
+ "HunyuanDiT2DModel",
294
+ "HunyuanDiT2DMultiControlNetModel",
295
+ "HunyuanImageTransformer2DModel",
296
+ "HunyuanVideo15Transformer3DModel",
297
+ "HunyuanVideoFramepackTransformer3DModel",
298
+ "HunyuanVideoTransformer3DModel",
299
+ "I2VGenXLUNet",
300
+ "Ideogram4Transformer2DModel",
301
+ "JoyImageEditPlusTransformer3DModel",
302
+ "JoyImageEditTransformer3DModel",
303
+ "Kandinsky3UNet",
304
+ "Kandinsky5Transformer3DModel",
305
+ "Krea2Transformer2DModel",
306
+ "LatteTransformer3DModel",
307
+ "LongCatAudioDiTTransformer",
308
+ "LongCatAudioDiTVae",
309
+ "LongCatImageTransformer2DModel",
310
+ "LTX2VideoTransformer3DModel",
311
+ "LTXVideoTransformer3DModel",
312
+ "Lumina2Transformer2DModel",
313
+ "LuminaNextDiT2DModel",
314
+ "MiniMaxH3Transformer3DModel",
315
+ "MochiTransformer3DModel",
316
+ "ModelMixin",
317
+ "MotifVideoTransformer3DModel",
318
+ "MotionAdapter",
319
+ "MultiAdapter",
320
+ "MultiControlNetModel",
321
+ "NucleusMoEImageTransformer2DModel",
322
+ "OmniGenTransformer2DModel",
323
+ "OvisImageTransformer2DModel",
324
+ "ParallelConfig",
325
+ "PixArtTransformer2DModel",
326
+ "PriorTransformer",
327
+ "PRXTransformer2DModel",
328
+ "QwenImageControlNetModel",
329
+ "QwenImageMultiControlNetModel",
330
+ "QwenImageTransformer2DModel",
331
+ "SanaControlNetModel",
332
+ "SanaTransformer2DModel",
333
+ "SanaVideoTransformer3DModel",
334
+ "SD3ControlNetModel",
335
+ "SD3MultiControlNetModel",
336
+ "SD3Transformer2DModel",
337
+ "SkyReelsV2Transformer3DModel",
338
+ "SparseControlNetModel",
339
+ "StableAudioDiTModel",
340
+ "StableCascadeUNet",
341
+ "T2IAdapter",
342
+ "T5FilmDecoder",
343
+ "Transformer2DModel",
344
+ "TransformerTemporalModel",
345
+ "UNet1DModel",
346
+ "UNet2DConditionModel",
347
+ "UNet2DModel",
348
+ "UNet3DConditionModel",
349
+ "UNetControlNetXSModel",
350
+ "UNetMotionModel",
351
+ "UNetSpatioTemporalConditionModel",
352
+ "UVit2DModel",
353
+ "VQModel",
354
+ "WanAnimateTransformer3DModel",
355
+ "WanTransformer3DModel",
356
+ "WanVACETransformer3DModel",
357
+ "ZImageControlNetModel",
358
+ "ZImageTransformer2DModel",
359
+ "attention_backend",
360
+ ]
361
+ )
362
+ _import_structure["modular_pipelines"].extend(
363
+ [
364
+ "AutoPipelineBlocks",
365
+ "ComponentsManager",
366
+ "ComponentSpec",
367
+ "ConditionalPipelineBlocks",
368
+ "ConfigSpec",
369
+ "InputParam",
370
+ "LoopSequentialPipelineBlocks",
371
+ "ModularPipeline",
372
+ "ModularPipelineBlocks",
373
+ "OutputParam",
374
+ "SequentialPipelineBlocks",
375
+ ]
376
+ )
377
+ _import_structure["optimization"] = [
378
+ "get_constant_schedule",
379
+ "get_constant_schedule_with_warmup",
380
+ "get_cosine_schedule_with_warmup",
381
+ "get_cosine_with_hard_restarts_schedule_with_warmup",
382
+ "get_linear_schedule_with_warmup",
383
+ "get_polynomial_decay_schedule_with_warmup",
384
+ "get_scheduler",
385
+ ]
386
+ _import_structure["pipelines"].extend(
387
+ [
388
+ "AudioPipelineOutput",
389
+ "AutoPipelineForImage2Image",
390
+ "AutoPipelineForInpainting",
391
+ "AutoPipelineForText2Audio",
392
+ "AutoPipelineForText2Image",
393
+ "ConsistencyModelPipeline",
394
+ "DanceDiffusionPipeline",
395
+ "DDIMPipeline",
396
+ "DDPMPipeline",
397
+ "DiffusionPipeline",
398
+ "DiTPipeline",
399
+ "ImagePipelineOutput",
400
+ "KarrasVePipeline",
401
+ "LDMPipeline",
402
+ "LDMSuperResolutionPipeline",
403
+ "PNDMPipeline",
404
+ "RePaintPipeline",
405
+ "ScoreSdeVePipeline",
406
+ "StableDiffusionMixin",
407
+ ]
408
+ )
409
+ _import_structure["quantizers"] = ["DiffusersQuantizer"]
410
+ _import_structure["schedulers"].extend(
411
+ [
412
+ "AmusedScheduler",
413
+ "BlockRefinementScheduler",
414
+ "BlockRefinementSchedulerOutput",
415
+ "CMStochasticIterativeScheduler",
416
+ "CogVideoXDDIMScheduler",
417
+ "CogVideoXDPMScheduler",
418
+ "DDIMInverseScheduler",
419
+ "DDIMParallelScheduler",
420
+ "DDIMScheduler",
421
+ "DDPMParallelScheduler",
422
+ "DDPMScheduler",
423
+ "DDPMWuerstchenScheduler",
424
+ "DEISMultistepScheduler",
425
+ "DiscreteDDIMScheduler",
426
+ "DiscreteDDIMSchedulerOutput",
427
+ "DPMSolverMultistepInverseScheduler",
428
+ "DPMSolverMultistepScheduler",
429
+ "DPMSolverSinglestepScheduler",
430
+ "EDMDPMSolverMultistepScheduler",
431
+ "EDMEulerScheduler",
432
+ "EntropyBoundScheduler",
433
+ "EntropyBoundSchedulerOutput",
434
+ "EulerAncestralDiscreteScheduler",
435
+ "EulerDiscreteScheduler",
436
+ "FlowMapEulerDiscreteScheduler",
437
+ "FlowMatchEulerDiscreteScheduler",
438
+ "FlowMatchHeunDiscreteScheduler",
439
+ "FlowMatchLCMScheduler",
440
+ "HeliosDMDScheduler",
441
+ "HeliosScheduler",
442
+ "HeunDiscreteScheduler",
443
+ "IPNDMScheduler",
444
+ "KarrasVeScheduler",
445
+ "KDPM2AncestralDiscreteScheduler",
446
+ "KDPM2DiscreteScheduler",
447
+ "LCMScheduler",
448
+ "LTXEulerAncestralRFScheduler",
449
+ "MiniMaxH3Scheduler",
450
+ "PNDMScheduler",
451
+ "RePaintScheduler",
452
+ "SASolverScheduler",
453
+ "SchedulerMixin",
454
+ "SCMScheduler",
455
+ "ScoreSdeVeScheduler",
456
+ "TCDScheduler",
457
+ "UnCLIPScheduler",
458
+ "UniPCMultistepScheduler",
459
+ "VQDiffusionScheduler",
460
+ ]
461
+ )
462
+ _import_structure["training_utils"] = ["EMAModel"]
463
+ _import_structure["video_processor"] = ["VideoProcessor"]
464
+
465
+ try:
466
+ if not (is_torch_available() and is_scipy_available()):
467
+ raise OptionalDependencyNotAvailable()
468
+ except OptionalDependencyNotAvailable:
469
+ from .utils import dummy_torch_and_scipy_objects # noqa F403
470
+
471
+ _import_structure["utils.dummy_torch_and_scipy_objects"] = [
472
+ name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_")
473
+ ]
474
+
475
+ else:
476
+ _import_structure["schedulers"].extend(["LMSDiscreteScheduler"])
477
+
478
+ try:
479
+ if not (is_torch_available() and is_torchsde_available()):
480
+ raise OptionalDependencyNotAvailable()
481
+ except OptionalDependencyNotAvailable:
482
+ from .utils import dummy_torch_and_torchsde_objects # noqa F403
483
+
484
+ _import_structure["utils.dummy_torch_and_torchsde_objects"] = [
485
+ name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_")
486
+ ]
487
+
488
+ else:
489
+ _import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"])
490
+
491
+ try:
492
+ if not (is_torch_available() and is_transformers_available()):
493
+ raise OptionalDependencyNotAvailable()
494
+ except OptionalDependencyNotAvailable:
495
+ from .utils import dummy_torch_and_transformers_objects # noqa F403
496
+
497
+ _import_structure["utils.dummy_torch_and_transformers_objects"] = [
498
+ name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
499
+ ]
500
+
501
+ else:
502
+ _import_structure["modular_pipelines"].extend(
503
+ [
504
+ "AnimaAutoBlocks",
505
+ "AnimaModularPipeline",
506
+ "Cosmos3DistilledBlocks",
507
+ "Cosmos3DistilledModularPipeline",
508
+ "Cosmos3OmniBlocks",
509
+ "Cosmos3OmniModularPipeline",
510
+ "ErnieImageAutoBlocks",
511
+ "ErnieImageModularPipeline",
512
+ "Flux2AutoBlocks",
513
+ "Flux2KleinAutoBlocks",
514
+ "Flux2KleinBaseAutoBlocks",
515
+ "Flux2KleinBaseModularPipeline",
516
+ "Flux2KleinModularPipeline",
517
+ "Flux2ModularPipeline",
518
+ "FluxAutoBlocks",
519
+ "FluxKontextAutoBlocks",
520
+ "FluxKontextModularPipeline",
521
+ "FluxModularPipeline",
522
+ "HeliosAutoBlocks",
523
+ "HeliosModularPipeline",
524
+ "HeliosPyramidAutoBlocks",
525
+ "HeliosPyramidDistilledAutoBlocks",
526
+ "HeliosPyramidDistilledModularPipeline",
527
+ "HeliosPyramidModularPipeline",
528
+ "HunyuanVideo15AutoBlocks",
529
+ "HunyuanVideo15ModularPipeline",
530
+ "Ideogram4AutoBlocks",
531
+ "Ideogram4ModularPipeline",
532
+ "Krea2AutoBlocks",
533
+ "Krea2ModularPipeline",
534
+ "Krea2TurboAutoBlocks",
535
+ "Krea2TurboModularPipeline",
536
+ "LTXAutoBlocks",
537
+ "LTXModularPipeline",
538
+ "MiniMaxH3Blocks",
539
+ "MiniMaxH3ModularPipeline",
540
+ "QwenImageAutoBlocks",
541
+ "QwenImageEditAutoBlocks",
542
+ "QwenImageEditModularPipeline",
543
+ "QwenImageEditPlusAutoBlocks",
544
+ "QwenImageEditPlusModularPipeline",
545
+ "QwenImageLayeredAutoBlocks",
546
+ "QwenImageLayeredModularPipeline",
547
+ "QwenImageModularPipeline",
548
+ "StableDiffusion3AutoBlocks",
549
+ "StableDiffusion3ModularPipeline",
550
+ "StableDiffusionXLAutoBlocks",
551
+ "StableDiffusionXLModularPipeline",
552
+ "Wan22Blocks",
553
+ "Wan22Image2VideoBlocks",
554
+ "Wan22Image2VideoModularPipeline",
555
+ "Wan22ModularPipeline",
556
+ "WanBlocks",
557
+ "WanImage2VideoAutoBlocks",
558
+ "WanImage2VideoModularPipeline",
559
+ "WanModularPipeline",
560
+ "ZImageAutoBlocks",
561
+ "ZImageModularPipeline",
562
+ ]
563
+ )
564
+ _import_structure["pipelines"].extend(
565
+ [
566
+ "AceStepAudioTokenDetokenizer",
567
+ "AceStepAudioTokenizer",
568
+ "AceStepConditionEncoder",
569
+ "AceStepPipeline",
570
+ "AllegroPipeline",
571
+ "AltDiffusionImg2ImgPipeline",
572
+ "AltDiffusionPipeline",
573
+ "AmusedImg2ImgPipeline",
574
+ "AmusedInpaintPipeline",
575
+ "AmusedPipeline",
576
+ "AnimateDiffControlNetPipeline",
577
+ "AnimateDiffPAGPipeline",
578
+ "AnimateDiffPipeline",
579
+ "AnimateDiffSDXLPipeline",
580
+ "AnimateDiffSparseControlNetPipeline",
581
+ "AnimateDiffVideoToVideoControlNetPipeline",
582
+ "AnimateDiffVideoToVideoPipeline",
583
+ "AnyFlowFARPipeline",
584
+ "AnyFlowPipeline",
585
+ "AudioLDM2Pipeline",
586
+ "AudioLDM2ProjectionModel",
587
+ "AudioLDM2UNet2DConditionModel",
588
+ "AudioLDMPipeline",
589
+ "AuraFlowPipeline",
590
+ "BlipDiffusionControlNetPipeline",
591
+ "BlipDiffusionPipeline",
592
+ "BriaFiboEditPipeline",
593
+ "BriaFiboPipeline",
594
+ "BriaPipeline",
595
+ "ChromaImg2ImgPipeline",
596
+ "ChromaInpaintPipeline",
597
+ "ChromaPipeline",
598
+ "ChronoEditPipeline",
599
+ "CLIPImageProjection",
600
+ "CogVideoXFunControlPipeline",
601
+ "CogVideoXImageToVideoPipeline",
602
+ "CogVideoXPipeline",
603
+ "CogVideoXVideoToVideoPipeline",
604
+ "CogView3PlusPipeline",
605
+ "CogView4ControlPipeline",
606
+ "CogView4Pipeline",
607
+ "ConsisIDPipeline",
608
+ "Cosmos2_5_PredictBasePipeline",
609
+ "Cosmos2_5_TransferPipeline",
610
+ "Cosmos2TextToImagePipeline",
611
+ "Cosmos2VideoToWorldPipeline",
612
+ "Cosmos3OmniPipeline",
613
+ "CosmosActionCondition",
614
+ "CosmosTextToWorldPipeline",
615
+ "CosmosVideoToWorldPipeline",
616
+ "CycleDiffusionPipeline",
617
+ "DiffusionGemmaPipeline",
618
+ "DiffusionGemmaPipelineOutput",
619
+ "DreamLiteMobilePipeline",
620
+ "DreamLitePipeline",
621
+ "DreamLitePipelineOutput",
622
+ "EasyAnimateControlPipeline",
623
+ "EasyAnimateInpaintPipeline",
624
+ "EasyAnimatePipeline",
625
+ "ErnieImagePipeline",
626
+ "Flux2KleinInpaintPipeline",
627
+ "Flux2KleinKVPipeline",
628
+ "Flux2KleinPipeline",
629
+ "Flux2Pipeline",
630
+ "FluxControlImg2ImgPipeline",
631
+ "FluxControlInpaintPipeline",
632
+ "FluxControlNetImg2ImgPipeline",
633
+ "FluxControlNetInpaintPipeline",
634
+ "FluxControlNetPipeline",
635
+ "FluxControlPipeline",
636
+ "FluxFillPipeline",
637
+ "FluxImg2ImgPipeline",
638
+ "FluxInpaintPipeline",
639
+ "FluxKontextInpaintPipeline",
640
+ "FluxKontextPipeline",
641
+ "FluxPipeline",
642
+ "FluxPriorReduxPipeline",
643
+ "GlmImagePipeline",
644
+ "HeliosPipeline",
645
+ "HeliosPyramidPipeline",
646
+ "HiDreamImagePipeline",
647
+ "HunyuanDiTControlNetPipeline",
648
+ "HunyuanDiTPAGPipeline",
649
+ "HunyuanDiTPipeline",
650
+ "HunyuanImagePipeline",
651
+ "HunyuanImageRefinerPipeline",
652
+ "HunyuanSkyreelsImageToVideoPipeline",
653
+ "HunyuanVideo15ImageToVideoPipeline",
654
+ "HunyuanVideo15Pipeline",
655
+ "HunyuanVideoFramepackPipeline",
656
+ "HunyuanVideoImageToVideoPipeline",
657
+ "HunyuanVideoPipeline",
658
+ "I2VGenXLPipeline",
659
+ "Ideogram4Pipeline",
660
+ "Ideogram4PromptEnhancerHead",
661
+ "IFImg2ImgPipeline",
662
+ "IFImg2ImgSuperResolutionPipeline",
663
+ "IFInpaintingPipeline",
664
+ "IFInpaintingSuperResolutionPipeline",
665
+ "IFPipeline",
666
+ "IFSuperResolutionPipeline",
667
+ "ImageTextPipelineOutput",
668
+ "JoyImageEditPipeline",
669
+ "JoyImageEditPipelineOutput",
670
+ "JoyImageEditPlusPipeline",
671
+ "JoyImageEditPlusPipelineOutput",
672
+ "Kandinsky3Img2ImgPipeline",
673
+ "Kandinsky3Pipeline",
674
+ "Kandinsky5I2IPipeline",
675
+ "Kandinsky5I2VPipeline",
676
+ "Kandinsky5T2IPipeline",
677
+ "Kandinsky5T2VPipeline",
678
+ "KandinskyCombinedPipeline",
679
+ "KandinskyImg2ImgCombinedPipeline",
680
+ "KandinskyImg2ImgPipeline",
681
+ "KandinskyInpaintCombinedPipeline",
682
+ "KandinskyInpaintPipeline",
683
+ "KandinskyPipeline",
684
+ "KandinskyPriorPipeline",
685
+ "KandinskyV22CombinedPipeline",
686
+ "KandinskyV22ControlnetImg2ImgPipeline",
687
+ "KandinskyV22ControlnetPipeline",
688
+ "KandinskyV22Img2ImgCombinedPipeline",
689
+ "KandinskyV22Img2ImgPipeline",
690
+ "KandinskyV22InpaintCombinedPipeline",
691
+ "KandinskyV22InpaintPipeline",
692
+ "KandinskyV22Pipeline",
693
+ "KandinskyV22PriorEmb2EmbPipeline",
694
+ "KandinskyV22PriorPipeline",
695
+ "Krea2Pipeline",
696
+ "LatentConsistencyModelImg2ImgPipeline",
697
+ "LatentConsistencyModelPipeline",
698
+ "LattePipeline",
699
+ "LDMTextToImagePipeline",
700
+ "LEditsPPPipelineStableDiffusion",
701
+ "LEditsPPPipelineStableDiffusionXL",
702
+ "LLaDA2Pipeline",
703
+ "LLaDA2PipelineOutput",
704
+ "LongCatAudioDiTPipeline",
705
+ "LongCatImageEditPipeline",
706
+ "LongCatImagePipeline",
707
+ "LTX2ConditionPipeline",
708
+ "LTX2HDRPipeline",
709
+ "LTX2ImageToVideoPipeline",
710
+ "LTX2InContextPipeline",
711
+ "LTX2LatentUpsamplePipeline",
712
+ "LTX2Pipeline",
713
+ "LTXConditionPipeline",
714
+ "LTXI2VLongMultiPromptPipeline",
715
+ "LTXImageToVideoPipeline",
716
+ "LTXLatentUpsamplePipeline",
717
+ "LTXPipeline",
718
+ "LucyEditPipeline",
719
+ "Lumina2Pipeline",
720
+ "Lumina2Text2ImgPipeline",
721
+ "LuminaPipeline",
722
+ "LuminaText2ImgPipeline",
723
+ "MarigoldDepthPipeline",
724
+ "MarigoldIntrinsicsPipeline",
725
+ "MarigoldNormalsPipeline",
726
+ "MochiPipeline",
727
+ "MotifVideoImage2VideoPipeline",
728
+ "MotifVideoPipeline",
729
+ "MotifVideoPipelineOutput",
730
+ "MusicLDMPipeline",
731
+ "NucleusMoEImagePipeline",
732
+ "OmniGenPipeline",
733
+ "OvisImagePipeline",
734
+ "PaintByExamplePipeline",
735
+ "PIAPipeline",
736
+ "PixArtAlphaPipeline",
737
+ "PixArtSigmaPAGPipeline",
738
+ "PixArtSigmaPipeline",
739
+ "PRXPipeline",
740
+ "PRXPixelPipeline",
741
+ "QwenImageControlNetInpaintPipeline",
742
+ "QwenImageControlNetPipeline",
743
+ "QwenImageEditInpaintPipeline",
744
+ "QwenImageEditPipeline",
745
+ "QwenImageEditPlusPipeline",
746
+ "QwenImageImg2ImgPipeline",
747
+ "QwenImageInpaintPipeline",
748
+ "QwenImageLayeredPipeline",
749
+ "QwenImagePipeline",
750
+ "ReduxImageEncoder",
751
+ "SanaControlNetPipeline",
752
+ "SanaImageToVideoPipeline",
753
+ "SanaPAGPipeline",
754
+ "SanaPipeline",
755
+ "SanaSprintImg2ImgPipeline",
756
+ "SanaSprintPipeline",
757
+ "SanaVideoPipeline",
758
+ "SanaVideoPipeline",
759
+ "SemanticStableDiffusionPipeline",
760
+ "ShapEImg2ImgPipeline",
761
+ "ShapEPipeline",
762
+ "SkyReelsV2DiffusionForcingImageToVideoPipeline",
763
+ "SkyReelsV2DiffusionForcingPipeline",
764
+ "SkyReelsV2DiffusionForcingVideoToVideoPipeline",
765
+ "SkyReelsV2ImageToVideoPipeline",
766
+ "SkyReelsV2Pipeline",
767
+ "StableAudioPipeline",
768
+ "StableAudioProjectionModel",
769
+ "StableCascadeCombinedPipeline",
770
+ "StableCascadeDecoderPipeline",
771
+ "StableCascadePriorPipeline",
772
+ "StableDiffusion3ControlNetInpaintingPipeline",
773
+ "StableDiffusion3ControlNetPipeline",
774
+ "StableDiffusion3Img2ImgPipeline",
775
+ "StableDiffusion3InpaintPipeline",
776
+ "StableDiffusion3PAGImg2ImgPipeline",
777
+ "StableDiffusion3PAGImg2ImgPipeline",
778
+ "StableDiffusion3PAGPipeline",
779
+ "StableDiffusion3Pipeline",
780
+ "StableDiffusionAdapterPipeline",
781
+ "StableDiffusionAttendAndExcitePipeline",
782
+ "StableDiffusionControlNetImg2ImgPipeline",
783
+ "StableDiffusionControlNetInpaintPipeline",
784
+ "StableDiffusionControlNetPAGInpaintPipeline",
785
+ "StableDiffusionControlNetPAGPipeline",
786
+ "StableDiffusionControlNetPipeline",
787
+ "StableDiffusionControlNetXSPipeline",
788
+ "StableDiffusionDepth2ImgPipeline",
789
+ "StableDiffusionDiffEditPipeline",
790
+ "StableDiffusionGLIGENPipeline",
791
+ "StableDiffusionGLIGENTextImagePipeline",
792
+ "StableDiffusionImageVariationPipeline",
793
+ "StableDiffusionImg2ImgPipeline",
794
+ "StableDiffusionInpaintPipeline",
795
+ "StableDiffusionInpaintPipelineLegacy",
796
+ "StableDiffusionInstructPix2PixPipeline",
797
+ "StableDiffusionLatentUpscalePipeline",
798
+ "StableDiffusionLDM3DPipeline",
799
+ "StableDiffusionModelEditingPipeline",
800
+ "StableDiffusionPAGImg2ImgPipeline",
801
+ "StableDiffusionPAGInpaintPipeline",
802
+ "StableDiffusionPAGPipeline",
803
+ "StableDiffusionPanoramaPipeline",
804
+ "StableDiffusionParadigmsPipeline",
805
+ "StableDiffusionPipeline",
806
+ "StableDiffusionPipelineSafe",
807
+ "StableDiffusionPix2PixZeroPipeline",
808
+ "StableDiffusionSAGPipeline",
809
+ "StableDiffusionUpscalePipeline",
810
+ "StableDiffusionXLAdapterPipeline",
811
+ "StableDiffusionXLControlNetImg2ImgPipeline",
812
+ "StableDiffusionXLControlNetInpaintPipeline",
813
+ "StableDiffusionXLControlNetPAGImg2ImgPipeline",
814
+ "StableDiffusionXLControlNetPAGPipeline",
815
+ "StableDiffusionXLControlNetPipeline",
816
+ "StableDiffusionXLControlNetUnionImg2ImgPipeline",
817
+ "StableDiffusionXLControlNetUnionInpaintPipeline",
818
+ "StableDiffusionXLControlNetUnionPipeline",
819
+ "StableDiffusionXLControlNetXSPipeline",
820
+ "StableDiffusionXLImg2ImgPipeline",
821
+ "StableDiffusionXLInpaintPipeline",
822
+ "StableDiffusionXLInstructPix2PixPipeline",
823
+ "StableDiffusionXLPAGImg2ImgPipeline",
824
+ "StableDiffusionXLPAGInpaintPipeline",
825
+ "StableDiffusionXLPAGPipeline",
826
+ "StableDiffusionXLPipeline",
827
+ "StableUnCLIPImg2ImgPipeline",
828
+ "StableUnCLIPPipeline",
829
+ "StableVideoDiffusionPipeline",
830
+ "TextToVideoSDPipeline",
831
+ "TextToVideoZeroPipeline",
832
+ "TextToVideoZeroSDXLPipeline",
833
+ "UnCLIPImageVariationPipeline",
834
+ "UnCLIPPipeline",
835
+ "UniDiffuserModel",
836
+ "UniDiffuserPipeline",
837
+ "UniDiffuserTextDecoder",
838
+ "VersatileDiffusionDualGuidedPipeline",
839
+ "VersatileDiffusionImageVariationPipeline",
840
+ "VersatileDiffusionPipeline",
841
+ "VersatileDiffusionTextToImagePipeline",
842
+ "VideoToVideoSDPipeline",
843
+ "VisualClozeGenerationPipeline",
844
+ "VisualClozePipeline",
845
+ "VQDiffusionPipeline",
846
+ "WanAnimatePipeline",
847
+ "WanImageToVideoPipeline",
848
+ "WanPipeline",
849
+ "WanVACEPipeline",
850
+ "WanVideoToVideoPipeline",
851
+ "WuerstchenCombinedPipeline",
852
+ "WuerstchenDecoderPipeline",
853
+ "WuerstchenPriorPipeline",
854
+ "ZImageControlNetInpaintPipeline",
855
+ "ZImageControlNetPipeline",
856
+ "ZImageImg2ImgPipeline",
857
+ "ZImageInpaintPipeline",
858
+ "ZImageOmniPipeline",
859
+ "ZImagePipeline",
860
+ ]
861
+ )
862
+
863
+
864
+ try:
865
+ if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
866
+ raise OptionalDependencyNotAvailable()
867
+ except OptionalDependencyNotAvailable:
868
+ from .utils import dummy_torch_and_transformers_and_opencv_objects # noqa F403
869
+
870
+ _import_structure["utils.dummy_torch_and_transformers_and_opencv_objects"] = [
871
+ name for name in dir(dummy_torch_and_transformers_and_opencv_objects) if not name.startswith("_")
872
+ ]
873
+
874
+ else:
875
+ _import_structure["pipelines"].extend(["ConsisIDPipeline"])
876
+
877
+ try:
878
+ if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
879
+ raise OptionalDependencyNotAvailable()
880
+ except OptionalDependencyNotAvailable:
881
+ from .utils import dummy_torch_and_transformers_and_sentencepiece_objects # noqa F403
882
+
883
+ _import_structure["utils.dummy_torch_and_transformers_and_sentencepiece_objects"] = [
884
+ name for name in dir(dummy_torch_and_transformers_and_sentencepiece_objects) if not name.startswith("_")
885
+ ]
886
+
887
+ else:
888
+ _import_structure["pipelines"].extend(["KolorsImg2ImgPipeline", "KolorsPAGPipeline", "KolorsPipeline"])
889
+
890
+ try:
891
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
892
+ raise OptionalDependencyNotAvailable()
893
+ except OptionalDependencyNotAvailable:
894
+ from .utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403
895
+
896
+ _import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [
897
+ name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_")
898
+ ]
899
+
900
+ else:
901
+ _import_structure["pipelines"].extend(
902
+ [
903
+ "OnnxStableDiffusionImg2ImgPipeline",
904
+ "OnnxStableDiffusionInpaintPipeline",
905
+ "OnnxStableDiffusionInpaintPipelineLegacy",
906
+ "OnnxStableDiffusionPipeline",
907
+ "OnnxStableDiffusionUpscalePipeline",
908
+ "StableDiffusionOnnxPipeline",
909
+ ]
910
+ )
911
+
912
+ try:
913
+ if not (is_torch_available() and is_librosa_available()):
914
+ raise OptionalDependencyNotAvailable()
915
+ except OptionalDependencyNotAvailable:
916
+ from .utils import dummy_torch_and_librosa_objects # noqa F403
917
+
918
+ _import_structure["utils.dummy_torch_and_librosa_objects"] = [
919
+ name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_")
920
+ ]
921
+
922
+ else:
923
+ _import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"])
924
+
925
+ try:
926
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
927
+ raise OptionalDependencyNotAvailable()
928
+ except OptionalDependencyNotAvailable:
929
+ from .utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403
930
+
931
+ _import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [
932
+ name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_")
933
+ ]
934
+
935
+
936
+ else:
937
+ _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"])
938
+
939
+ try:
940
+ if not (is_note_seq_available()):
941
+ raise OptionalDependencyNotAvailable()
942
+ except OptionalDependencyNotAvailable:
943
+ from .utils import dummy_note_seq_objects # noqa F403
944
+
945
+ _import_structure["utils.dummy_note_seq_objects"] = [
946
+ name for name in dir(dummy_note_seq_objects) if not name.startswith("_")
947
+ ]
948
+
949
+
950
+ else:
951
+ _import_structure["pipelines"].extend(["MidiProcessor"])
952
+
953
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
954
+ from .configuration_utils import ConfigMixin
955
+ from .quantizers import PipelineQuantizationConfig
956
+
957
+ try:
958
+ if not is_bitsandbytes_available():
959
+ raise OptionalDependencyNotAvailable()
960
+ except OptionalDependencyNotAvailable:
961
+ from .utils.dummy_bitsandbytes_objects import *
962
+ else:
963
+ from .quantizers.quantization_config import BitsAndBytesConfig
964
+
965
+ try:
966
+ if not is_gguf_available():
967
+ raise OptionalDependencyNotAvailable()
968
+ except OptionalDependencyNotAvailable:
969
+ from .utils.dummy_gguf_objects import *
970
+ else:
971
+ from .quantizers.quantization_config import GGUFQuantizationConfig
972
+
973
+ try:
974
+ if not is_torchao_available():
975
+ raise OptionalDependencyNotAvailable()
976
+ except OptionalDependencyNotAvailable:
977
+ from .utils.dummy_torchao_objects import *
978
+ else:
979
+ from .quantizers.quantization_config import TorchAoConfig
980
+
981
+ try:
982
+ if not is_optimum_quanto_available():
983
+ raise OptionalDependencyNotAvailable()
984
+ except OptionalDependencyNotAvailable:
985
+ from .utils.dummy_optimum_quanto_objects import *
986
+ else:
987
+ from .quantizers.quantization_config import QuantoConfig
988
+
989
+ try:
990
+ if not is_nvidia_modelopt_available():
991
+ raise OptionalDependencyNotAvailable()
992
+ except OptionalDependencyNotAvailable:
993
+ from .utils.dummy_nvidia_modelopt_objects import *
994
+ else:
995
+ from .quantizers.quantization_config import NVIDIAModelOptConfig
996
+
997
+ try:
998
+ if not is_torch_available():
999
+ raise OptionalDependencyNotAvailable()
1000
+ except OptionalDependencyNotAvailable:
1001
+ from .utils.dummy_nunchaku_lite_objects import *
1002
+ else:
1003
+ from .quantizers.quantization_config import NunchakuLiteQuantizationConfig
1004
+
1005
+ try:
1006
+ if not is_auto_round_available():
1007
+ raise OptionalDependencyNotAvailable()
1008
+ except OptionalDependencyNotAvailable:
1009
+ from .utils.dummy_auto_round_objects import *
1010
+ else:
1011
+ from .quantizers.quantization_config import AutoRoundConfig
1012
+
1013
+ try:
1014
+ if not is_sdnq_available():
1015
+ raise OptionalDependencyNotAvailable()
1016
+ except OptionalDependencyNotAvailable:
1017
+ from .utils.dummy_sdnq_objects import *
1018
+ else:
1019
+ from .quantizers.quantization_config import SDNQConfig
1020
+
1021
+ try:
1022
+ if not is_onnx_available():
1023
+ raise OptionalDependencyNotAvailable()
1024
+ except OptionalDependencyNotAvailable:
1025
+ from .utils.dummy_onnx_objects import * # noqa F403
1026
+ else:
1027
+ from .pipelines import OnnxRuntimeModel
1028
+
1029
+ try:
1030
+ if not is_torch_available():
1031
+ raise OptionalDependencyNotAvailable()
1032
+ except OptionalDependencyNotAvailable:
1033
+ from .utils.dummy_pt_objects import * # noqa F403
1034
+ else:
1035
+ from .guiders import (
1036
+ AdaptiveProjectedGuidance,
1037
+ AdaptiveProjectedMixGuidance,
1038
+ AutoGuidance,
1039
+ BaseGuidance,
1040
+ ClassifierFreeGuidance,
1041
+ ClassifierFreeZeroStarGuidance,
1042
+ FrequencyDecoupledGuidance,
1043
+ PerturbedAttentionGuidance,
1044
+ SkipLayerGuidance,
1045
+ SmoothedEnergyGuidance,
1046
+ TangentialClassifierFreeGuidance,
1047
+ )
1048
+ from .hooks import (
1049
+ FasterCacheConfig,
1050
+ FirstBlockCacheConfig,
1051
+ HookRegistry,
1052
+ LayerSkipConfig,
1053
+ MagCacheConfig,
1054
+ PyramidAttentionBroadcastConfig,
1055
+ SmoothedEnergyGuidanceConfig,
1056
+ TaylorSeerCacheConfig,
1057
+ TextKVCacheConfig,
1058
+ apply_faster_cache,
1059
+ apply_first_block_cache,
1060
+ apply_layer_skip,
1061
+ apply_mag_cache,
1062
+ apply_pyramid_attention_broadcast,
1063
+ apply_taylorseer_cache,
1064
+ apply_text_kv_cache,
1065
+ )
1066
+ from .image_processor import (
1067
+ InpaintProcessor,
1068
+ IPAdapterMaskProcessor,
1069
+ PixArtImageProcessor,
1070
+ VaeImageProcessor,
1071
+ VaeImageProcessorLDM3D,
1072
+ )
1073
+ from .models import (
1074
+ AceStepTransformer1DModel,
1075
+ AllegroTransformer3DModel,
1076
+ AnimaTextConditioner,
1077
+ AnyFlowFARTransformer3DModel,
1078
+ AnyFlowTransformer3DModel,
1079
+ AsymmetricAutoencoderKL,
1080
+ AttentionBackendName,
1081
+ AuraFlowTransformer2DModel,
1082
+ AutoencoderDC,
1083
+ AutoencoderKL,
1084
+ AutoencoderKLAllegro,
1085
+ AutoencoderKLCogVideoX,
1086
+ AutoencoderKLCosmos,
1087
+ AutoencoderKLFlux2,
1088
+ AutoencoderKLHunyuanImage,
1089
+ AutoencoderKLHunyuanImageRefiner,
1090
+ AutoencoderKLHunyuanVideo,
1091
+ AutoencoderKLHunyuanVideo15,
1092
+ AutoencoderKLKVAE,
1093
+ AutoencoderKLKVAEVideo,
1094
+ AutoencoderKLLTX2Audio,
1095
+ AutoencoderKLLTX2Video,
1096
+ AutoencoderKLLTXVideo,
1097
+ AutoencoderKLMagvit,
1098
+ AutoencoderKLMiniMaxH3,
1099
+ AutoencoderKLMiniMaxH3Audio,
1100
+ AutoencoderKLMochi,
1101
+ AutoencoderKLQwenImage,
1102
+ AutoencoderKLTemporalDecoder,
1103
+ AutoencoderKLWan,
1104
+ AutoencoderOobleck,
1105
+ AutoencoderRAE,
1106
+ AutoencoderTiny,
1107
+ AutoencoderVidTok,
1108
+ AutoModel,
1109
+ BriaFiboTransformer2DModel,
1110
+ BriaTransformer2DModel,
1111
+ CacheMixin,
1112
+ ChromaTransformer2DModel,
1113
+ ChronoEditTransformer3DModel,
1114
+ CogVideoXTransformer3DModel,
1115
+ CogView3PlusTransformer2DModel,
1116
+ CogView4Transformer2DModel,
1117
+ ConsisIDTransformer3DModel,
1118
+ ConsistencyDecoderVAE,
1119
+ ContextParallelConfig,
1120
+ ControlNetModel,
1121
+ ControlNetUnionModel,
1122
+ ControlNetXSAdapter,
1123
+ Cosmos3AVAEAudioTokenizer,
1124
+ Cosmos3OmniTransformer,
1125
+ CosmosControlNetModel,
1126
+ CosmosTransformer3DModel,
1127
+ DiTTransformer2DModel,
1128
+ DreamLiteTransformer2DModel,
1129
+ DreamLiteUNetModel,
1130
+ EasyAnimateTransformer3DModel,
1131
+ ErnieImageTransformer2DModel,
1132
+ Flux2Transformer2DModel,
1133
+ FluxControlNetModel,
1134
+ FluxMultiControlNetModel,
1135
+ FluxTransformer2DModel,
1136
+ GlmImageTransformer2DModel,
1137
+ HeliosTransformer3DModel,
1138
+ HiDreamImageTransformer2DModel,
1139
+ HunyuanDiT2DControlNetModel,
1140
+ HunyuanDiT2DModel,
1141
+ HunyuanDiT2DMultiControlNetModel,
1142
+ HunyuanImageTransformer2DModel,
1143
+ HunyuanVideo15Transformer3DModel,
1144
+ HunyuanVideoFramepackTransformer3DModel,
1145
+ HunyuanVideoTransformer3DModel,
1146
+ I2VGenXLUNet,
1147
+ Ideogram4Transformer2DModel,
1148
+ JoyImageEditPlusTransformer3DModel,
1149
+ JoyImageEditTransformer3DModel,
1150
+ Kandinsky3UNet,
1151
+ Kandinsky5Transformer3DModel,
1152
+ Krea2Transformer2DModel,
1153
+ LatteTransformer3DModel,
1154
+ LongCatAudioDiTTransformer,
1155
+ LongCatAudioDiTVae,
1156
+ LongCatImageTransformer2DModel,
1157
+ LTX2VideoTransformer3DModel,
1158
+ LTXVideoTransformer3DModel,
1159
+ Lumina2Transformer2DModel,
1160
+ LuminaNextDiT2DModel,
1161
+ MiniMaxH3Transformer3DModel,
1162
+ MochiTransformer3DModel,
1163
+ ModelMixin,
1164
+ MotifVideoTransformer3DModel,
1165
+ MotionAdapter,
1166
+ MultiAdapter,
1167
+ MultiControlNetModel,
1168
+ NucleusMoEImageTransformer2DModel,
1169
+ OmniGenTransformer2DModel,
1170
+ OvisImageTransformer2DModel,
1171
+ ParallelConfig,
1172
+ PixArtTransformer2DModel,
1173
+ PriorTransformer,
1174
+ PRXTransformer2DModel,
1175
+ QwenImageControlNetModel,
1176
+ QwenImageMultiControlNetModel,
1177
+ QwenImageTransformer2DModel,
1178
+ SanaControlNetModel,
1179
+ SanaTransformer2DModel,
1180
+ SanaVideoTransformer3DModel,
1181
+ SD3ControlNetModel,
1182
+ SD3MultiControlNetModel,
1183
+ SD3Transformer2DModel,
1184
+ SkyReelsV2Transformer3DModel,
1185
+ SparseControlNetModel,
1186
+ StableAudioDiTModel,
1187
+ T2IAdapter,
1188
+ T5FilmDecoder,
1189
+ Transformer2DModel,
1190
+ TransformerTemporalModel,
1191
+ UNet1DModel,
1192
+ UNet2DConditionModel,
1193
+ UNet2DModel,
1194
+ UNet3DConditionModel,
1195
+ UNetControlNetXSModel,
1196
+ UNetMotionModel,
1197
+ UNetSpatioTemporalConditionModel,
1198
+ UVit2DModel,
1199
+ VQModel,
1200
+ WanAnimateTransformer3DModel,
1201
+ WanTransformer3DModel,
1202
+ WanVACETransformer3DModel,
1203
+ ZImageControlNetModel,
1204
+ ZImageTransformer2DModel,
1205
+ attention_backend,
1206
+ )
1207
+ from .modular_pipelines import (
1208
+ AutoPipelineBlocks,
1209
+ ComponentsManager,
1210
+ ComponentSpec,
1211
+ ConditionalPipelineBlocks,
1212
+ ConfigSpec,
1213
+ InputParam,
1214
+ LoopSequentialPipelineBlocks,
1215
+ ModularPipeline,
1216
+ ModularPipelineBlocks,
1217
+ OutputParam,
1218
+ SequentialPipelineBlocks,
1219
+ )
1220
+ from .optimization import (
1221
+ get_constant_schedule,
1222
+ get_constant_schedule_with_warmup,
1223
+ get_cosine_schedule_with_warmup,
1224
+ get_cosine_with_hard_restarts_schedule_with_warmup,
1225
+ get_linear_schedule_with_warmup,
1226
+ get_polynomial_decay_schedule_with_warmup,
1227
+ get_scheduler,
1228
+ )
1229
+ from .pipelines import (
1230
+ AudioPipelineOutput,
1231
+ AutoPipelineForImage2Image,
1232
+ AutoPipelineForInpainting,
1233
+ AutoPipelineForText2Audio,
1234
+ AutoPipelineForText2Image,
1235
+ BlipDiffusionControlNetPipeline,
1236
+ BlipDiffusionPipeline,
1237
+ CLIPImageProjection,
1238
+ ConsistencyModelPipeline,
1239
+ DanceDiffusionPipeline,
1240
+ DDIMPipeline,
1241
+ DDPMPipeline,
1242
+ DiffusionPipeline,
1243
+ DiTPipeline,
1244
+ ImagePipelineOutput,
1245
+ KarrasVePipeline,
1246
+ LDMPipeline,
1247
+ LDMSuperResolutionPipeline,
1248
+ PNDMPipeline,
1249
+ RePaintPipeline,
1250
+ ScoreSdeVePipeline,
1251
+ StableDiffusionMixin,
1252
+ )
1253
+ from .quantizers import DiffusersQuantizer
1254
+ from .schedulers import (
1255
+ AmusedScheduler,
1256
+ BlockRefinementScheduler,
1257
+ BlockRefinementSchedulerOutput,
1258
+ CMStochasticIterativeScheduler,
1259
+ CogVideoXDDIMScheduler,
1260
+ CogVideoXDPMScheduler,
1261
+ DDIMInverseScheduler,
1262
+ DDIMParallelScheduler,
1263
+ DDIMScheduler,
1264
+ DDPMParallelScheduler,
1265
+ DDPMScheduler,
1266
+ DDPMWuerstchenScheduler,
1267
+ DEISMultistepScheduler,
1268
+ DiscreteDDIMScheduler,
1269
+ DiscreteDDIMSchedulerOutput,
1270
+ DPMSolverMultistepInverseScheduler,
1271
+ DPMSolverMultistepScheduler,
1272
+ DPMSolverSinglestepScheduler,
1273
+ EDMDPMSolverMultistepScheduler,
1274
+ EDMEulerScheduler,
1275
+ EntropyBoundScheduler,
1276
+ EntropyBoundSchedulerOutput,
1277
+ EulerAncestralDiscreteScheduler,
1278
+ EulerDiscreteScheduler,
1279
+ FlowMapEulerDiscreteScheduler,
1280
+ FlowMatchEulerDiscreteScheduler,
1281
+ FlowMatchHeunDiscreteScheduler,
1282
+ FlowMatchLCMScheduler,
1283
+ HeliosDMDScheduler,
1284
+ HeliosScheduler,
1285
+ HeunDiscreteScheduler,
1286
+ IPNDMScheduler,
1287
+ KarrasVeScheduler,
1288
+ KDPM2AncestralDiscreteScheduler,
1289
+ KDPM2DiscreteScheduler,
1290
+ LCMScheduler,
1291
+ LTXEulerAncestralRFScheduler,
1292
+ MiniMaxH3Scheduler,
1293
+ PNDMScheduler,
1294
+ RePaintScheduler,
1295
+ SASolverScheduler,
1296
+ SchedulerMixin,
1297
+ SCMScheduler,
1298
+ ScoreSdeVeScheduler,
1299
+ TCDScheduler,
1300
+ UnCLIPScheduler,
1301
+ UniPCMultistepScheduler,
1302
+ VQDiffusionScheduler,
1303
+ )
1304
+ from .training_utils import EMAModel
1305
+ from .video_processor import VideoProcessor
1306
+
1307
+ try:
1308
+ if not (is_torch_available() and is_scipy_available()):
1309
+ raise OptionalDependencyNotAvailable()
1310
+ except OptionalDependencyNotAvailable:
1311
+ from .utils.dummy_torch_and_scipy_objects import * # noqa F403
1312
+ else:
1313
+ from .schedulers import LMSDiscreteScheduler
1314
+
1315
+ try:
1316
+ if not (is_torch_available() and is_torchsde_available()):
1317
+ raise OptionalDependencyNotAvailable()
1318
+ except OptionalDependencyNotAvailable:
1319
+ from .utils.dummy_torch_and_torchsde_objects import * # noqa F403
1320
+ else:
1321
+ from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler
1322
+
1323
+ try:
1324
+ if not (is_torch_available() and is_transformers_available()):
1325
+ raise OptionalDependencyNotAvailable()
1326
+ except OptionalDependencyNotAvailable:
1327
+ from .utils.dummy_torch_and_transformers_objects import * # noqa F403
1328
+ else:
1329
+ from .modular_pipelines import (
1330
+ AnimaAutoBlocks,
1331
+ AnimaModularPipeline,
1332
+ Cosmos3DistilledBlocks,
1333
+ Cosmos3DistilledModularPipeline,
1334
+ Cosmos3OmniBlocks,
1335
+ Cosmos3OmniModularPipeline,
1336
+ ErnieImageAutoBlocks,
1337
+ ErnieImageModularPipeline,
1338
+ Flux2AutoBlocks,
1339
+ Flux2KleinAutoBlocks,
1340
+ Flux2KleinBaseAutoBlocks,
1341
+ Flux2KleinBaseModularPipeline,
1342
+ Flux2KleinModularPipeline,
1343
+ Flux2ModularPipeline,
1344
+ FluxAutoBlocks,
1345
+ FluxKontextAutoBlocks,
1346
+ FluxKontextModularPipeline,
1347
+ FluxModularPipeline,
1348
+ HeliosAutoBlocks,
1349
+ HeliosModularPipeline,
1350
+ HeliosPyramidAutoBlocks,
1351
+ HeliosPyramidDistilledAutoBlocks,
1352
+ HeliosPyramidDistilledModularPipeline,
1353
+ HeliosPyramidModularPipeline,
1354
+ HunyuanVideo15AutoBlocks,
1355
+ HunyuanVideo15ModularPipeline,
1356
+ Ideogram4AutoBlocks,
1357
+ Ideogram4ModularPipeline,
1358
+ Krea2AutoBlocks,
1359
+ Krea2ModularPipeline,
1360
+ Krea2TurboAutoBlocks,
1361
+ Krea2TurboModularPipeline,
1362
+ LTXAutoBlocks,
1363
+ LTXModularPipeline,
1364
+ MiniMaxH3Blocks,
1365
+ MiniMaxH3ModularPipeline,
1366
+ QwenImageAutoBlocks,
1367
+ QwenImageEditAutoBlocks,
1368
+ QwenImageEditModularPipeline,
1369
+ QwenImageEditPlusAutoBlocks,
1370
+ QwenImageEditPlusModularPipeline,
1371
+ QwenImageLayeredAutoBlocks,
1372
+ QwenImageLayeredModularPipeline,
1373
+ QwenImageModularPipeline,
1374
+ StableDiffusion3AutoBlocks,
1375
+ StableDiffusion3ModularPipeline,
1376
+ StableDiffusionXLAutoBlocks,
1377
+ StableDiffusionXLModularPipeline,
1378
+ Wan22Blocks,
1379
+ Wan22Image2VideoBlocks,
1380
+ Wan22Image2VideoModularPipeline,
1381
+ Wan22ModularPipeline,
1382
+ WanBlocks,
1383
+ WanImage2VideoAutoBlocks,
1384
+ WanImage2VideoModularPipeline,
1385
+ WanModularPipeline,
1386
+ ZImageAutoBlocks,
1387
+ ZImageModularPipeline,
1388
+ )
1389
+ from .pipelines import (
1390
+ AceStepAudioTokenDetokenizer,
1391
+ AceStepAudioTokenizer,
1392
+ AceStepConditionEncoder,
1393
+ AceStepPipeline,
1394
+ AllegroPipeline,
1395
+ AltDiffusionImg2ImgPipeline,
1396
+ AltDiffusionPipeline,
1397
+ AmusedImg2ImgPipeline,
1398
+ AmusedInpaintPipeline,
1399
+ AmusedPipeline,
1400
+ AnimateDiffControlNetPipeline,
1401
+ AnimateDiffPAGPipeline,
1402
+ AnimateDiffPipeline,
1403
+ AnimateDiffSDXLPipeline,
1404
+ AnimateDiffSparseControlNetPipeline,
1405
+ AnimateDiffVideoToVideoControlNetPipeline,
1406
+ AnimateDiffVideoToVideoPipeline,
1407
+ AnyFlowFARPipeline,
1408
+ AnyFlowPipeline,
1409
+ AudioLDM2Pipeline,
1410
+ AudioLDM2ProjectionModel,
1411
+ AudioLDM2UNet2DConditionModel,
1412
+ AudioLDMPipeline,
1413
+ AuraFlowPipeline,
1414
+ BriaFiboEditPipeline,
1415
+ BriaFiboPipeline,
1416
+ BriaPipeline,
1417
+ ChromaImg2ImgPipeline,
1418
+ ChromaInpaintPipeline,
1419
+ ChromaPipeline,
1420
+ ChronoEditPipeline,
1421
+ CLIPImageProjection,
1422
+ CogVideoXFunControlPipeline,
1423
+ CogVideoXImageToVideoPipeline,
1424
+ CogVideoXPipeline,
1425
+ CogVideoXVideoToVideoPipeline,
1426
+ CogView3PlusPipeline,
1427
+ CogView4ControlPipeline,
1428
+ CogView4Pipeline,
1429
+ ConsisIDPipeline,
1430
+ Cosmos2_5_PredictBasePipeline,
1431
+ Cosmos2_5_TransferPipeline,
1432
+ Cosmos2TextToImagePipeline,
1433
+ Cosmos2VideoToWorldPipeline,
1434
+ Cosmos3OmniPipeline,
1435
+ CosmosActionCondition,
1436
+ CosmosTextToWorldPipeline,
1437
+ CosmosVideoToWorldPipeline,
1438
+ CycleDiffusionPipeline,
1439
+ DiffusionGemmaPipeline,
1440
+ DiffusionGemmaPipelineOutput,
1441
+ DreamLiteMobilePipeline,
1442
+ DreamLitePipeline,
1443
+ DreamLitePipelineOutput,
1444
+ EasyAnimateControlPipeline,
1445
+ EasyAnimateInpaintPipeline,
1446
+ EasyAnimatePipeline,
1447
+ ErnieImagePipeline,
1448
+ Flux2KleinInpaintPipeline,
1449
+ Flux2KleinKVPipeline,
1450
+ Flux2KleinPipeline,
1451
+ Flux2Pipeline,
1452
+ FluxControlImg2ImgPipeline,
1453
+ FluxControlInpaintPipeline,
1454
+ FluxControlNetImg2ImgPipeline,
1455
+ FluxControlNetInpaintPipeline,
1456
+ FluxControlNetPipeline,
1457
+ FluxControlPipeline,
1458
+ FluxFillPipeline,
1459
+ FluxImg2ImgPipeline,
1460
+ FluxInpaintPipeline,
1461
+ FluxKontextInpaintPipeline,
1462
+ FluxKontextPipeline,
1463
+ FluxPipeline,
1464
+ FluxPriorReduxPipeline,
1465
+ GlmImagePipeline,
1466
+ HeliosPipeline,
1467
+ HeliosPyramidPipeline,
1468
+ HiDreamImagePipeline,
1469
+ HunyuanDiTControlNetPipeline,
1470
+ HunyuanDiTPAGPipeline,
1471
+ HunyuanDiTPipeline,
1472
+ HunyuanImagePipeline,
1473
+ HunyuanImageRefinerPipeline,
1474
+ HunyuanSkyreelsImageToVideoPipeline,
1475
+ HunyuanVideo15ImageToVideoPipeline,
1476
+ HunyuanVideo15Pipeline,
1477
+ HunyuanVideoFramepackPipeline,
1478
+ HunyuanVideoImageToVideoPipeline,
1479
+ HunyuanVideoPipeline,
1480
+ I2VGenXLPipeline,
1481
+ Ideogram4Pipeline,
1482
+ Ideogram4PromptEnhancerHead,
1483
+ IFImg2ImgPipeline,
1484
+ IFImg2ImgSuperResolutionPipeline,
1485
+ IFInpaintingPipeline,
1486
+ IFInpaintingSuperResolutionPipeline,
1487
+ IFPipeline,
1488
+ IFSuperResolutionPipeline,
1489
+ ImageTextPipelineOutput,
1490
+ JoyImageEditPipeline,
1491
+ JoyImageEditPipelineOutput,
1492
+ JoyImageEditPlusPipeline,
1493
+ JoyImageEditPlusPipelineOutput,
1494
+ Kandinsky3Img2ImgPipeline,
1495
+ Kandinsky3Pipeline,
1496
+ Kandinsky5I2IPipeline,
1497
+ Kandinsky5I2VPipeline,
1498
+ Kandinsky5T2IPipeline,
1499
+ Kandinsky5T2VPipeline,
1500
+ KandinskyCombinedPipeline,
1501
+ KandinskyImg2ImgCombinedPipeline,
1502
+ KandinskyImg2ImgPipeline,
1503
+ KandinskyInpaintCombinedPipeline,
1504
+ KandinskyInpaintPipeline,
1505
+ KandinskyPipeline,
1506
+ KandinskyPriorPipeline,
1507
+ KandinskyV22CombinedPipeline,
1508
+ KandinskyV22ControlnetImg2ImgPipeline,
1509
+ KandinskyV22ControlnetPipeline,
1510
+ KandinskyV22Img2ImgCombinedPipeline,
1511
+ KandinskyV22Img2ImgPipeline,
1512
+ KandinskyV22InpaintCombinedPipeline,
1513
+ KandinskyV22InpaintPipeline,
1514
+ KandinskyV22Pipeline,
1515
+ KandinskyV22PriorEmb2EmbPipeline,
1516
+ KandinskyV22PriorPipeline,
1517
+ Krea2Pipeline,
1518
+ LatentConsistencyModelImg2ImgPipeline,
1519
+ LatentConsistencyModelPipeline,
1520
+ LattePipeline,
1521
+ LDMTextToImagePipeline,
1522
+ LEditsPPPipelineStableDiffusion,
1523
+ LEditsPPPipelineStableDiffusionXL,
1524
+ LLaDA2Pipeline,
1525
+ LLaDA2PipelineOutput,
1526
+ LongCatAudioDiTPipeline,
1527
+ LongCatImageEditPipeline,
1528
+ LongCatImagePipeline,
1529
+ LTX2ConditionPipeline,
1530
+ LTX2HDRPipeline,
1531
+ LTX2ImageToVideoPipeline,
1532
+ LTX2InContextPipeline,
1533
+ LTX2LatentUpsamplePipeline,
1534
+ LTX2Pipeline,
1535
+ LTXConditionPipeline,
1536
+ LTXI2VLongMultiPromptPipeline,
1537
+ LTXImageToVideoPipeline,
1538
+ LTXLatentUpsamplePipeline,
1539
+ LTXPipeline,
1540
+ LucyEditPipeline,
1541
+ Lumina2Pipeline,
1542
+ Lumina2Text2ImgPipeline,
1543
+ LuminaPipeline,
1544
+ LuminaText2ImgPipeline,
1545
+ MarigoldDepthPipeline,
1546
+ MarigoldIntrinsicsPipeline,
1547
+ MarigoldNormalsPipeline,
1548
+ MochiPipeline,
1549
+ MotifVideoImage2VideoPipeline,
1550
+ MotifVideoPipeline,
1551
+ MotifVideoPipelineOutput,
1552
+ MusicLDMPipeline,
1553
+ NucleusMoEImagePipeline,
1554
+ OmniGenPipeline,
1555
+ OvisImagePipeline,
1556
+ PaintByExamplePipeline,
1557
+ PIAPipeline,
1558
+ PixArtAlphaPipeline,
1559
+ PixArtSigmaPAGPipeline,
1560
+ PixArtSigmaPipeline,
1561
+ PRXPipeline,
1562
+ PRXPixelPipeline,
1563
+ QwenImageControlNetInpaintPipeline,
1564
+ QwenImageControlNetPipeline,
1565
+ QwenImageEditInpaintPipeline,
1566
+ QwenImageEditPipeline,
1567
+ QwenImageEditPlusPipeline,
1568
+ QwenImageImg2ImgPipeline,
1569
+ QwenImageInpaintPipeline,
1570
+ QwenImageLayeredPipeline,
1571
+ QwenImagePipeline,
1572
+ ReduxImageEncoder,
1573
+ SanaControlNetPipeline,
1574
+ SanaImageToVideoPipeline,
1575
+ SanaPAGPipeline,
1576
+ SanaPipeline,
1577
+ SanaSprintImg2ImgPipeline,
1578
+ SanaSprintPipeline,
1579
+ SanaVideoPipeline,
1580
+ SemanticStableDiffusionPipeline,
1581
+ ShapEImg2ImgPipeline,
1582
+ ShapEPipeline,
1583
+ SkyReelsV2DiffusionForcingImageToVideoPipeline,
1584
+ SkyReelsV2DiffusionForcingPipeline,
1585
+ SkyReelsV2DiffusionForcingVideoToVideoPipeline,
1586
+ SkyReelsV2ImageToVideoPipeline,
1587
+ SkyReelsV2Pipeline,
1588
+ StableAudioPipeline,
1589
+ StableAudioProjectionModel,
1590
+ StableCascadeCombinedPipeline,
1591
+ StableCascadeDecoderPipeline,
1592
+ StableCascadePriorPipeline,
1593
+ StableDiffusion3ControlNetInpaintingPipeline,
1594
+ StableDiffusion3ControlNetPipeline,
1595
+ StableDiffusion3Img2ImgPipeline,
1596
+ StableDiffusion3InpaintPipeline,
1597
+ StableDiffusion3PAGImg2ImgPipeline,
1598
+ StableDiffusion3PAGPipeline,
1599
+ StableDiffusion3Pipeline,
1600
+ StableDiffusionAdapterPipeline,
1601
+ StableDiffusionAttendAndExcitePipeline,
1602
+ StableDiffusionControlNetImg2ImgPipeline,
1603
+ StableDiffusionControlNetInpaintPipeline,
1604
+ StableDiffusionControlNetPAGInpaintPipeline,
1605
+ StableDiffusionControlNetPAGPipeline,
1606
+ StableDiffusionControlNetPipeline,
1607
+ StableDiffusionControlNetXSPipeline,
1608
+ StableDiffusionDepth2ImgPipeline,
1609
+ StableDiffusionDiffEditPipeline,
1610
+ StableDiffusionGLIGENPipeline,
1611
+ StableDiffusionGLIGENTextImagePipeline,
1612
+ StableDiffusionImageVariationPipeline,
1613
+ StableDiffusionImg2ImgPipeline,
1614
+ StableDiffusionInpaintPipeline,
1615
+ StableDiffusionInpaintPipelineLegacy,
1616
+ StableDiffusionInstructPix2PixPipeline,
1617
+ StableDiffusionLatentUpscalePipeline,
1618
+ StableDiffusionLDM3DPipeline,
1619
+ StableDiffusionModelEditingPipeline,
1620
+ StableDiffusionPAGImg2ImgPipeline,
1621
+ StableDiffusionPAGInpaintPipeline,
1622
+ StableDiffusionPAGPipeline,
1623
+ StableDiffusionPanoramaPipeline,
1624
+ StableDiffusionParadigmsPipeline,
1625
+ StableDiffusionPipeline,
1626
+ StableDiffusionPipelineSafe,
1627
+ StableDiffusionPix2PixZeroPipeline,
1628
+ StableDiffusionSAGPipeline,
1629
+ StableDiffusionUpscalePipeline,
1630
+ StableDiffusionXLAdapterPipeline,
1631
+ StableDiffusionXLControlNetImg2ImgPipeline,
1632
+ StableDiffusionXLControlNetInpaintPipeline,
1633
+ StableDiffusionXLControlNetPAGImg2ImgPipeline,
1634
+ StableDiffusionXLControlNetPAGPipeline,
1635
+ StableDiffusionXLControlNetPipeline,
1636
+ StableDiffusionXLControlNetUnionImg2ImgPipeline,
1637
+ StableDiffusionXLControlNetUnionInpaintPipeline,
1638
+ StableDiffusionXLControlNetUnionPipeline,
1639
+ StableDiffusionXLControlNetXSPipeline,
1640
+ StableDiffusionXLImg2ImgPipeline,
1641
+ StableDiffusionXLInpaintPipeline,
1642
+ StableDiffusionXLInstructPix2PixPipeline,
1643
+ StableDiffusionXLPAGImg2ImgPipeline,
1644
+ StableDiffusionXLPAGInpaintPipeline,
1645
+ StableDiffusionXLPAGPipeline,
1646
+ StableDiffusionXLPipeline,
1647
+ StableUnCLIPImg2ImgPipeline,
1648
+ StableUnCLIPPipeline,
1649
+ StableVideoDiffusionPipeline,
1650
+ TextToVideoSDPipeline,
1651
+ TextToVideoZeroPipeline,
1652
+ TextToVideoZeroSDXLPipeline,
1653
+ UnCLIPImageVariationPipeline,
1654
+ UnCLIPPipeline,
1655
+ UniDiffuserModel,
1656
+ UniDiffuserPipeline,
1657
+ UniDiffuserTextDecoder,
1658
+ VersatileDiffusionDualGuidedPipeline,
1659
+ VersatileDiffusionImageVariationPipeline,
1660
+ VersatileDiffusionPipeline,
1661
+ VersatileDiffusionTextToImagePipeline,
1662
+ VideoToVideoSDPipeline,
1663
+ VisualClozeGenerationPipeline,
1664
+ VisualClozePipeline,
1665
+ VQDiffusionPipeline,
1666
+ WanAnimatePipeline,
1667
+ WanImageToVideoPipeline,
1668
+ WanPipeline,
1669
+ WanVACEPipeline,
1670
+ WanVideoToVideoPipeline,
1671
+ WuerstchenCombinedPipeline,
1672
+ WuerstchenDecoderPipeline,
1673
+ WuerstchenPriorPipeline,
1674
+ ZImageControlNetInpaintPipeline,
1675
+ ZImageControlNetPipeline,
1676
+ ZImageImg2ImgPipeline,
1677
+ ZImageInpaintPipeline,
1678
+ ZImageOmniPipeline,
1679
+ ZImagePipeline,
1680
+ )
1681
+
1682
+ try:
1683
+ if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
1684
+ raise OptionalDependencyNotAvailable()
1685
+ except OptionalDependencyNotAvailable:
1686
+ from .utils.dummy_torch_and_transformers_and_sentencepiece_objects import * # noqa F403
1687
+ else:
1688
+ from .pipelines import KolorsImg2ImgPipeline, KolorsPAGPipeline, KolorsPipeline
1689
+
1690
+ try:
1691
+ if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
1692
+ raise OptionalDependencyNotAvailable()
1693
+ except OptionalDependencyNotAvailable:
1694
+ from .utils.dummy_torch_and_transformers_and_opencv_objects import * # noqa F403
1695
+ else:
1696
+ from .pipelines import ConsisIDPipeline
1697
+
1698
+ try:
1699
+ if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
1700
+ raise OptionalDependencyNotAvailable()
1701
+ except OptionalDependencyNotAvailable:
1702
+ from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403
1703
+ else:
1704
+ from .pipelines import (
1705
+ OnnxStableDiffusionImg2ImgPipeline,
1706
+ OnnxStableDiffusionInpaintPipeline,
1707
+ OnnxStableDiffusionInpaintPipelineLegacy,
1708
+ OnnxStableDiffusionPipeline,
1709
+ OnnxStableDiffusionUpscalePipeline,
1710
+ StableDiffusionOnnxPipeline,
1711
+ )
1712
+
1713
+ try:
1714
+ if not (is_torch_available() and is_librosa_available()):
1715
+ raise OptionalDependencyNotAvailable()
1716
+ except OptionalDependencyNotAvailable:
1717
+ from .utils.dummy_torch_and_librosa_objects import * # noqa F403
1718
+ else:
1719
+ from .pipelines import AudioDiffusionPipeline, Mel
1720
+
1721
+ try:
1722
+ if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
1723
+ raise OptionalDependencyNotAvailable()
1724
+ except OptionalDependencyNotAvailable:
1725
+ from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403
1726
+ else:
1727
+ from .pipelines import SpectrogramDiffusionPipeline
1728
+
1729
+ try:
1730
+ if not (is_note_seq_available()):
1731
+ raise OptionalDependencyNotAvailable()
1732
+ except OptionalDependencyNotAvailable:
1733
+ from .utils.dummy_note_seq_objects import * # noqa F403
1734
+ else:
1735
+ from .pipelines import MidiProcessor
1736
+
1737
+ else:
1738
+ import sys
1739
+
1740
+ sys.modules[__name__] = _LazyModule(
1741
+ __name__,
1742
+ globals()["__file__"],
1743
+ _import_structure,
1744
+ module_spec=__spec__,
1745
+ extra_objects={"__version__": __version__},
1746
+ )
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py ADDED
@@ -0,0 +1,244 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any
2
+
3
+ from .configuration_utils import ConfigMixin, register_to_config
4
+ from .utils import CONFIG_NAME
5
+
6
+
7
+ class PipelineCallback(ConfigMixin):
8
+ """
9
+ Base class for all the official callbacks used in a pipeline. This class provides a structure for implementing
10
+ custom callbacks and ensures that all callbacks have a consistent interface.
11
+
12
+ Please implement the following:
13
+ `tensor_inputs`: This should return a list of tensor inputs specific to your callback. You will only be able to
14
+ include
15
+ variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
16
+ `callback_fn`: This method defines the core functionality of your callback.
17
+ """
18
+
19
+ config_name = CONFIG_NAME
20
+
21
+ @register_to_config
22
+ def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None):
23
+ super().__init__()
24
+
25
+ if (cutoff_step_ratio is None and cutoff_step_index is None) or (
26
+ cutoff_step_ratio is not None and cutoff_step_index is not None
27
+ ):
28
+ raise ValueError("Either cutoff_step_ratio or cutoff_step_index should be provided, not both or none.")
29
+
30
+ if cutoff_step_ratio is not None and (
31
+ not isinstance(cutoff_step_ratio, float) or not (0.0 <= cutoff_step_ratio <= 1.0)
32
+ ):
33
+ raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.")
34
+
35
+ @property
36
+ def tensor_inputs(self) -> list[str]:
37
+ raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}")
38
+
39
+ def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> dict[str, Any]:
40
+ raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}")
41
+
42
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
43
+ return self.callback_fn(pipeline, step_index, timestep, callback_kwargs)
44
+
45
+
46
+ class MultiPipelineCallbacks:
47
+ """
48
+ This class is designed to handle multiple pipeline callbacks. It accepts a list of PipelineCallback objects and
49
+ provides a unified interface for calling all of them.
50
+ """
51
+
52
+ def __init__(self, callbacks: list[PipelineCallback]):
53
+ self.callbacks = callbacks
54
+
55
+ @property
56
+ def tensor_inputs(self) -> list[str]:
57
+ return [input for callback in self.callbacks for input in callback.tensor_inputs]
58
+
59
+ def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
60
+ """
61
+ Calls all the callbacks in order with the given arguments and returns the final callback_kwargs.
62
+ """
63
+ for callback in self.callbacks:
64
+ callback_kwargs = callback(pipeline, step_index, timestep, callback_kwargs)
65
+
66
+ return callback_kwargs
67
+
68
+
69
+ class SDCFGCutoffCallback(PipelineCallback):
70
+ """
71
+ Callback function for Stable Diffusion Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
72
+ `cutoff_step_index`), this callback will disable the CFG.
73
+
74
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
75
+ """
76
+
77
+ tensor_inputs = ["prompt_embeds"]
78
+
79
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
80
+ cutoff_step_ratio = self.config.cutoff_step_ratio
81
+ cutoff_step_index = self.config.cutoff_step_index
82
+
83
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
84
+ cutoff_step = (
85
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
86
+ )
87
+
88
+ if step_index == cutoff_step:
89
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
90
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
91
+
92
+ pipeline._guidance_scale = 0.0
93
+
94
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
95
+ return callback_kwargs
96
+
97
+
98
+ class SDXLCFGCutoffCallback(PipelineCallback):
99
+ """
100
+ Callback function for the base Stable Diffusion XL Pipelines. After certain number of steps (set by
101
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
102
+
103
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
104
+ """
105
+
106
+ tensor_inputs = [
107
+ "prompt_embeds",
108
+ "add_text_embeds",
109
+ "add_time_ids",
110
+ ]
111
+
112
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
113
+ cutoff_step_ratio = self.config.cutoff_step_ratio
114
+ cutoff_step_index = self.config.cutoff_step_index
115
+
116
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
117
+ cutoff_step = (
118
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
119
+ )
120
+
121
+ if step_index == cutoff_step:
122
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
123
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
124
+
125
+ add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
126
+ add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens
127
+
128
+ add_time_ids = callback_kwargs[self.tensor_inputs[2]]
129
+ add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector
130
+
131
+ pipeline._guidance_scale = 0.0
132
+
133
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
134
+ callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
135
+ callback_kwargs[self.tensor_inputs[2]] = add_time_ids
136
+
137
+ return callback_kwargs
138
+
139
+
140
+ class SDXLControlnetCFGCutoffCallback(PipelineCallback):
141
+ """
142
+ Callback function for the Controlnet Stable Diffusion XL Pipelines. After certain number of steps (set by
143
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
144
+
145
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
146
+ """
147
+
148
+ tensor_inputs = [
149
+ "prompt_embeds",
150
+ "add_text_embeds",
151
+ "add_time_ids",
152
+ "image",
153
+ ]
154
+
155
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
156
+ cutoff_step_ratio = self.config.cutoff_step_ratio
157
+ cutoff_step_index = self.config.cutoff_step_index
158
+
159
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
160
+ cutoff_step = (
161
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
162
+ )
163
+
164
+ if step_index == cutoff_step:
165
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
166
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
167
+
168
+ add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
169
+ add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens
170
+
171
+ add_time_ids = callback_kwargs[self.tensor_inputs[2]]
172
+ add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector
173
+
174
+ # For Controlnet
175
+ image = callback_kwargs[self.tensor_inputs[3]]
176
+ image = image[-1:]
177
+
178
+ pipeline._guidance_scale = 0.0
179
+
180
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
181
+ callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
182
+ callback_kwargs[self.tensor_inputs[2]] = add_time_ids
183
+ callback_kwargs[self.tensor_inputs[3]] = image
184
+
185
+ return callback_kwargs
186
+
187
+
188
+ class IPAdapterScaleCutoffCallback(PipelineCallback):
189
+ """
190
+ Callback function for any pipeline that inherits `IPAdapterMixin`. After certain number of steps (set by
191
+ `cutoff_step_ratio` or `cutoff_step_index`), this callback will set the IP Adapter scale to `0.0`.
192
+
193
+ Note: This callback mutates the IP Adapter attention processors by setting the scale to 0.0 after the cutoff step.
194
+ """
195
+
196
+ tensor_inputs = []
197
+
198
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
199
+ cutoff_step_ratio = self.config.cutoff_step_ratio
200
+ cutoff_step_index = self.config.cutoff_step_index
201
+
202
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
203
+ cutoff_step = (
204
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
205
+ )
206
+
207
+ if step_index == cutoff_step:
208
+ pipeline.set_ip_adapter_scale(0.0)
209
+ return callback_kwargs
210
+
211
+
212
+ class SD3CFGCutoffCallback(PipelineCallback):
213
+ """
214
+ Callback function for Stable Diffusion 3 Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
215
+ `cutoff_step_index`), this callback will disable the CFG.
216
+
217
+ Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
218
+ """
219
+
220
+ tensor_inputs = ["prompt_embeds", "pooled_prompt_embeds"]
221
+
222
+ def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
223
+ cutoff_step_ratio = self.config.cutoff_step_ratio
224
+ cutoff_step_index = self.config.cutoff_step_index
225
+
226
+ # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
227
+ cutoff_step = (
228
+ cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
229
+ )
230
+
231
+ if step_index == cutoff_step:
232
+ prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
233
+ prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
234
+
235
+ pooled_prompt_embeds = callback_kwargs[self.tensor_inputs[1]]
236
+ pooled_prompt_embeds = pooled_prompt_embeds[
237
+ -1:
238
+ ] # "-1" denotes the embeddings for conditional pooled text tokens.
239
+
240
+ pipeline._guidance_scale = 0.0
241
+
242
+ callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
243
+ callback_kwargs[self.tensor_inputs[1]] = pooled_prompt_embeds
244
+ return callback_kwargs
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py ADDED
@@ -0,0 +1,752 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2025 The HuggingFace Inc. team.
3
+ # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """ConfigMixin base class and utilities."""
17
+
18
+ import functools
19
+ import importlib
20
+ import inspect
21
+ import json
22
+ import os
23
+ import re
24
+ from collections import OrderedDict
25
+ from pathlib import Path
26
+ from typing import Any
27
+
28
+ import numpy as np
29
+ from huggingface_hub import DDUFEntry, create_repo, hf_hub_download
30
+ from huggingface_hub.utils import (
31
+ EntryNotFoundError,
32
+ HfHubHTTPError,
33
+ RepositoryNotFoundError,
34
+ RevisionNotFoundError,
35
+ validate_hf_hub_args,
36
+ )
37
+ from typing_extensions import Self
38
+
39
+ from . import __version__
40
+ from .utils import (
41
+ HUGGINGFACE_CO_RESOLVE_ENDPOINT,
42
+ DummyObject,
43
+ deprecate,
44
+ extract_commit_hash,
45
+ http_user_agent,
46
+ logging,
47
+ )
48
+
49
+
50
+ logger = logging.get_logger(__name__)
51
+
52
+ _re_configuration_file = re.compile(r"config\.(.*)\.json")
53
+
54
+
55
+ class FrozenDict(OrderedDict):
56
+ def __init__(self, *args, **kwargs):
57
+ super().__init__(*args, **kwargs)
58
+
59
+ for key, value in self.items():
60
+ setattr(self, key, value)
61
+
62
+ self.__frozen = True
63
+
64
+ def __delitem__(self, *args, **kwargs):
65
+ raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
66
+
67
+ def setdefault(self, *args, **kwargs):
68
+ raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
69
+
70
+ def pop(self, *args, **kwargs):
71
+ raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
72
+
73
+ def update(self, *args, **kwargs):
74
+ raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
75
+
76
+ def __setattr__(self, name, value):
77
+ if hasattr(self, "__frozen") and self.__frozen:
78
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
79
+ super().__setattr__(name, value)
80
+
81
+ def __setitem__(self, name, value):
82
+ if hasattr(self, "__frozen") and self.__frozen:
83
+ raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
84
+ super().__setitem__(name, value)
85
+
86
+
87
+ class ConfigMixin:
88
+ r"""
89
+ Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
90
+ provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
91
+ saving classes that inherit from [`ConfigMixin`].
92
+
93
+ Class attributes:
94
+ - **config_name** (`str`) -- A filename under which the config should stored when calling
95
+ [`~ConfigMixin.save_config`] (should be overridden by parent class).
96
+ - **ignore_for_config** (`list[str]`) -- A list of attributes that should not be saved in the config (should be
97
+ overridden by subclass).
98
+ - **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
99
+ - **_deprecated_kwargs** (`list[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
100
+ should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
101
+ subclass).
102
+ """
103
+
104
+ config_name = None
105
+ ignore_for_config = []
106
+ has_compatibles = False
107
+
108
+ _deprecated_kwargs = []
109
+ _auto_class = None
110
+
111
+ @classmethod
112
+ def register_for_auto_class(cls, auto_class="AutoModel"):
113
+ """
114
+ Register this class with the given auto class so that it can be loaded with `AutoModel.from_pretrained(...,
115
+ trust_remote_code=True)`.
116
+
117
+ When the config is saved, the resulting `config.json` will include an `auto_map` entry mapping the auto class
118
+ to this class's module and class name.
119
+
120
+ Args:
121
+ auto_class (`str` or type, *optional*, defaults to `"AutoModel"`):
122
+ The auto class to register this class with. Can be a string (e.g. `"AutoModel"`) or the class itself.
123
+ Currently only `"AutoModel"` is supported.
124
+
125
+ Example:
126
+
127
+ ```python
128
+ from diffusers import ModelMixin, ConfigMixin
129
+
130
+
131
+ class MyCustomModel(ModelMixin, ConfigMixin): ...
132
+
133
+
134
+ MyCustomModel.register_for_auto_class("AutoModel")
135
+ ```
136
+ """
137
+ if auto_class != "AutoModel":
138
+ raise ValueError(f"Only 'AutoModel' is supported, got '{auto_class}'.")
139
+
140
+ cls._auto_class = auto_class
141
+
142
+ def register_to_config(self, **kwargs):
143
+ if self.config_name is None:
144
+ raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
145
+ # Special case for `kwargs` used in deprecation warning added to schedulers
146
+ # TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
147
+ # or solve in a more general way.
148
+ kwargs.pop("kwargs", None)
149
+
150
+ if not hasattr(self, "_internal_dict"):
151
+ internal_dict = kwargs
152
+ else:
153
+ previous_dict = dict(self._internal_dict)
154
+ internal_dict = {**self._internal_dict, **kwargs}
155
+ logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
156
+
157
+ self._internal_dict = FrozenDict(internal_dict)
158
+
159
+ def __getattr__(self, name: str) -> Any:
160
+ """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
161
+ config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
162
+
163
+ This function is mostly copied from PyTorch's __getattr__ overwrite:
164
+ https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
165
+ """
166
+
167
+ is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
168
+ is_attribute = name in self.__dict__
169
+
170
+ if is_in_config and not is_attribute:
171
+ 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}'."
172
+ deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
173
+ return self._internal_dict[name]
174
+
175
+ raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
176
+
177
+ def save_config(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
178
+ """
179
+ Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
180
+ [`~ConfigMixin.from_config`] class method.
181
+
182
+ Args:
183
+ save_directory (`str` or `os.PathLike`):
184
+ Directory where the configuration JSON file is saved (will be created if it does not exist).
185
+ push_to_hub (`bool`, *optional*, defaults to `False`):
186
+ Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
187
+ repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
188
+ namespace).
189
+ kwargs (`dict[str, Any]`, *optional*):
190
+ Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
191
+ """
192
+ if os.path.isfile(save_directory):
193
+ raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
194
+
195
+ os.makedirs(save_directory, exist_ok=True)
196
+
197
+ # If we save using the predefined names, we can load using `from_config`
198
+ output_config_file = os.path.join(save_directory, self.config_name)
199
+
200
+ self.to_json_file(output_config_file)
201
+ logger.info(f"Configuration saved in {output_config_file}")
202
+
203
+ if push_to_hub:
204
+ commit_message = kwargs.pop("commit_message", None)
205
+ private = kwargs.pop("private", None)
206
+ create_pr = kwargs.pop("create_pr", False)
207
+ token = kwargs.pop("token", None)
208
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
209
+ repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
210
+ subfolder = kwargs.pop("subfolder", None)
211
+
212
+ self._upload_folder(
213
+ save_directory,
214
+ repo_id,
215
+ token=token,
216
+ commit_message=commit_message,
217
+ create_pr=create_pr,
218
+ subfolder=subfolder,
219
+ )
220
+
221
+ @classmethod
222
+ def from_config(
223
+ cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs
224
+ ) -> Self | tuple[Self, dict[str, Any]]:
225
+ r"""
226
+ Instantiate a Python class from a config dictionary.
227
+
228
+ Parameters:
229
+ config (`dict[str, Any]`):
230
+ A config dictionary from which the Python class is instantiated. Make sure to only load configuration
231
+ files of compatible classes.
232
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
233
+ Whether kwargs that are not consumed by the Python class should be returned or not.
234
+ kwargs (remaining dictionary of keyword arguments, *optional*):
235
+ Can be used to update the configuration object (after it is loaded) and initiate the Python class.
236
+ `**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
237
+ overwrite the same named arguments in `config`.
238
+
239
+ Returns:
240
+ [`ModelMixin`] or [`SchedulerMixin`]:
241
+ A model or scheduler object instantiated from a config dictionary.
242
+
243
+ Examples:
244
+
245
+ ```python
246
+ >>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
247
+
248
+ >>> # Download scheduler from huggingface.co and cache.
249
+ >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
250
+
251
+ >>> # Instantiate DDIM scheduler class with same config as DDPM
252
+ >>> scheduler = DDIMScheduler.from_config(scheduler.config)
253
+
254
+ >>> # Instantiate PNDM scheduler class with same config as DDPM
255
+ >>> scheduler = PNDMScheduler.from_config(scheduler.config)
256
+ ```
257
+ """
258
+ # <===== TO BE REMOVED WITH DEPRECATION
259
+ # TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
260
+ if "pretrained_model_name_or_path" in kwargs:
261
+ config = kwargs.pop("pretrained_model_name_or_path")
262
+
263
+ if config is None:
264
+ raise ValueError("Please make sure to provide a config as the first positional argument.")
265
+ # ======>
266
+
267
+ if not isinstance(config, dict):
268
+ deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
269
+ if "Scheduler" in cls.__name__:
270
+ deprecation_message += (
271
+ f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
272
+ " Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
273
+ " be removed in v1.0.0."
274
+ )
275
+ elif "Model" in cls.__name__:
276
+ deprecation_message += (
277
+ f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
278
+ f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
279
+ " instead. This functionality will be removed in v1.0.0."
280
+ )
281
+ deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
282
+ config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
283
+
284
+ init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
285
+
286
+ # Allow dtype to be specified on initialization
287
+ if "dtype" in unused_kwargs:
288
+ init_dict["dtype"] = unused_kwargs.pop("dtype")
289
+
290
+ # add possible deprecated kwargs
291
+ for deprecated_kwarg in cls._deprecated_kwargs:
292
+ if deprecated_kwarg in unused_kwargs:
293
+ init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
294
+
295
+ # Return model and optionally state and/or unused_kwargs
296
+ model = cls(**init_dict)
297
+
298
+ # make sure to also save config parameters that might be used for compatible classes
299
+ # update _class_name
300
+ if "_class_name" in hidden_dict:
301
+ hidden_dict["_class_name"] = cls.__name__
302
+
303
+ model.register_to_config(**hidden_dict)
304
+
305
+ # add hidden kwargs of compatible classes to unused_kwargs
306
+ unused_kwargs = {**unused_kwargs, **hidden_dict}
307
+
308
+ if return_unused_kwargs:
309
+ return (model, unused_kwargs)
310
+ else:
311
+ return model
312
+
313
+ @classmethod
314
+ def get_config_dict(cls, *args, **kwargs):
315
+ deprecation_message = (
316
+ f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
317
+ " removed in version v1.0.0"
318
+ )
319
+ deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
320
+ return cls.load_config(*args, **kwargs)
321
+
322
+ @classmethod
323
+ @validate_hf_hub_args
324
+ def load_config(
325
+ cls,
326
+ pretrained_model_name_or_path: str | os.PathLike,
327
+ return_unused_kwargs=False,
328
+ return_commit_hash=False,
329
+ **kwargs,
330
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
331
+ r"""
332
+ Load a model or scheduler configuration.
333
+
334
+ Parameters:
335
+ pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
336
+ Can be either:
337
+
338
+ - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
339
+ the Hub.
340
+ - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
341
+ [`~ConfigMixin.save_config`].
342
+
343
+ cache_dir (`str | os.PathLike`, *optional*):
344
+ Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
345
+ is not used.
346
+ force_download (`bool`, *optional*, defaults to `False`):
347
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
348
+ cached versions if they exist.
349
+ proxies (`dict[str, str]`, *optional*):
350
+ A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
351
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
352
+ output_loading_info(`bool`, *optional*, defaults to `False`):
353
+ Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
354
+ local_files_only (`bool`, *optional*, defaults to `False`):
355
+ Whether to only load local model weights and configuration files or not. If set to `True`, the model
356
+ won't be downloaded from the Hub.
357
+ token (`str` or *bool*, *optional*):
358
+ The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
359
+ `diffusers-cli login` (stored in `~/.huggingface`) is used.
360
+ revision (`str`, *optional*, defaults to `"main"`):
361
+ The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
362
+ allowed by Git.
363
+ subfolder (`str`, *optional*, defaults to `""`):
364
+ The subfolder location of a model file within a larger model repository on the Hub or locally.
365
+ return_unused_kwargs (`bool`, *optional*, defaults to `False):
366
+ Whether unused keyword arguments of the config are returned.
367
+ return_commit_hash (`bool`, *optional*, defaults to `False):
368
+ Whether the `commit_hash` of the loaded configuration are returned.
369
+
370
+ Returns:
371
+ `dict`:
372
+ A dictionary of all the parameters stored in a JSON configuration file.
373
+
374
+ """
375
+ cache_dir = kwargs.pop("cache_dir", None)
376
+ local_dir = kwargs.pop("local_dir", None)
377
+ local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
378
+ force_download = kwargs.pop("force_download", False)
379
+ proxies = kwargs.pop("proxies", None)
380
+ token = kwargs.pop("token", None)
381
+ local_files_only = kwargs.pop("local_files_only", False)
382
+ revision = kwargs.pop("revision", None)
383
+ _ = kwargs.pop("mirror", None)
384
+ subfolder = kwargs.pop("subfolder", None)
385
+ user_agent = kwargs.pop("user_agent", {})
386
+ dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
387
+
388
+ user_agent = {**user_agent, "file_type": "config"}
389
+ user_agent = http_user_agent(user_agent)
390
+
391
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
392
+
393
+ if cls.config_name is None:
394
+ raise ValueError(
395
+ "`self.config_name` is not defined. Note that one should not load a config from "
396
+ "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
397
+ )
398
+ # Custom path for now
399
+ if dduf_entries:
400
+ if subfolder is not None:
401
+ raise ValueError(
402
+ "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). "
403
+ "Please check the DDUF structure"
404
+ )
405
+ config_file = cls._get_config_file_from_dduf(pretrained_model_name_or_path, dduf_entries)
406
+ elif os.path.isfile(pretrained_model_name_or_path):
407
+ config_file = pretrained_model_name_or_path
408
+ elif os.path.isdir(pretrained_model_name_or_path):
409
+ if subfolder is not None and os.path.isfile(
410
+ os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
411
+ ):
412
+ config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
413
+ elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
414
+ # Load from a PyTorch checkpoint
415
+ config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
416
+ else:
417
+ raise EnvironmentError(
418
+ f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
419
+ )
420
+ else:
421
+ try:
422
+ # Load from URL or cache if already cached
423
+ config_file = hf_hub_download(
424
+ pretrained_model_name_or_path,
425
+ filename=cls.config_name,
426
+ cache_dir=cache_dir,
427
+ force_download=force_download,
428
+ proxies=proxies,
429
+ local_files_only=local_files_only,
430
+ token=token,
431
+ user_agent=user_agent,
432
+ subfolder=subfolder,
433
+ revision=revision,
434
+ local_dir=local_dir,
435
+ local_dir_use_symlinks=local_dir_use_symlinks,
436
+ )
437
+ except RepositoryNotFoundError:
438
+ raise EnvironmentError(
439
+ f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
440
+ " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
441
+ " token having permission to this repo with `token` or log in with `hf auth login`."
442
+ )
443
+ except RevisionNotFoundError:
444
+ raise EnvironmentError(
445
+ f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
446
+ " this model name. Check the model page at"
447
+ f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
448
+ )
449
+ except EntryNotFoundError:
450
+ raise EnvironmentError(
451
+ f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
452
+ )
453
+ except HfHubHTTPError as err:
454
+ raise EnvironmentError(
455
+ "There was a specific connection error when trying to load"
456
+ f" {pretrained_model_name_or_path}:\n{err}"
457
+ )
458
+ except ValueError:
459
+ raise EnvironmentError(
460
+ f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
461
+ f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
462
+ f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
463
+ " run the library in offline mode at"
464
+ " 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
465
+ )
466
+ except EnvironmentError:
467
+ raise EnvironmentError(
468
+ f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
469
+ "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
470
+ f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
471
+ f"containing a {cls.config_name} file"
472
+ )
473
+ try:
474
+ config_dict = cls._dict_from_json_file(config_file, dduf_entries=dduf_entries)
475
+
476
+ commit_hash = extract_commit_hash(config_file)
477
+ except (json.JSONDecodeError, UnicodeDecodeError):
478
+ raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.")
479
+
480
+ if not (return_unused_kwargs or return_commit_hash):
481
+ return config_dict
482
+
483
+ outputs = (config_dict,)
484
+
485
+ if return_unused_kwargs:
486
+ outputs += (kwargs,)
487
+
488
+ if return_commit_hash:
489
+ outputs += (commit_hash,)
490
+
491
+ return outputs
492
+
493
+ @staticmethod
494
+ def _get_init_keys(input_class):
495
+ return set(dict(inspect.signature(input_class.__init__).parameters).keys())
496
+
497
+ @classmethod
498
+ def extract_init_dict(cls, config_dict, **kwargs):
499
+ # Skip keys that were not present in the original config, so default __init__ values were used
500
+ used_defaults = config_dict.get("_use_default_values", [])
501
+ config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"}
502
+
503
+ # 0. Copy origin config dict
504
+ original_dict = dict(config_dict.items())
505
+
506
+ # 1. Retrieve expected config attributes from __init__ signature
507
+ expected_keys = cls._get_init_keys(cls)
508
+ expected_keys.remove("self")
509
+ # remove general kwargs if present in dict
510
+ if "kwargs" in expected_keys:
511
+ expected_keys.remove("kwargs")
512
+
513
+ # 2. Remove attributes that cannot be expected from expected config attributes
514
+ # remove keys to be ignored
515
+ if len(cls.ignore_for_config) > 0:
516
+ expected_keys = expected_keys - set(cls.ignore_for_config)
517
+
518
+ # load diffusers library to import compatible and original scheduler
519
+ diffusers_library = importlib.import_module(__name__.split(".")[0])
520
+
521
+ if cls.has_compatibles:
522
+ compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]
523
+ else:
524
+ compatible_classes = []
525
+
526
+ expected_keys_comp_cls = set()
527
+ for c in compatible_classes:
528
+ expected_keys_c = cls._get_init_keys(c)
529
+ expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c)
530
+ expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls)
531
+ config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls}
532
+
533
+ # remove attributes from orig class that cannot be expected
534
+ orig_cls_name = config_dict.pop("_class_name", cls.__name__)
535
+ if (
536
+ isinstance(orig_cls_name, str)
537
+ and orig_cls_name != cls.__name__
538
+ and hasattr(diffusers_library, orig_cls_name)
539
+ ):
540
+ orig_cls = getattr(diffusers_library, orig_cls_name)
541
+ unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys
542
+ config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig}
543
+ elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)):
544
+ raise ValueError(
545
+ "Make sure that the `_class_name` is of type string or list of string (for custom pipelines)."
546
+ )
547
+
548
+ # remove private attributes
549
+ config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")}
550
+
551
+ # remove quantization_config
552
+ config_dict = {k: v for k, v in config_dict.items() if k != "quantization_config"}
553
+
554
+ # 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments
555
+ init_dict = {}
556
+ for key in expected_keys:
557
+ # if config param is passed to kwarg and is present in config dict
558
+ # it should overwrite existing config dict key
559
+ if key in kwargs and key in config_dict:
560
+ config_dict[key] = kwargs.pop(key)
561
+
562
+ if key in kwargs:
563
+ # overwrite key
564
+ init_dict[key] = kwargs.pop(key)
565
+ elif key in config_dict:
566
+ # use value from config dict
567
+ init_dict[key] = config_dict.pop(key)
568
+
569
+ # 4. Give nice warning if unexpected values have been passed
570
+ if len(config_dict) > 0:
571
+ logger.warning(
572
+ f"The config attributes {config_dict} were passed to {cls.__name__}, "
573
+ "but are not expected and will be ignored. Please verify your "
574
+ f"{cls.config_name} configuration file."
575
+ )
576
+
577
+ # 5. Give nice info if config attributes are initialized to default because they have not been passed
578
+ passed_keys = set(init_dict.keys())
579
+ if len(expected_keys - passed_keys) > 0:
580
+ logger.info(
581
+ f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
582
+ )
583
+
584
+ # 6. Define unused keyword arguments
585
+ unused_kwargs = {**config_dict, **kwargs}
586
+
587
+ # 7. Define "hidden" config parameters that were saved for compatible classes
588
+ hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict}
589
+
590
+ return init_dict, unused_kwargs, hidden_config_dict
591
+
592
+ @classmethod
593
+ def _dict_from_json_file(cls, json_file: str | os.PathLike, dduf_entries: dict[str, DDUFEntry] | None = None):
594
+ if dduf_entries:
595
+ text = dduf_entries[json_file].read_text()
596
+ else:
597
+ with open(json_file, "r", encoding="utf-8") as reader:
598
+ text = reader.read()
599
+ return json.loads(text)
600
+
601
+ def __repr__(self):
602
+ return f"{self.__class__.__name__} {self.to_json_string()}"
603
+
604
+ @property
605
+ def config(self) -> dict[str, Any]:
606
+ """
607
+ Returns the config of the class as a frozen dictionary
608
+
609
+ Returns:
610
+ `dict[str, Any]`: Config of the class.
611
+ """
612
+ return self._internal_dict
613
+
614
+ def to_json_string(self) -> str:
615
+ """
616
+ Serializes the configuration instance to a JSON string.
617
+
618
+ Returns:
619
+ `str`:
620
+ String containing all the attributes that make up the configuration instance in JSON format.
621
+ """
622
+ config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
623
+ config_dict["_class_name"] = self.__class__.__name__
624
+ config_dict["_diffusers_version"] = __version__
625
+
626
+ def to_json_saveable(value):
627
+ if isinstance(value, np.ndarray):
628
+ value = value.tolist()
629
+ elif isinstance(value, Path):
630
+ value = value.as_posix()
631
+ elif hasattr(value, "to_dict") and callable(value.to_dict):
632
+ value = value.to_dict()
633
+ elif isinstance(value, list):
634
+ value = [to_json_saveable(v) for v in value]
635
+ return value
636
+
637
+ if "quantization_config" in config_dict:
638
+ config_dict["quantization_config"] = (
639
+ config_dict.quantization_config.to_dict()
640
+ if not isinstance(config_dict.quantization_config, dict)
641
+ else config_dict.quantization_config
642
+ )
643
+
644
+ config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
645
+ # Don't save "_ignore_files" or "_use_default_values"
646
+ config_dict.pop("_ignore_files", None)
647
+ config_dict.pop("_use_default_values", None)
648
+ # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable.
649
+ _ = config_dict.pop("_pre_quantization_dtype", None)
650
+
651
+ if getattr(self, "_auto_class", None) is not None:
652
+ module = self.__class__.__module__.split(".")[-1]
653
+ auto_map = config_dict.get("auto_map", {})
654
+ auto_map[self._auto_class] = f"{module}.{self.__class__.__name__}"
655
+ config_dict["auto_map"] = auto_map
656
+
657
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
658
+
659
+ def to_json_file(self, json_file_path: str | os.PathLike):
660
+ """
661
+ Save the configuration instance's parameters to a JSON file.
662
+
663
+ Args:
664
+ json_file_path (`str` or `os.PathLike`):
665
+ Path to the JSON file to save a configuration instance's parameters.
666
+ """
667
+ with open(json_file_path, "w", encoding="utf-8") as writer:
668
+ writer.write(self.to_json_string())
669
+
670
+ @classmethod
671
+ def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: dict[str, DDUFEntry]):
672
+ # paths inside a DDUF file must always be "/"
673
+ config_file = (
674
+ cls.config_name
675
+ if pretrained_model_name_or_path == ""
676
+ else "/".join([pretrained_model_name_or_path, cls.config_name])
677
+ )
678
+ if config_file not in dduf_entries:
679
+ raise ValueError(
680
+ f"We did not manage to find the file {config_file} in the dduf file. We only have the following files {dduf_entries.keys()}"
681
+ )
682
+ return config_file
683
+
684
+
685
+ def register_to_config(init):
686
+ r"""
687
+ Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are
688
+ automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that
689
+ shouldn't be registered in the config, use the `ignore_for_config` class variable
690
+
691
+ Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init!
692
+ """
693
+
694
+ @functools.wraps(init)
695
+ def inner_init(self, *args, **kwargs):
696
+ # Ignore private kwargs in the init.
697
+ init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
698
+ config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")}
699
+ if not isinstance(self, ConfigMixin):
700
+ raise RuntimeError(
701
+ f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
702
+ "not inherit from `ConfigMixin`."
703
+ )
704
+
705
+ ignore = getattr(self, "ignore_for_config", [])
706
+ # Get positional arguments aligned with kwargs
707
+ new_kwargs = {}
708
+ signature = inspect.signature(init)
709
+ parameters = {
710
+ name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore
711
+ }
712
+ for arg, name in zip(args, parameters.keys()):
713
+ new_kwargs[name] = arg
714
+
715
+ # Then add all kwargs
716
+ new_kwargs.update(
717
+ {
718
+ k: init_kwargs.get(k, default)
719
+ for k, default in parameters.items()
720
+ if k not in ignore and k not in new_kwargs
721
+ }
722
+ )
723
+
724
+ # Take note of the parameters that were not present in the loaded config
725
+ if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
726
+ new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
727
+
728
+ new_kwargs = {**config_init_kwargs, **new_kwargs}
729
+ getattr(self, "register_to_config")(**new_kwargs)
730
+ init(self, *args, **init_kwargs)
731
+
732
+ return inner_init
733
+
734
+
735
+ class LegacyConfigMixin(ConfigMixin):
736
+ r"""
737
+ A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
738
+ pipeline-specific classes (like `DiTTransformer2DModel`).
739
+ """
740
+
741
+ @classmethod
742
+ def from_config(cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs):
743
+ # To prevent dependency import problem.
744
+ from .models.model_loading_utils import _fetch_remapped_cls_from_config
745
+
746
+ # resolve remapping
747
+ remapped_class = _fetch_remapped_cls_from_config(config, cls)
748
+
749
+ if remapped_class is cls:
750
+ return super(LegacyConfigMixin, remapped_class).from_config(config, return_unused_kwargs, **kwargs)
751
+ else:
752
+ return remapped_class.from_config(config, return_unused_kwargs, **kwargs)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from .dependency_versions_table import deps
16
+ from .utils.versions import require_version, require_version_core
17
+
18
+
19
+ # define which module versions we always want to check at run time
20
+ # (usually the ones defined in `install_requires` in setup.py)
21
+ #
22
+ # order specific notes:
23
+ # - tqdm must be checked before tokenizers
24
+
25
+ pkgs_to_check_at_runtime = "python requests filelock numpy".split()
26
+ for pkg in pkgs_to_check_at_runtime:
27
+ if pkg in deps:
28
+ require_version_core(deps[pkg])
29
+ else:
30
+ raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
31
+
32
+
33
+ def dep_version_check(pkg, hint=None):
34
+ require_version(deps[pkg], hint)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # THIS FILE HAS BEEN AUTOGENERATED. To update:
2
+ # 1. modify the `_deps` dict in setup.py
3
+ # 2. run `make deps_table_update`
4
+ deps = {
5
+ "Pillow": "Pillow",
6
+ "accelerate": "accelerate>=0.31.0",
7
+ "datasets": "datasets",
8
+ "filelock": "filelock",
9
+ "ftfy": "ftfy",
10
+ "hf-doc-builder": "hf-doc-builder>=0.3.0",
11
+ "httpx": "httpx<1.0.0",
12
+ "huggingface-hub": "huggingface-hub>=1.23.0,<2.0",
13
+ "requests-mock": "requests-mock==1.10.0",
14
+ "importlib_metadata": "importlib_metadata",
15
+ "invisible-watermark": "invisible-watermark>=0.2.0",
16
+ "isort": "isort>=5.5.4",
17
+ "Jinja2": "Jinja2",
18
+ "torchsde": "torchsde",
19
+ "note_seq": "note_seq",
20
+ "librosa": "librosa",
21
+ "llvmlite": "llvmlite>=0.40.0",
22
+ "numba": "numba>=0.57.0",
23
+ "numpy": "numpy",
24
+ "parameterized": "parameterized",
25
+ "peft": "peft>=0.17.0",
26
+ "protobuf": "protobuf>=3.20.3,<4",
27
+ "pytest": "pytest",
28
+ "pytest-timeout": "pytest-timeout",
29
+ "pytest-xdist": "pytest-xdist",
30
+ "python": "python>=3.10.0",
31
+ "ruff": "ruff==0.9.10",
32
+ "safetensors": "safetensors>=0.8.0",
33
+ "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
34
+ "GitPython": "GitPython<3.1.19",
35
+ "scipy": "scipy",
36
+ "onnx": "onnx",
37
+ "optimum_quanto": "optimum_quanto>=0.2.6",
38
+ "gguf": "gguf>=0.10.0",
39
+ "auto-round": "auto-round>=0.13.0",
40
+ "torchao": "torchao>=0.7.0",
41
+ "bitsandbytes": "bitsandbytes>=0.43.3",
42
+ "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1",
43
+ "sdnq": "sdnq>=0.2.2",
44
+ "regex": "regex!=2019.12.17",
45
+ "requests": "requests",
46
+ "tensorboard": "tensorboard",
47
+ "tiktoken": "tiktoken>=0.7.0",
48
+ "torch": "torch>=2.6",
49
+ "torchvision": "torchvision",
50
+ "transformers": "transformers>=4.41.2",
51
+ "urllib3": "urllib3<=2.0.0",
52
+ "black": "black",
53
+ "phonemizer": "phonemizer",
54
+ "opencv-python": "opencv-python",
55
+ "timm": "timm",
56
+ "flashpack": "flashpack",
57
+ }
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py ADDED
@@ -0,0 +1,1468 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import math
16
+ import warnings
17
+
18
+ import numpy as np
19
+ import PIL.Image
20
+ import torch
21
+ import torch.nn.functional as F
22
+ from PIL import Image, ImageFilter, ImageOps
23
+
24
+ from .configuration_utils import ConfigMixin, register_to_config
25
+ from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate
26
+
27
+
28
+ PipelineImageInput = (
29
+ PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image] | list[np.ndarray] | list[torch.Tensor]
30
+ )
31
+
32
+ PipelineDepthInput = PipelineImageInput
33
+
34
+
35
+ def is_valid_image(image) -> bool:
36
+ r"""
37
+ Checks if the input is a valid image.
38
+
39
+ A valid image can be:
40
+ - A `PIL.Image.Image`.
41
+ - A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image).
42
+
43
+ Args:
44
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
45
+ The image to validate. It can be a PIL image, a NumPy array, or a torch tensor.
46
+
47
+ Returns:
48
+ `bool`:
49
+ `True` if the input is a valid image, `False` otherwise.
50
+ """
51
+ return isinstance(image, PIL.Image.Image) or isinstance(image, (np.ndarray, torch.Tensor)) and image.ndim in (2, 3)
52
+
53
+
54
+ def is_valid_image_imagelist(images):
55
+ r"""
56
+ Checks if the input is a valid image or list of images.
57
+
58
+ The input can be one of the following formats:
59
+ - A 4D tensor or numpy array (batch of images).
60
+ - A valid single image: `PIL.Image.Image`, 2D `np.ndarray` or `torch.Tensor` (grayscale image), 3D `np.ndarray` or
61
+ `torch.Tensor`.
62
+ - A list of valid images.
63
+
64
+ Args:
65
+ images (`np.ndarray | torch.Tensor | PIL.Image.Image | list`):
66
+ The image(s) to check. Can be a batch of images (4D tensor/array), a single image, or a list of valid
67
+ images.
68
+
69
+ Returns:
70
+ `bool`:
71
+ `True` if the input is valid, `False` otherwise.
72
+ """
73
+ if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
74
+ return True
75
+ elif is_valid_image(images):
76
+ return True
77
+ elif isinstance(images, list):
78
+ return all(is_valid_image(image) for image in images)
79
+ return False
80
+
81
+
82
+ class VaeImageProcessor(ConfigMixin):
83
+ """
84
+ Image processor for VAE.
85
+
86
+ Args:
87
+ do_resize (`bool`, *optional*, defaults to `True`):
88
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
89
+ `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
90
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
91
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
92
+ resample (`str`, *optional*, defaults to `lanczos`):
93
+ Resampling filter to use when resizing the image.
94
+ do_normalize (`bool`, *optional*, defaults to `True`):
95
+ Whether to normalize the image to [-1,1].
96
+ do_binarize (`bool`, *optional*, defaults to `False`):
97
+ Whether to binarize the image to 0/1.
98
+ do_convert_rgb (`bool`, *optional*, defaults to be `False`):
99
+ Whether to convert the images to RGB format.
100
+ do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
101
+ Whether to convert the images to grayscale format.
102
+ """
103
+
104
+ config_name = CONFIG_NAME
105
+
106
+ @register_to_config
107
+ def __init__(
108
+ self,
109
+ do_resize: bool = True,
110
+ vae_scale_factor: int = 8,
111
+ vae_latent_channels: int = 4,
112
+ resample: str = "lanczos",
113
+ reducing_gap: int | None = None,
114
+ do_normalize: bool = True,
115
+ do_binarize: bool = False,
116
+ do_convert_rgb: bool = False,
117
+ do_convert_grayscale: bool = False,
118
+ ):
119
+ super().__init__()
120
+ if do_convert_rgb and do_convert_grayscale:
121
+ raise ValueError(
122
+ "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`,"
123
+ " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.",
124
+ " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`",
125
+ )
126
+
127
+ @staticmethod
128
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
129
+ r"""
130
+ Convert a numpy image or a batch of images to a PIL image.
131
+
132
+ Args:
133
+ images (`np.ndarray`):
134
+ The image array to convert to PIL format.
135
+
136
+ Returns:
137
+ `list[PIL.Image.Image]`:
138
+ A list of PIL images.
139
+ """
140
+ if images.ndim == 3:
141
+ images = images[None, ...]
142
+ images = (images * 255).round().astype("uint8")
143
+ if images.shape[-1] == 1:
144
+ # special case for grayscale (single channel) images
145
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
146
+ else:
147
+ pil_images = [Image.fromarray(image) for image in images]
148
+
149
+ return pil_images
150
+
151
+ @staticmethod
152
+ def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
153
+ r"""
154
+ Convert a PIL image or a list of PIL images to NumPy arrays.
155
+
156
+ Args:
157
+ images (`PIL.Image.Image` or `list[PIL.Image.Image]`):
158
+ The PIL image or list of images to convert to NumPy format.
159
+
160
+ Returns:
161
+ `np.ndarray`:
162
+ A NumPy array representation of the images.
163
+ """
164
+ if not isinstance(images, list):
165
+ images = [images]
166
+ images = [np.array(image).astype(np.float32) / 255.0 for image in images]
167
+ images = np.stack(images, axis=0)
168
+
169
+ return images
170
+
171
+ @staticmethod
172
+ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
173
+ r"""
174
+ Convert a NumPy image to a PyTorch tensor.
175
+
176
+ Args:
177
+ images (`np.ndarray`):
178
+ The NumPy image array to convert to PyTorch format.
179
+
180
+ Returns:
181
+ `torch.Tensor`:
182
+ A PyTorch tensor representation of the images.
183
+ """
184
+ if images.ndim == 3:
185
+ images = images[..., None]
186
+
187
+ images = torch.from_numpy(images.transpose(0, 3, 1, 2))
188
+ return images
189
+
190
+ @staticmethod
191
+ def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
192
+ r"""
193
+ Convert a PyTorch tensor to a NumPy image.
194
+
195
+ Args:
196
+ images (`torch.Tensor`):
197
+ The PyTorch tensor to convert to NumPy format.
198
+
199
+ Returns:
200
+ `np.ndarray`:
201
+ A NumPy array representation of the images.
202
+ """
203
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
204
+ return images
205
+
206
+ @staticmethod
207
+ def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
208
+ r"""
209
+ Normalize an image array to [-1,1].
210
+
211
+ Args:
212
+ images (`np.ndarray` or `torch.Tensor`):
213
+ The image array to normalize.
214
+
215
+ Returns:
216
+ `np.ndarray` or `torch.Tensor`:
217
+ The normalized image array.
218
+ """
219
+ return 2.0 * images - 1.0
220
+
221
+ @staticmethod
222
+ def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
223
+ r"""
224
+ Denormalize an image array to [0,1].
225
+
226
+ Args:
227
+ images (`np.ndarray` or `torch.Tensor`):
228
+ The image array to denormalize.
229
+
230
+ Returns:
231
+ `np.ndarray` or `torch.Tensor`:
232
+ The denormalized image array.
233
+ """
234
+ return (images * 0.5 + 0.5).clamp(0, 1)
235
+
236
+ @staticmethod
237
+ def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:
238
+ r"""
239
+ Converts a PIL image to RGB format.
240
+
241
+ Args:
242
+ image (`PIL.Image.Image`):
243
+ The PIL image to convert to RGB.
244
+
245
+ Returns:
246
+ `PIL.Image.Image`:
247
+ The RGB-converted PIL image.
248
+ """
249
+ image = image.convert("RGB")
250
+
251
+ return image
252
+
253
+ @staticmethod
254
+ def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image:
255
+ r"""
256
+ Converts a given PIL image to grayscale.
257
+
258
+ Args:
259
+ image (`PIL.Image.Image`):
260
+ The input image to convert.
261
+
262
+ Returns:
263
+ `PIL.Image.Image`:
264
+ The image converted to grayscale.
265
+ """
266
+ image = image.convert("L")
267
+
268
+ return image
269
+
270
+ @staticmethod
271
+ def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image:
272
+ r"""
273
+ Applies Gaussian blur to an image.
274
+
275
+ Args:
276
+ image (`PIL.Image.Image`):
277
+ The PIL image to convert to grayscale.
278
+
279
+ Returns:
280
+ `PIL.Image.Image`:
281
+ The grayscale-converted PIL image.
282
+ """
283
+ image = image.filter(ImageFilter.GaussianBlur(blur_factor))
284
+
285
+ return image
286
+
287
+ @staticmethod
288
+ def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0):
289
+ r"""
290
+ Finds a rectangular region that contains all masked ares in an image, and expands region to match the aspect
291
+ ratio of the original image; for example, if user drew mask in a 128x32 region, and the dimensions for
292
+ processing are 512x512, the region will be expanded to 128x128.
293
+
294
+ Args:
295
+ mask_image (PIL.Image.Image): Mask image.
296
+ width (int): Width of the image to be processed.
297
+ height (int): Height of the image to be processed.
298
+ pad (int, optional): Padding to be added to the crop region. Defaults to 0.
299
+
300
+ Returns:
301
+ tuple: (x1, y1, x2, y2) represent a rectangular region that contains all masked ares in an image and
302
+ matches the original aspect ratio.
303
+ """
304
+
305
+ mask_image = mask_image.convert("L")
306
+ mask = np.array(mask_image)
307
+
308
+ # 1. find a rectangular region that contains all masked ares in an image
309
+ h, w = mask.shape
310
+ crop_left = 0
311
+ for i in range(w):
312
+ if not (mask[:, i] == 0).all():
313
+ break
314
+ crop_left += 1
315
+
316
+ crop_right = 0
317
+ for i in reversed(range(w)):
318
+ if not (mask[:, i] == 0).all():
319
+ break
320
+ crop_right += 1
321
+
322
+ crop_top = 0
323
+ for i in range(h):
324
+ if not (mask[i] == 0).all():
325
+ break
326
+ crop_top += 1
327
+
328
+ crop_bottom = 0
329
+ for i in reversed(range(h)):
330
+ if not (mask[i] == 0).all():
331
+ break
332
+ crop_bottom += 1
333
+
334
+ # 2. add padding to the crop region
335
+ x1, y1, x2, y2 = (
336
+ int(max(crop_left - pad, 0)),
337
+ int(max(crop_top - pad, 0)),
338
+ int(min(w - crop_right + pad, w)),
339
+ int(min(h - crop_bottom + pad, h)),
340
+ )
341
+
342
+ # 3. expands crop region to match the aspect ratio of the image to be processed
343
+ ratio_crop_region = (x2 - x1) / (y2 - y1)
344
+ ratio_processing = width / height
345
+
346
+ if ratio_crop_region > ratio_processing:
347
+ desired_height = (x2 - x1) / ratio_processing
348
+ desired_height_diff = int(desired_height - (y2 - y1))
349
+ y1 -= desired_height_diff // 2
350
+ y2 += desired_height_diff - desired_height_diff // 2
351
+ if y2 >= mask_image.height:
352
+ diff = y2 - mask_image.height
353
+ y2 -= diff
354
+ y1 -= diff
355
+ if y1 < 0:
356
+ y2 -= y1
357
+ y1 -= y1
358
+ if y2 >= mask_image.height:
359
+ y2 = mask_image.height
360
+ else:
361
+ desired_width = (y2 - y1) * ratio_processing
362
+ desired_width_diff = int(desired_width - (x2 - x1))
363
+ x1 -= desired_width_diff // 2
364
+ x2 += desired_width_diff - desired_width_diff // 2
365
+ if x2 >= mask_image.width:
366
+ diff = x2 - mask_image.width
367
+ x2 -= diff
368
+ x1 -= diff
369
+ if x1 < 0:
370
+ x2 -= x1
371
+ x1 -= x1
372
+ if x2 >= mask_image.width:
373
+ x2 = mask_image.width
374
+
375
+ return x1, y1, x2, y2
376
+
377
+ def _resize_and_fill(
378
+ self,
379
+ image: PIL.Image.Image,
380
+ width: int,
381
+ height: int,
382
+ ) -> PIL.Image.Image:
383
+ r"""
384
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
385
+ the image within the dimensions, filling empty with data from image.
386
+
387
+ Args:
388
+ image (`PIL.Image.Image`):
389
+ The image to resize and fill.
390
+ width (`int`):
391
+ The width to resize the image to.
392
+ height (`int`):
393
+ The height to resize the image to.
394
+
395
+ Returns:
396
+ `PIL.Image.Image`:
397
+ The resized and filled image.
398
+ """
399
+
400
+ ratio = width / height
401
+ src_ratio = image.width / image.height
402
+
403
+ src_w = width if ratio < src_ratio else image.width * height // image.height
404
+ src_h = height if ratio >= src_ratio else image.height * width // image.width
405
+
406
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
407
+ res = Image.new("RGB", (width, height))
408
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
409
+
410
+ if ratio < src_ratio:
411
+ fill_height = height // 2 - src_h // 2
412
+ if fill_height > 0:
413
+ res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
414
+ res.paste(
415
+ resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)),
416
+ box=(0, fill_height + src_h),
417
+ )
418
+ elif ratio > src_ratio:
419
+ fill_width = width // 2 - src_w // 2
420
+ if fill_width > 0:
421
+ res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
422
+ res.paste(
423
+ resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)),
424
+ box=(fill_width + src_w, 0),
425
+ )
426
+
427
+ return res
428
+
429
+ def _resize_and_crop(
430
+ self,
431
+ image: PIL.Image.Image,
432
+ width: int,
433
+ height: int,
434
+ ) -> PIL.Image.Image:
435
+ r"""
436
+ Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
437
+ the image within the dimensions, cropping the excess.
438
+
439
+ Args:
440
+ image (`PIL.Image.Image`):
441
+ The image to resize and crop.
442
+ width (`int`):
443
+ The width to resize the image to.
444
+ height (`int`):
445
+ The height to resize the image to.
446
+
447
+ Returns:
448
+ `PIL.Image.Image`:
449
+ The resized and cropped image.
450
+ """
451
+ ratio = width / height
452
+ src_ratio = image.width / image.height
453
+
454
+ src_w = width if ratio > src_ratio else image.width * height // image.height
455
+ src_h = height if ratio <= src_ratio else image.height * width // image.width
456
+
457
+ resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
458
+ res = Image.new("RGB", (width, height))
459
+ res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
460
+ return res
461
+
462
+ def resize(
463
+ self,
464
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
465
+ height: int,
466
+ width: int,
467
+ resize_mode: str = "default", # "default", "fill", "crop"
468
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
469
+ """
470
+ Resize image.
471
+
472
+ Args:
473
+ image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
474
+ The image input, can be a PIL image, numpy array or pytorch tensor.
475
+ height (`int`):
476
+ The height to resize to.
477
+ width (`int`):
478
+ The width to resize to.
479
+ resize_mode (`str`, *optional*, defaults to `default`):
480
+ The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit
481
+ within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`,
482
+ will resize the image to fit within the specified width and height, maintaining the aspect ratio, and
483
+ then center the image within the dimensions, filling empty with data from image. If `crop`, will resize
484
+ the image to fit within the specified width and height, maintaining the aspect ratio, and then center
485
+ the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
486
+ supported for PIL image input.
487
+
488
+ Returns:
489
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
490
+ The resized image.
491
+ """
492
+ if resize_mode != "default" and not isinstance(image, PIL.Image.Image):
493
+ raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}")
494
+ if isinstance(image, PIL.Image.Image):
495
+ if resize_mode == "default":
496
+ image = image.resize(
497
+ (width, height),
498
+ resample=PIL_INTERPOLATION[self.config.resample],
499
+ reducing_gap=self.config.reducing_gap,
500
+ )
501
+ elif resize_mode == "fill":
502
+ image = self._resize_and_fill(image, width, height)
503
+ elif resize_mode == "crop":
504
+ image = self._resize_and_crop(image, width, height)
505
+ else:
506
+ raise ValueError(f"resize_mode {resize_mode} is not supported")
507
+
508
+ elif isinstance(image, torch.Tensor):
509
+ image = torch.nn.functional.interpolate(
510
+ image,
511
+ size=(height, width),
512
+ )
513
+ elif isinstance(image, np.ndarray):
514
+ image = self.numpy_to_pt(image)
515
+ image = torch.nn.functional.interpolate(
516
+ image,
517
+ size=(height, width),
518
+ )
519
+ image = self.pt_to_numpy(image)
520
+
521
+ return image
522
+
523
+ def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:
524
+ """
525
+ Create a mask.
526
+
527
+ Args:
528
+ image (`PIL.Image.Image`):
529
+ The image input, should be a PIL image.
530
+
531
+ Returns:
532
+ `PIL.Image.Image`:
533
+ The binarized image. Values less than 0.5 are set to 0, values greater than 0.5 are set to 1.
534
+ """
535
+ image[image < 0.5] = 0
536
+ image[image >= 0.5] = 1
537
+
538
+ return image
539
+
540
+ def _denormalize_conditionally(
541
+ self, images: torch.Tensor, do_denormalize: list[bool] | None = None
542
+ ) -> torch.Tensor:
543
+ r"""
544
+ Denormalize a batch of images based on a condition list.
545
+
546
+ Args:
547
+ images (`torch.Tensor`):
548
+ The input image tensor.
549
+ do_denormalize (`Optional[list[bool]`, *optional*, defaults to `None`):
550
+ A list of booleans indicating whether to denormalize each image in the batch. If `None`, will use the
551
+ value of `do_normalize` in the `VaeImageProcessor` config.
552
+ """
553
+ if do_denormalize is None:
554
+ return self.denormalize(images) if self.config.do_normalize else images
555
+
556
+ return torch.stack(
557
+ [self.denormalize(images[i]) if do_denormalize[i] else images[i] for i in range(images.shape[0])]
558
+ )
559
+
560
+ def get_default_height_width(
561
+ self,
562
+ image: PIL.Image.Image | np.ndarray | torch.Tensor,
563
+ height: int | None = None,
564
+ width: int | None = None,
565
+ ) -> tuple[int, int]:
566
+ r"""
567
+ Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
568
+
569
+ Args:
570
+ image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
571
+ The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
572
+ should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
573
+ tensor, it should have shape `[batch, channels, height, width]`.
574
+ height (`int | None`, *optional*, defaults to `None`):
575
+ The height of the preprocessed image. If `None`, the height of the `image` input will be used.
576
+ width (`int | None`, *optional*, defaults to `None`):
577
+ The width of the preprocessed image. If `None`, the width of the `image` input will be used.
578
+
579
+ Returns:
580
+ `tuple[int, int]`:
581
+ A tuple containing the height and width, both resized to the nearest integer multiple of
582
+ `vae_scale_factor`.
583
+ """
584
+
585
+ if height is None:
586
+ if isinstance(image, PIL.Image.Image):
587
+ height = image.height
588
+ elif isinstance(image, torch.Tensor):
589
+ height = image.shape[2]
590
+ else:
591
+ height = image.shape[1]
592
+
593
+ if width is None:
594
+ if isinstance(image, PIL.Image.Image):
595
+ width = image.width
596
+ elif isinstance(image, torch.Tensor):
597
+ width = image.shape[3]
598
+ else:
599
+ width = image.shape[2]
600
+
601
+ width, height = (
602
+ x - x % self.config.vae_scale_factor for x in (width, height)
603
+ ) # resize to integer multiple of vae_scale_factor
604
+
605
+ return height, width
606
+
607
+ def preprocess(
608
+ self,
609
+ image: PipelineImageInput,
610
+ height: int | None = None,
611
+ width: int | None = None,
612
+ resize_mode: str = "default", # "default", "fill", "crop"
613
+ crops_coords: tuple[int, int, int, int] | None = None,
614
+ ) -> torch.Tensor:
615
+ """
616
+ Preprocess the image input.
617
+
618
+ Args:
619
+ image (`PipelineImageInput`):
620
+ The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of
621
+ supported formats.
622
+ height (`int`, *optional*):
623
+ The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default
624
+ height.
625
+ width (`int`, *optional*):
626
+ The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width.
627
+ resize_mode (`str`, *optional*, defaults to `default`):
628
+ The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within
629
+ the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will
630
+ resize the image to fit within the specified width and height, maintaining the aspect ratio, and then
631
+ center the image within the dimensions, filling empty with data from image. If `crop`, will resize the
632
+ image to fit within the specified width and height, maintaining the aspect ratio, and then center the
633
+ image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
634
+ supported for PIL image input.
635
+ crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`):
636
+ The crop coordinates for each image in the batch. If `None`, will not crop the image.
637
+
638
+ Returns:
639
+ `torch.Tensor`:
640
+ The preprocessed image.
641
+ """
642
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
643
+
644
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
645
+ if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
646
+ if isinstance(image, torch.Tensor):
647
+ # if image is a pytorch tensor could have 2 possible shapes:
648
+ # 1. batch x height x width: we should insert the channel dimension at position 1
649
+ # 2. channel x height x width: we should insert batch dimension at position 0,
650
+ # however, since both channel and batch dimension has same size 1, it is same to insert at position 1
651
+ # for simplicity, we insert a dimension of size 1 at position 1 for both cases
652
+ image = image.unsqueeze(1)
653
+ else:
654
+ # if it is a numpy array, it could have 2 possible shapes:
655
+ # 1. batch x height x width: insert channel dimension on last position
656
+ # 2. height x width x channel: insert batch dimension on first position
657
+ if image.shape[-1] == 1:
658
+ image = np.expand_dims(image, axis=0)
659
+ else:
660
+ image = np.expand_dims(image, axis=-1)
661
+
662
+ if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4:
663
+ warnings.warn(
664
+ "Passing `image` as a list of 4d np.ndarray is deprecated."
665
+ "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
666
+ FutureWarning,
667
+ )
668
+ image = np.concatenate(image, axis=0)
669
+ if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
670
+ warnings.warn(
671
+ "Passing `image` as a list of 4d torch.Tensor is deprecated."
672
+ "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
673
+ FutureWarning,
674
+ )
675
+ image = torch.cat(image, axis=0)
676
+
677
+ if not is_valid_image_imagelist(image):
678
+ raise ValueError(
679
+ f"Input is in incorrect format. Currently, we only support {', '.join(str(x) for x in supported_formats)}"
680
+ )
681
+ if not isinstance(image, list):
682
+ image = [image]
683
+
684
+ if isinstance(image[0], PIL.Image.Image):
685
+ if crops_coords is not None:
686
+ image = [i.crop(crops_coords) for i in image]
687
+ if self.config.do_resize:
688
+ height, width = self.get_default_height_width(image[0], height, width)
689
+ image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
690
+ if self.config.do_convert_rgb:
691
+ image = [self.convert_to_rgb(i) for i in image]
692
+ elif self.config.do_convert_grayscale:
693
+ image = [self.convert_to_grayscale(i) for i in image]
694
+ image = self.pil_to_numpy(image) # to np
695
+ image = self.numpy_to_pt(image) # to pt
696
+
697
+ elif isinstance(image[0], np.ndarray):
698
+ image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
699
+
700
+ image = self.numpy_to_pt(image)
701
+
702
+ height, width = self.get_default_height_width(image, height, width)
703
+ if self.config.do_resize:
704
+ image = self.resize(image, height, width)
705
+
706
+ elif isinstance(image[0], torch.Tensor):
707
+ image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
708
+
709
+ if self.config.do_convert_grayscale and image.ndim == 3:
710
+ image = image.unsqueeze(1)
711
+
712
+ channel = image.shape[1]
713
+ # don't need any preprocess if the image is latents
714
+ if channel == self.config.vae_latent_channels:
715
+ return image
716
+
717
+ height, width = self.get_default_height_width(image, height, width)
718
+ if self.config.do_resize:
719
+ image = self.resize(image, height, width)
720
+
721
+ # expected range [0,1], normalize to [-1,1]
722
+ do_normalize = self.config.do_normalize
723
+ if do_normalize and image.min() < 0:
724
+ warnings.warn(
725
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
726
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
727
+ FutureWarning,
728
+ )
729
+ do_normalize = False
730
+ if do_normalize:
731
+ image = self.normalize(image)
732
+
733
+ if self.config.do_binarize:
734
+ image = self.binarize(image)
735
+
736
+ return image
737
+
738
+ def postprocess(
739
+ self,
740
+ image: torch.Tensor,
741
+ output_type: str = "pil",
742
+ do_denormalize: list[bool] | None = None,
743
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
744
+ """
745
+ Postprocess the image output from tensor to `output_type`.
746
+
747
+ Args:
748
+ image (`torch.Tensor`):
749
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
750
+ output_type (`str`, *optional*, defaults to `pil`):
751
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
752
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
753
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
754
+ `VaeImageProcessor` config.
755
+
756
+ Returns:
757
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
758
+ The postprocessed image.
759
+ """
760
+ if not isinstance(image, torch.Tensor):
761
+ raise ValueError(
762
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
763
+ )
764
+ if output_type not in ["latent", "pt", "np", "pil"]:
765
+ deprecation_message = (
766
+ 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: "
767
+ "`pil`, `np`, `pt`, `latent`"
768
+ )
769
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
770
+ output_type = "np"
771
+
772
+ if output_type == "latent":
773
+ return image
774
+
775
+ image = self._denormalize_conditionally(image, do_denormalize)
776
+
777
+ if output_type == "pt":
778
+ return image
779
+
780
+ image = self.pt_to_numpy(image)
781
+
782
+ if output_type == "np":
783
+ return image
784
+
785
+ if output_type == "pil":
786
+ return self.numpy_to_pil(image)
787
+
788
+ def apply_overlay(
789
+ self,
790
+ mask: PIL.Image.Image,
791
+ init_image: PIL.Image.Image,
792
+ image: PIL.Image.Image,
793
+ crop_coords: tuple[int, int, int, int] | None = None,
794
+ ) -> PIL.Image.Image:
795
+ r"""
796
+ Applies an overlay of the mask and the inpainted image on the original image.
797
+
798
+ Args:
799
+ mask (`PIL.Image.Image`):
800
+ The mask image that highlights regions to overlay.
801
+ init_image (`PIL.Image.Image`):
802
+ The original image to which the overlay is applied.
803
+ image (`PIL.Image.Image`):
804
+ The image to overlay onto the original.
805
+ crop_coords (`tuple[int, int, int, int]`, *optional*):
806
+ Coordinates to crop the image. If provided, the image will be cropped accordingly.
807
+
808
+ Returns:
809
+ `PIL.Image.Image`:
810
+ The final image with the overlay applied.
811
+ """
812
+
813
+ width, height = init_image.width, init_image.height
814
+
815
+ init_image_masked = PIL.Image.new("RGBa", (width, height))
816
+ init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L")))
817
+
818
+ init_image_masked = init_image_masked.convert("RGBA")
819
+
820
+ if crop_coords is not None:
821
+ x, y, x2, y2 = crop_coords
822
+ w = x2 - x
823
+ h = y2 - y
824
+ base_image = PIL.Image.new("RGBA", (width, height))
825
+ image = self.resize(image, height=h, width=w, resize_mode="crop")
826
+ base_image.paste(image, (x, y))
827
+ image = base_image.convert("RGB")
828
+
829
+ image = image.convert("RGBA")
830
+ image.alpha_composite(init_image_masked)
831
+ image = image.convert("RGB")
832
+
833
+ return image
834
+
835
+
836
+ class InpaintProcessor(ConfigMixin):
837
+ """
838
+ Image processor for inpainting image and mask.
839
+ """
840
+
841
+ config_name = CONFIG_NAME
842
+
843
+ @register_to_config
844
+ def __init__(
845
+ self,
846
+ do_resize: bool = True,
847
+ vae_scale_factor: int = 8,
848
+ vae_latent_channels: int = 4,
849
+ resample: str = "lanczos",
850
+ reducing_gap: int | None = None,
851
+ do_normalize: bool = True,
852
+ do_binarize: bool = False,
853
+ do_convert_grayscale: bool = False,
854
+ mask_do_normalize: bool = False,
855
+ mask_do_binarize: bool = True,
856
+ mask_do_convert_grayscale: bool = True,
857
+ ):
858
+ super().__init__()
859
+
860
+ self._image_processor = VaeImageProcessor(
861
+ do_resize=do_resize,
862
+ vae_scale_factor=vae_scale_factor,
863
+ vae_latent_channels=vae_latent_channels,
864
+ resample=resample,
865
+ reducing_gap=reducing_gap,
866
+ do_normalize=do_normalize,
867
+ do_binarize=do_binarize,
868
+ do_convert_grayscale=do_convert_grayscale,
869
+ )
870
+ self._mask_processor = VaeImageProcessor(
871
+ do_resize=do_resize,
872
+ vae_scale_factor=vae_scale_factor,
873
+ vae_latent_channels=vae_latent_channels,
874
+ resample=resample,
875
+ reducing_gap=reducing_gap,
876
+ do_normalize=mask_do_normalize,
877
+ do_binarize=mask_do_binarize,
878
+ do_convert_grayscale=mask_do_convert_grayscale,
879
+ )
880
+
881
+ def preprocess(
882
+ self,
883
+ image: PIL.Image.Image,
884
+ mask: PIL.Image.Image | None = None,
885
+ height: int | None = None,
886
+ width: int | None = None,
887
+ padding_mask_crop: int | None = None,
888
+ ) -> tuple[torch.Tensor, torch.Tensor]:
889
+ """
890
+ Preprocess the image and mask.
891
+ """
892
+ if mask is None and padding_mask_crop is not None:
893
+ raise ValueError("mask must be provided if padding_mask_crop is provided")
894
+
895
+ # if mask is None, same behavior as regular image processor
896
+ if mask is None:
897
+ return self._image_processor.preprocess(image, height=height, width=width)
898
+
899
+ if padding_mask_crop is not None:
900
+ crops_coords = self._image_processor.get_crop_region(mask, width, height, pad=padding_mask_crop)
901
+ resize_mode = "fill"
902
+ else:
903
+ crops_coords = None
904
+ resize_mode = "default"
905
+
906
+ processed_image = self._image_processor.preprocess(
907
+ image,
908
+ height=height,
909
+ width=width,
910
+ crops_coords=crops_coords,
911
+ resize_mode=resize_mode,
912
+ )
913
+
914
+ processed_mask = self._mask_processor.preprocess(
915
+ mask,
916
+ height=height,
917
+ width=width,
918
+ resize_mode=resize_mode,
919
+ crops_coords=crops_coords,
920
+ )
921
+
922
+ if crops_coords is not None:
923
+ postprocessing_kwargs = {
924
+ "crops_coords": crops_coords,
925
+ "original_image": image,
926
+ "original_mask": mask,
927
+ }
928
+ else:
929
+ postprocessing_kwargs = {
930
+ "crops_coords": None,
931
+ "original_image": None,
932
+ "original_mask": None,
933
+ }
934
+
935
+ return processed_image, processed_mask, postprocessing_kwargs
936
+
937
+ def postprocess(
938
+ self,
939
+ image: torch.Tensor,
940
+ output_type: str = "pil",
941
+ original_image: PIL.Image.Image | None = None,
942
+ original_mask: PIL.Image.Image | None = None,
943
+ crops_coords: tuple[int, int, int, int] | None = None,
944
+ ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
945
+ """
946
+ Postprocess the image, optionally apply mask overlay
947
+ """
948
+ image = self._image_processor.postprocess(
949
+ image,
950
+ output_type=output_type,
951
+ )
952
+ # optionally apply the mask overlay
953
+ if crops_coords is not None and (original_image is None or original_mask is None):
954
+ raise ValueError("original_image and original_mask must be provided if crops_coords is provided")
955
+
956
+ elif crops_coords is not None and output_type != "pil":
957
+ raise ValueError("output_type must be 'pil' if crops_coords is provided")
958
+
959
+ elif crops_coords is not None:
960
+ image = [
961
+ self._image_processor.apply_overlay(original_mask, original_image, i, crops_coords) for i in image
962
+ ]
963
+
964
+ return image
965
+
966
+
967
+ class VaeImageProcessorLDM3D(VaeImageProcessor):
968
+ """
969
+ Image processor for VAE LDM3D.
970
+
971
+ Args:
972
+ do_resize (`bool`, *optional*, defaults to `True`):
973
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
974
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
975
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
976
+ resample (`str`, *optional*, defaults to `lanczos`):
977
+ Resampling filter to use when resizing the image.
978
+ do_normalize (`bool`, *optional*, defaults to `True`):
979
+ Whether to normalize the image to [-1,1].
980
+ """
981
+
982
+ config_name = CONFIG_NAME
983
+
984
+ @register_to_config
985
+ def __init__(
986
+ self,
987
+ do_resize: bool = True,
988
+ vae_scale_factor: int = 8,
989
+ resample: str = "lanczos",
990
+ do_normalize: bool = True,
991
+ ):
992
+ super().__init__()
993
+
994
+ @staticmethod
995
+ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
996
+ r"""
997
+ Convert a NumPy image or a batch of images to a list of PIL images.
998
+
999
+ Args:
1000
+ images (`np.ndarray`):
1001
+ The input NumPy array of images, which can be a single image or a batch.
1002
+
1003
+ Returns:
1004
+ `list[PIL.Image.Image]`:
1005
+ A list of PIL images converted from the input NumPy array.
1006
+ """
1007
+ if images.ndim == 3:
1008
+ images = images[None, ...]
1009
+ images = (images * 255).round().astype("uint8")
1010
+ if images.shape[-1] == 1:
1011
+ # special case for grayscale (single channel) images
1012
+ pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
1013
+ else:
1014
+ pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
1015
+
1016
+ return pil_images
1017
+
1018
+ @staticmethod
1019
+ def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
1020
+ r"""
1021
+ Convert a PIL image or a list of PIL images to NumPy arrays.
1022
+
1023
+ Args:
1024
+ images (`list[PIL.Image.Image, PIL.Image.Image]`):
1025
+ The input image or list of images to be converted.
1026
+
1027
+ Returns:
1028
+ `np.ndarray`:
1029
+ A NumPy array of the converted images.
1030
+ """
1031
+ if not isinstance(images, list):
1032
+ images = [images]
1033
+
1034
+ images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]
1035
+ images = np.stack(images, axis=0)
1036
+ return images
1037
+
1038
+ @staticmethod
1039
+ def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
1040
+ r"""
1041
+ Convert an RGB-like depth image to a depth map.
1042
+ """
1043
+ # 1. Cast the tensor to a larger integer type (e.g., int32)
1044
+ # to safely perform the multiplication by 256.
1045
+ # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte.
1046
+ # 3. Cast the final result to the desired depth map type (uint16) if needed
1047
+ # before returning, though leaving it as int32/int64 is often safer
1048
+ # for return value from a library function.
1049
+
1050
+ if isinstance(image, torch.Tensor):
1051
+ # Cast to a safe dtype (e.g., int32 or int64) for the calculation
1052
+ original_dtype = image.dtype
1053
+ image_safe = image.to(torch.int32)
1054
+
1055
+ # Calculate the depth map
1056
+ depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
1057
+
1058
+ # You may want to cast the final result to uint16, but casting to a
1059
+ # larger int type (like int32) is sufficient to fix the overflow.
1060
+ # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required
1061
+ return depth_map.to(original_dtype)
1062
+
1063
+ elif isinstance(image, np.ndarray):
1064
+ # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
1065
+ original_dtype = image.dtype
1066
+ image_safe = image.astype(np.int32)
1067
+
1068
+ # Calculate the depth map
1069
+ depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
1070
+
1071
+ # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required
1072
+ return depth_map.astype(original_dtype)
1073
+ else:
1074
+ raise TypeError("Input image must be a torch.Tensor or np.ndarray")
1075
+
1076
+ def numpy_to_depth(self, images: np.ndarray) -> list[PIL.Image.Image]:
1077
+ r"""
1078
+ Convert a NumPy depth image or a batch of images to a list of PIL images.
1079
+
1080
+ Args:
1081
+ images (`np.ndarray`):
1082
+ The input NumPy array of depth images, which can be a single image or a batch.
1083
+
1084
+ Returns:
1085
+ `list[PIL.Image.Image]`:
1086
+ A list of PIL images converted from the input NumPy depth images.
1087
+ """
1088
+ if images.ndim == 3:
1089
+ images = images[None, ...]
1090
+ images_depth = images[:, :, :, 3:]
1091
+ if images.shape[-1] == 6:
1092
+ images_depth = (images_depth * 255).round().astype("uint8")
1093
+ pil_images = [
1094
+ Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
1095
+ ]
1096
+ elif images.shape[-1] == 4:
1097
+ images_depth = (images_depth * 65535.0).astype(np.uint16)
1098
+ pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
1099
+ else:
1100
+ raise Exception("Not supported")
1101
+
1102
+ return pil_images
1103
+
1104
+ def postprocess(
1105
+ self,
1106
+ image: torch.Tensor,
1107
+ output_type: str = "pil",
1108
+ do_denormalize: list[bool] | None = None,
1109
+ ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
1110
+ """
1111
+ Postprocess the image output from tensor to `output_type`.
1112
+
1113
+ Args:
1114
+ image (`torch.Tensor`):
1115
+ The image input, should be a pytorch tensor with shape `B x C x H x W`.
1116
+ output_type (`str`, *optional*, defaults to `pil`):
1117
+ The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
1118
+ do_denormalize (`list[bool]`, *optional*, defaults to `None`):
1119
+ Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
1120
+ `VaeImageProcessor` config.
1121
+
1122
+ Returns:
1123
+ `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
1124
+ The postprocessed image.
1125
+ """
1126
+ if not isinstance(image, torch.Tensor):
1127
+ raise ValueError(
1128
+ f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
1129
+ )
1130
+ if output_type not in ["latent", "pt", "np", "pil"]:
1131
+ deprecation_message = (
1132
+ 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: "
1133
+ "`pil`, `np`, `pt`, `latent`"
1134
+ )
1135
+ deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
1136
+ output_type = "np"
1137
+
1138
+ image = self._denormalize_conditionally(image, do_denormalize)
1139
+
1140
+ image = self.pt_to_numpy(image)
1141
+
1142
+ if output_type == "np":
1143
+ if image.shape[-1] == 6:
1144
+ image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
1145
+ else:
1146
+ image_depth = image[:, :, :, 3:]
1147
+ return image[:, :, :, :3], image_depth
1148
+
1149
+ if output_type == "pil":
1150
+ return self.numpy_to_pil(image), self.numpy_to_depth(image)
1151
+ else:
1152
+ raise Exception(f"This type {output_type} is not supported")
1153
+
1154
+ def preprocess(
1155
+ self,
1156
+ rgb: torch.Tensor | PIL.Image.Image | np.ndarray,
1157
+ depth: torch.Tensor | PIL.Image.Image | np.ndarray,
1158
+ height: int | None = None,
1159
+ width: int | None = None,
1160
+ target_res: int | None = None,
1161
+ ) -> torch.Tensor:
1162
+ r"""
1163
+ Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors.
1164
+
1165
+ Args:
1166
+ rgb (`torch.Tensor | PIL.Image.Image | np.ndarray`):
1167
+ The RGB input image, which can be a single image or a batch.
1168
+ depth (`torch.Tensor | PIL.Image.Image | np.ndarray`):
1169
+ The depth input image, which can be a single image or a batch.
1170
+ height (`int | None`, *optional*, defaults to `None`):
1171
+ The desired height of the processed image. If `None`, defaults to the height of the input image.
1172
+ width (`int | None`, *optional*, defaults to `None`):
1173
+ The desired width of the processed image. If `None`, defaults to the width of the input image.
1174
+ target_res (`int | None`, *optional*, defaults to `None`):
1175
+ Target resolution for resizing the images. If specified, overrides height and width.
1176
+
1177
+ Returns:
1178
+ `tuple[torch.Tensor, torch.Tensor]`:
1179
+ A tuple containing the processed RGB and depth images as PyTorch tensors.
1180
+ """
1181
+ supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
1182
+
1183
+ # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
1184
+ if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:
1185
+ raise Exception("This is not yet supported")
1186
+
1187
+ if isinstance(rgb, supported_formats):
1188
+ rgb = [rgb]
1189
+ depth = [depth]
1190
+ elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):
1191
+ raise ValueError(
1192
+ f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"
1193
+ )
1194
+
1195
+ if isinstance(rgb[0], PIL.Image.Image):
1196
+ if self.config.do_convert_rgb:
1197
+ raise Exception("This is not yet supported")
1198
+ # rgb = [self.convert_to_rgb(i) for i in rgb]
1199
+ # depth = [self.convert_to_depth(i) for i in depth] #TODO define convert_to_depth
1200
+ if self.config.do_resize or target_res:
1201
+ height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res
1202
+ rgb = [self.resize(i, height, width) for i in rgb]
1203
+ depth = [self.resize(i, height, width) for i in depth]
1204
+ rgb = self.pil_to_numpy(rgb) # to np
1205
+ rgb = self.numpy_to_pt(rgb) # to pt
1206
+
1207
+ depth = self.depth_pil_to_numpy(depth) # to np
1208
+ depth = self.numpy_to_pt(depth) # to pt
1209
+
1210
+ elif isinstance(rgb[0], np.ndarray):
1211
+ rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)
1212
+ rgb = self.numpy_to_pt(rgb)
1213
+ height, width = self.get_default_height_width(rgb, height, width)
1214
+ if self.config.do_resize:
1215
+ rgb = self.resize(rgb, height, width)
1216
+
1217
+ depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0)
1218
+ depth = self.numpy_to_pt(depth)
1219
+ height, width = self.get_default_height_width(depth, height, width)
1220
+ if self.config.do_resize:
1221
+ depth = self.resize(depth, height, width)
1222
+
1223
+ elif isinstance(rgb[0], torch.Tensor):
1224
+ raise Exception("This is not yet supported")
1225
+ # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)
1226
+
1227
+ # if self.config.do_convert_grayscale and rgb.ndim == 3:
1228
+ # rgb = rgb.unsqueeze(1)
1229
+
1230
+ # channel = rgb.shape[1]
1231
+
1232
+ # height, width = self.get_default_height_width(rgb, height, width)
1233
+ # if self.config.do_resize:
1234
+ # rgb = self.resize(rgb, height, width)
1235
+
1236
+ # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)
1237
+
1238
+ # if self.config.do_convert_grayscale and depth.ndim == 3:
1239
+ # depth = depth.unsqueeze(1)
1240
+
1241
+ # channel = depth.shape[1]
1242
+ # # don't need any preprocess if the image is latents
1243
+ # if depth == 4:
1244
+ # return rgb, depth
1245
+
1246
+ # height, width = self.get_default_height_width(depth, height, width)
1247
+ # if self.config.do_resize:
1248
+ # depth = self.resize(depth, height, width)
1249
+ # expected range [0,1], normalize to [-1,1]
1250
+ do_normalize = self.config.do_normalize
1251
+ if rgb.min() < 0 and do_normalize:
1252
+ warnings.warn(
1253
+ "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
1254
+ f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",
1255
+ FutureWarning,
1256
+ )
1257
+ do_normalize = False
1258
+
1259
+ if do_normalize:
1260
+ rgb = self.normalize(rgb)
1261
+ depth = self.normalize(depth)
1262
+
1263
+ if self.config.do_binarize:
1264
+ rgb = self.binarize(rgb)
1265
+ depth = self.binarize(depth)
1266
+
1267
+ return rgb, depth
1268
+
1269
+
1270
+ class IPAdapterMaskProcessor(VaeImageProcessor):
1271
+ """
1272
+ Image processor for IP Adapter image masks.
1273
+
1274
+ Args:
1275
+ do_resize (`bool`, *optional*, defaults to `True`):
1276
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
1277
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
1278
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
1279
+ resample (`str`, *optional*, defaults to `lanczos`):
1280
+ Resampling filter to use when resizing the image.
1281
+ do_normalize (`bool`, *optional*, defaults to `False`):
1282
+ Whether to normalize the image to [-1,1].
1283
+ do_binarize (`bool`, *optional*, defaults to `True`):
1284
+ Whether to binarize the image to 0/1.
1285
+ do_convert_grayscale (`bool`, *optional*, defaults to be `True`):
1286
+ Whether to convert the images to grayscale format.
1287
+
1288
+ """
1289
+
1290
+ config_name = CONFIG_NAME
1291
+
1292
+ @register_to_config
1293
+ def __init__(
1294
+ self,
1295
+ do_resize: bool = True,
1296
+ vae_scale_factor: int = 8,
1297
+ resample: str = "lanczos",
1298
+ do_normalize: bool = False,
1299
+ do_binarize: bool = True,
1300
+ do_convert_grayscale: bool = True,
1301
+ ):
1302
+ super().__init__(
1303
+ do_resize=do_resize,
1304
+ vae_scale_factor=vae_scale_factor,
1305
+ resample=resample,
1306
+ do_normalize=do_normalize,
1307
+ do_binarize=do_binarize,
1308
+ do_convert_grayscale=do_convert_grayscale,
1309
+ )
1310
+
1311
+ @staticmethod
1312
+ def downsample(mask: torch.Tensor, batch_size: int, num_queries: int, value_embed_dim: int):
1313
+ """
1314
+ Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the
1315
+ aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.
1316
+
1317
+ Args:
1318
+ mask (`torch.Tensor`):
1319
+ The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
1320
+ batch_size (`int`):
1321
+ The batch size.
1322
+ num_queries (`int`):
1323
+ The number of queries.
1324
+ value_embed_dim (`int`):
1325
+ The dimensionality of the value embeddings.
1326
+
1327
+ Returns:
1328
+ `torch.Tensor`:
1329
+ The downsampled mask tensor.
1330
+
1331
+ """
1332
+ o_h = mask.shape[1]
1333
+ o_w = mask.shape[2]
1334
+ ratio = o_w / o_h
1335
+ mask_h = int(math.sqrt(num_queries / ratio))
1336
+ mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
1337
+ mask_w = num_queries // mask_h
1338
+
1339
+ mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)
1340
+
1341
+ # Repeat batch_size times
1342
+ if mask_downsample.shape[0] < batch_size:
1343
+ mask_downsample = mask_downsample.repeat(batch_size, 1, 1)
1344
+
1345
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)
1346
+
1347
+ downsampled_area = mask_h * mask_w
1348
+ # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
1349
+ # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
1350
+ if downsampled_area < num_queries:
1351
+ warnings.warn(
1352
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
1353
+ "Please update your masks or adjust the output size for optimal performance.",
1354
+ UserWarning,
1355
+ )
1356
+ mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
1357
+ # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
1358
+ if downsampled_area > num_queries:
1359
+ warnings.warn(
1360
+ "The aspect ratio of the mask does not match the aspect ratio of the output image. "
1361
+ "Please update your masks or adjust the output size for optimal performance.",
1362
+ UserWarning,
1363
+ )
1364
+ mask_downsample = mask_downsample[:, :num_queries]
1365
+
1366
+ # Repeat last dimension to match SDPA output shape
1367
+ mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
1368
+ 1, 1, value_embed_dim
1369
+ )
1370
+
1371
+ return mask_downsample
1372
+
1373
+
1374
+ class PixArtImageProcessor(VaeImageProcessor):
1375
+ """
1376
+ Image processor for PixArt image resize and crop.
1377
+
1378
+ Args:
1379
+ do_resize (`bool`, *optional*, defaults to `True`):
1380
+ Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
1381
+ `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
1382
+ vae_scale_factor (`int`, *optional*, defaults to `8`):
1383
+ VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
1384
+ resample (`str`, *optional*, defaults to `lanczos`):
1385
+ Resampling filter to use when resizing the image.
1386
+ do_normalize (`bool`, *optional*, defaults to `True`):
1387
+ Whether to normalize the image to [-1,1].
1388
+ do_binarize (`bool`, *optional*, defaults to `False`):
1389
+ Whether to binarize the image to 0/1.
1390
+ do_convert_rgb (`bool`, *optional*, defaults to be `False`):
1391
+ Whether to convert the images to RGB format.
1392
+ do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
1393
+ Whether to convert the images to grayscale format.
1394
+ """
1395
+
1396
+ @register_to_config
1397
+ def __init__(
1398
+ self,
1399
+ do_resize: bool = True,
1400
+ vae_scale_factor: int = 8,
1401
+ resample: str = "lanczos",
1402
+ do_normalize: bool = True,
1403
+ do_binarize: bool = False,
1404
+ do_convert_grayscale: bool = False,
1405
+ ):
1406
+ super().__init__(
1407
+ do_resize=do_resize,
1408
+ vae_scale_factor=vae_scale_factor,
1409
+ resample=resample,
1410
+ do_normalize=do_normalize,
1411
+ do_binarize=do_binarize,
1412
+ do_convert_grayscale=do_convert_grayscale,
1413
+ )
1414
+
1415
+ @staticmethod
1416
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
1417
+ r"""
1418
+ Returns the binned height and width based on the aspect ratio.
1419
+
1420
+ Args:
1421
+ height (`int`): The height of the image.
1422
+ width (`int`): The width of the image.
1423
+ ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
1424
+
1425
+ Returns:
1426
+ `tuple[int, int]`: The closest binned height and width.
1427
+ """
1428
+ ar = float(height / width)
1429
+ closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
1430
+ default_hw = ratios[closest_ratio]
1431
+ return int(default_hw[0]), int(default_hw[1])
1432
+
1433
+ @staticmethod
1434
+ def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
1435
+ r"""
1436
+ Resizes and crops a tensor of images to the specified dimensions.
1437
+
1438
+ Args:
1439
+ samples (`torch.Tensor`):
1440
+ A tensor of shape (N, C, H, W) where N is the batch size, C is the number of channels, H is the height,
1441
+ and W is the width.
1442
+ new_width (`int`): The desired width of the output images.
1443
+ new_height (`int`): The desired height of the output images.
1444
+
1445
+ Returns:
1446
+ `torch.Tensor`: A tensor containing the resized and cropped images.
1447
+ """
1448
+ orig_height, orig_width = samples.shape[2], samples.shape[3]
1449
+
1450
+ # Check if resizing is needed
1451
+ if orig_height != new_height or orig_width != new_width:
1452
+ ratio = max(new_height / orig_height, new_width / orig_width)
1453
+ resized_width = int(orig_width * ratio)
1454
+ resized_height = int(orig_height * ratio)
1455
+
1456
+ # Resize
1457
+ samples = F.interpolate(
1458
+ samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
1459
+ )
1460
+
1461
+ # Center Crop
1462
+ start_x = (resized_width - new_width) // 2
1463
+ end_x = start_x + new_width
1464
+ start_y = (resized_height - new_height) // 2
1465
+ end_y = start_y + new_height
1466
+ samples = samples[:, :, start_y:end_y, start_x:end_x]
1467
+
1468
+ return samples
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ..utils import (
4
+ DIFFUSERS_SLOW_IMPORT,
5
+ OptionalDependencyNotAvailable,
6
+ _LazyModule,
7
+ get_objects_from_module,
8
+ is_torch_available,
9
+ is_transformers_available,
10
+ logging,
11
+ )
12
+
13
+
14
+ logger = logging.get_logger(__name__)
15
+ logger.warning(
16
+ "Modular Diffusers is currently an experimental feature under active development. The API is subject to breaking changes in future releases."
17
+ )
18
+
19
+ # These modules contain pipelines from multiple libraries/frameworks
20
+ _dummy_objects = {}
21
+ _import_structure = {}
22
+
23
+ try:
24
+ if not is_torch_available():
25
+ raise OptionalDependencyNotAvailable()
26
+ except OptionalDependencyNotAvailable:
27
+ from ..utils import dummy_pt_objects # noqa F403
28
+
29
+ _dummy_objects.update(get_objects_from_module(dummy_pt_objects))
30
+ else:
31
+ _import_structure["modular_pipeline"] = [
32
+ "ModularPipelineBlocks",
33
+ "ModularPipeline",
34
+ "AutoPipelineBlocks",
35
+ "SequentialPipelineBlocks",
36
+ "ConditionalPipelineBlocks",
37
+ "LoopSequentialPipelineBlocks",
38
+ "PipelineState",
39
+ "BlockState",
40
+ ]
41
+ _import_structure["modular_pipeline_utils"] = [
42
+ "ComponentSpec",
43
+ "ConfigSpec",
44
+ "InputParam",
45
+ "OutputParam",
46
+ "InsertableDict",
47
+ ]
48
+ _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"]
49
+ _import_structure["stable_diffusion_3"] = ["StableDiffusion3AutoBlocks", "StableDiffusion3ModularPipeline"]
50
+ _import_structure["wan"] = [
51
+ "WanBlocks",
52
+ "Wan22Blocks",
53
+ "WanImage2VideoAutoBlocks",
54
+ "Wan22Image2VideoBlocks",
55
+ "WanModularPipeline",
56
+ "Wan22ModularPipeline",
57
+ "WanImage2VideoModularPipeline",
58
+ "Wan22Image2VideoModularPipeline",
59
+ ]
60
+ _import_structure["helios"] = [
61
+ "HeliosAutoBlocks",
62
+ "HeliosModularPipeline",
63
+ "HeliosPyramidAutoBlocks",
64
+ "HeliosPyramidDistilledAutoBlocks",
65
+ "HeliosPyramidDistilledModularPipeline",
66
+ "HeliosPyramidModularPipeline",
67
+ ]
68
+ _import_structure["flux"] = [
69
+ "FluxAutoBlocks",
70
+ "FluxModularPipeline",
71
+ "FluxKontextAutoBlocks",
72
+ "FluxKontextModularPipeline",
73
+ ]
74
+ _import_structure["flux2"] = [
75
+ "Flux2AutoBlocks",
76
+ "Flux2KleinAutoBlocks",
77
+ "Flux2KleinBaseAutoBlocks",
78
+ "Flux2ModularPipeline",
79
+ "Flux2KleinModularPipeline",
80
+ "Flux2KleinBaseModularPipeline",
81
+ ]
82
+ _import_structure["ideogram4"] = [
83
+ "Ideogram4AutoBlocks",
84
+ "Ideogram4ModularPipeline",
85
+ ]
86
+ _import_structure["krea2"] = [
87
+ "Krea2AutoBlocks",
88
+ "Krea2ModularPipeline",
89
+ "Krea2TurboAutoBlocks",
90
+ "Krea2TurboModularPipeline",
91
+ ]
92
+ _import_structure["qwenimage"] = [
93
+ "QwenImageAutoBlocks",
94
+ "QwenImageModularPipeline",
95
+ "QwenImageEditModularPipeline",
96
+ "QwenImageEditAutoBlocks",
97
+ "QwenImageEditPlusModularPipeline",
98
+ "QwenImageEditPlusAutoBlocks",
99
+ "QwenImageLayeredModularPipeline",
100
+ "QwenImageLayeredAutoBlocks",
101
+ ]
102
+ _import_structure["anima"] = [
103
+ "AnimaAutoBlocks",
104
+ "AnimaModularPipeline",
105
+ ]
106
+ _import_structure["cosmos"] = [
107
+ "Cosmos3DistilledBlocks",
108
+ "Cosmos3DistilledModularPipeline",
109
+ "Cosmos3OmniBlocks",
110
+ "Cosmos3OmniModularPipeline",
111
+ ]
112
+ _import_structure["ernie_image"] = [
113
+ "ErnieImageAutoBlocks",
114
+ "ErnieImageModularPipeline",
115
+ ]
116
+ _import_structure["hunyuan_video1_5"] = [
117
+ "HunyuanVideo15AutoBlocks",
118
+ "HunyuanVideo15ModularPipeline",
119
+ ]
120
+ _import_structure["ltx"] = [
121
+ "LTXAutoBlocks",
122
+ "LTXModularPipeline",
123
+ ]
124
+ _import_structure["minimax_h3"] = [
125
+ "MiniMaxH3Blocks",
126
+ "MiniMaxH3ModularPipeline",
127
+ ]
128
+ _import_structure["z_image"] = [
129
+ "ZImageAutoBlocks",
130
+ "ZImageModularPipeline",
131
+ ]
132
+ _import_structure["components_manager"] = ["ComponentsManager"]
133
+
134
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
135
+ try:
136
+ if not is_torch_available():
137
+ raise OptionalDependencyNotAvailable()
138
+ except OptionalDependencyNotAvailable:
139
+ from ..utils.dummy_pt_objects import * # noqa F403
140
+ else:
141
+ from .anima import AnimaAutoBlocks, AnimaModularPipeline
142
+ from .components_manager import ComponentsManager
143
+ from .cosmos import (
144
+ Cosmos3DistilledBlocks,
145
+ Cosmos3DistilledModularPipeline,
146
+ Cosmos3OmniBlocks,
147
+ Cosmos3OmniModularPipeline,
148
+ )
149
+ from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline
150
+ from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
151
+ from .flux2 import (
152
+ Flux2AutoBlocks,
153
+ Flux2KleinAutoBlocks,
154
+ Flux2KleinBaseAutoBlocks,
155
+ Flux2KleinBaseModularPipeline,
156
+ Flux2KleinModularPipeline,
157
+ Flux2ModularPipeline,
158
+ )
159
+ from .helios import (
160
+ HeliosAutoBlocks,
161
+ HeliosModularPipeline,
162
+ HeliosPyramidAutoBlocks,
163
+ HeliosPyramidDistilledAutoBlocks,
164
+ HeliosPyramidDistilledModularPipeline,
165
+ HeliosPyramidModularPipeline,
166
+ )
167
+ from .hunyuan_video1_5 import (
168
+ HunyuanVideo15AutoBlocks,
169
+ HunyuanVideo15ModularPipeline,
170
+ )
171
+ from .ideogram4 import (
172
+ Ideogram4AutoBlocks,
173
+ Ideogram4ModularPipeline,
174
+ )
175
+ from .krea2 import (
176
+ Krea2AutoBlocks,
177
+ Krea2ModularPipeline,
178
+ Krea2TurboAutoBlocks,
179
+ Krea2TurboModularPipeline,
180
+ )
181
+ from .ltx import LTXAutoBlocks, LTXModularPipeline
182
+ from .minimax_h3 import (
183
+ MiniMaxH3Blocks,
184
+ MiniMaxH3ModularPipeline,
185
+ )
186
+ from .modular_pipeline import (
187
+ AutoPipelineBlocks,
188
+ BlockState,
189
+ ConditionalPipelineBlocks,
190
+ LoopSequentialPipelineBlocks,
191
+ ModularPipeline,
192
+ ModularPipelineBlocks,
193
+ PipelineState,
194
+ SequentialPipelineBlocks,
195
+ )
196
+ from .modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam
197
+ from .qwenimage import (
198
+ QwenImageAutoBlocks,
199
+ QwenImageEditAutoBlocks,
200
+ QwenImageEditModularPipeline,
201
+ QwenImageEditPlusAutoBlocks,
202
+ QwenImageEditPlusModularPipeline,
203
+ QwenImageLayeredAutoBlocks,
204
+ QwenImageLayeredModularPipeline,
205
+ QwenImageModularPipeline,
206
+ )
207
+ from .stable_diffusion_3 import StableDiffusion3AutoBlocks, StableDiffusion3ModularPipeline
208
+ from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline
209
+ from .wan import (
210
+ Wan22Blocks,
211
+ Wan22Image2VideoBlocks,
212
+ Wan22Image2VideoModularPipeline,
213
+ Wan22ModularPipeline,
214
+ WanBlocks,
215
+ WanImage2VideoAutoBlocks,
216
+ WanImage2VideoModularPipeline,
217
+ WanModularPipeline,
218
+ )
219
+ from .z_image import ZImageAutoBlocks, ZImageModularPipeline
220
+ else:
221
+ import sys
222
+
223
+ sys.modules[__name__] = _LazyModule(
224
+ __name__,
225
+ globals()["__file__"],
226
+ _import_structure,
227
+ module_spec=__spec__,
228
+ )
229
+ for name, value in _dummy_objects.items():
230
+ setattr(sys.modules[__name__], name, value)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ...utils import (
4
+ DIFFUSERS_SLOW_IMPORT,
5
+ OptionalDependencyNotAvailable,
6
+ _LazyModule,
7
+ get_objects_from_module,
8
+ is_torch_available,
9
+ is_transformers_available,
10
+ )
11
+
12
+
13
+ _dummy_objects = {}
14
+ _import_structure = {}
15
+
16
+ try:
17
+ if not (is_transformers_available() and is_torch_available()):
18
+ raise OptionalDependencyNotAvailable()
19
+ except OptionalDependencyNotAvailable:
20
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
21
+
22
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
23
+ else:
24
+ _import_structure["modular_blocks_anima"] = ["AnimaAutoBlocks"]
25
+ _import_structure["modular_pipeline"] = ["AnimaModularPipeline"]
26
+
27
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
28
+ try:
29
+ if not (is_transformers_available() and is_torch_available()):
30
+ raise OptionalDependencyNotAvailable()
31
+ except OptionalDependencyNotAvailable:
32
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
33
+ else:
34
+ from .modular_blocks_anima import AnimaAutoBlocks
35
+ from .modular_pipeline import AnimaModularPipeline
36
+ else:
37
+ import sys
38
+
39
+ sys.modules[__name__] = _LazyModule(
40
+ __name__,
41
+ globals()["__file__"],
42
+ _import_structure,
43
+ module_spec=__spec__,
44
+ )
45
+
46
+ for name, value in _dummy_objects.items():
47
+ setattr(sys.modules[__name__], name, value)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py ADDED
@@ -0,0 +1,714 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import inspect
16
+
17
+ import numpy as np
18
+ import torch
19
+
20
+ from ...models import AnimaTextConditioner, CosmosTransformer3DModel
21
+ from ...schedulers import FlowMatchEulerDiscreteScheduler
22
+ from ...utils.torch_utils import randn_tensor
23
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
24
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
25
+ from .modular_pipeline import AnimaModularPipeline
26
+
27
+
28
+ def retrieve_timesteps(
29
+ scheduler,
30
+ num_inference_steps: int | None = None,
31
+ device: str | torch.device | None = None,
32
+ timesteps: list[int] | None = None,
33
+ sigmas: list[float] | None = None,
34
+ **kwargs,
35
+ ):
36
+ if timesteps is not None and sigmas is not None:
37
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
38
+ if timesteps is not None:
39
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
40
+ if not accepts_timesteps:
41
+ raise ValueError(
42
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
43
+ f" timestep schedules. Please check whether you are using the correct scheduler."
44
+ )
45
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
46
+ timesteps = scheduler.timesteps
47
+ num_inference_steps = len(timesteps)
48
+ elif sigmas is not None:
49
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
50
+ if not accept_sigmas:
51
+ raise ValueError(
52
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
53
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
54
+ )
55
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
56
+ timesteps = scheduler.timesteps
57
+ num_inference_steps = len(timesteps)
58
+ else:
59
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
60
+ timesteps = scheduler.timesteps
61
+ return timesteps, num_inference_steps
62
+
63
+
64
+ # Copied from diffusers.modular_pipelines.z_image.before_denoise.repeat_tensor_to_batch_size
65
+ def repeat_tensor_to_batch_size(
66
+ input_name: str,
67
+ input_tensor: torch.Tensor,
68
+ batch_size: int,
69
+ num_images_per_prompt: int = 1,
70
+ ) -> torch.Tensor:
71
+ """Repeat tensor elements to match the final batch size.
72
+
73
+ This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
74
+ by repeating each element along dimension 0.
75
+
76
+ The input tensor must have batch size 1 or batch_size. The function will:
77
+ - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
78
+ - If batch size equals batch_size: repeat each element num_images_per_prompt times
79
+
80
+ Args:
81
+ input_name (str): Name of the input tensor (used for error messages)
82
+ input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
83
+ batch_size (int): The base batch size (number of prompts)
84
+ num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
85
+
86
+ Returns:
87
+ torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
88
+
89
+ Raises:
90
+ ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
91
+
92
+ Examples:
93
+ tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
94
+ batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
95
+ [4, 3]
96
+
97
+ tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
98
+ tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
99
+ - shape: [4, 3]
100
+ """
101
+ # make sure input is a tensor
102
+ if not isinstance(input_tensor, torch.Tensor):
103
+ raise ValueError(f"`{input_name}` must be a tensor")
104
+
105
+ # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
106
+ if input_tensor.shape[0] == 1:
107
+ repeat_by = batch_size * num_images_per_prompt
108
+ elif input_tensor.shape[0] == batch_size:
109
+ repeat_by = num_images_per_prompt
110
+ else:
111
+ raise ValueError(
112
+ f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
113
+ )
114
+
115
+ # expand the tensor to match the batch_size * num_images_per_prompt
116
+ input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
117
+
118
+ return input_tensor
119
+
120
+
121
+ class AnimaTextConditioningStep(ModularPipelineBlocks):
122
+ model_name = "anima"
123
+
124
+ @property
125
+ def description(self) -> str:
126
+ return "Map Qwen text encoder states and T5 token ids to Cosmos text conditioning for Anima."
127
+
128
+ @property
129
+ def expected_components(self) -> list[ComponentSpec]:
130
+ return [
131
+ ComponentSpec("text_conditioner", AnimaTextConditioner),
132
+ ComponentSpec("transformer", CosmosTransformer3DModel),
133
+ ]
134
+
135
+ @property
136
+ def inputs(self) -> list[InputParam]:
137
+ return [
138
+ InputParam(
139
+ "qwen_prompt_embeds",
140
+ required=True,
141
+ type_hint=torch.Tensor,
142
+ description="Qwen prompt embeddings generated by the text encoder step.",
143
+ ),
144
+ InputParam(
145
+ "qwen_attention_mask",
146
+ required=True,
147
+ type_hint=torch.Tensor,
148
+ description="Qwen prompt attention mask generated by the text encoder step.",
149
+ ),
150
+ InputParam(
151
+ "t5_input_ids",
152
+ required=True,
153
+ type_hint=torch.Tensor,
154
+ description="T5 prompt token ids generated by the text encoder step.",
155
+ ),
156
+ InputParam(
157
+ "t5_attention_mask",
158
+ required=True,
159
+ type_hint=torch.Tensor,
160
+ description="T5 prompt attention mask generated by the text encoder step.",
161
+ ),
162
+ InputParam(
163
+ "negative_qwen_prompt_embeds",
164
+ type_hint=torch.Tensor,
165
+ description="Negative Qwen prompt embeddings generated by the text encoder step.",
166
+ ),
167
+ InputParam(
168
+ "negative_qwen_attention_mask",
169
+ type_hint=torch.Tensor,
170
+ description="Negative Qwen prompt attention mask generated by the text encoder step.",
171
+ ),
172
+ InputParam(
173
+ "negative_t5_input_ids",
174
+ type_hint=torch.Tensor,
175
+ description="Negative T5 prompt token ids generated by the text encoder step.",
176
+ ),
177
+ InputParam(
178
+ "negative_t5_attention_mask",
179
+ type_hint=torch.Tensor,
180
+ description="Negative T5 prompt attention mask generated by the text encoder step.",
181
+ ),
182
+ ]
183
+
184
+ @property
185
+ def intermediate_outputs(self) -> list[OutputParam]:
186
+ return [
187
+ OutputParam(
188
+ "prompt_embeds",
189
+ type_hint=torch.Tensor,
190
+ description="Conditioned prompt embeddings generated by the Anima text conditioner.",
191
+ ),
192
+ OutputParam(
193
+ "negative_prompt_embeds",
194
+ type_hint=torch.Tensor,
195
+ description="Conditioned negative prompt embeddings generated by the Anima text conditioner.",
196
+ ),
197
+ ]
198
+
199
+ @staticmethod
200
+ def _condition_prompt_embeds(
201
+ components: AnimaModularPipeline,
202
+ qwen_prompt_embeds: torch.Tensor,
203
+ qwen_attention_mask: torch.Tensor,
204
+ t5_input_ids: torch.Tensor,
205
+ t5_attention_mask: torch.Tensor,
206
+ device: torch.device,
207
+ conditioning_dtype: torch.dtype,
208
+ output_dtype: torch.dtype,
209
+ ) -> torch.Tensor:
210
+ prompt_embeds = components.text_conditioner(
211
+ source_hidden_states=qwen_prompt_embeds.to(device=device, dtype=conditioning_dtype),
212
+ target_input_ids=t5_input_ids.to(device),
213
+ target_attention_mask=t5_attention_mask.to(device),
214
+ source_attention_mask=qwen_attention_mask.to(device),
215
+ )
216
+ return prompt_embeds.to(dtype=output_dtype, device=device)
217
+
218
+ @torch.no_grad()
219
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
220
+ block_state = self.get_block_state(state)
221
+ device = components._execution_device
222
+ conditioning_dtype = components.text_conditioner.dtype
223
+ output_dtype = components.transformer.dtype
224
+
225
+ block_state.prompt_embeds = self._condition_prompt_embeds(
226
+ components,
227
+ qwen_prompt_embeds=block_state.qwen_prompt_embeds,
228
+ qwen_attention_mask=block_state.qwen_attention_mask,
229
+ t5_input_ids=block_state.t5_input_ids,
230
+ t5_attention_mask=block_state.t5_attention_mask,
231
+ device=device,
232
+ conditioning_dtype=conditioning_dtype,
233
+ output_dtype=output_dtype,
234
+ )
235
+
236
+ block_state.negative_prompt_embeds = None
237
+ if block_state.negative_qwen_prompt_embeds is not None:
238
+ block_state.negative_prompt_embeds = self._condition_prompt_embeds(
239
+ components,
240
+ qwen_prompt_embeds=block_state.negative_qwen_prompt_embeds,
241
+ qwen_attention_mask=block_state.negative_qwen_attention_mask,
242
+ t5_input_ids=block_state.negative_t5_input_ids,
243
+ t5_attention_mask=block_state.negative_t5_attention_mask,
244
+ device=device,
245
+ conditioning_dtype=conditioning_dtype,
246
+ output_dtype=output_dtype,
247
+ )
248
+
249
+ self.set_block_state(state, block_state)
250
+ return components, state
251
+
252
+
253
+ class AnimaTextInputStep(ModularPipelineBlocks):
254
+ model_name = "anima"
255
+
256
+ @property
257
+ def description(self) -> str:
258
+ return "Input processing step that expands Anima prompt embeddings for the requested image batch."
259
+
260
+ @property
261
+ def expected_components(self) -> list[ComponentSpec]:
262
+ return [ComponentSpec("transformer", CosmosTransformer3DModel)]
263
+
264
+ @property
265
+ def inputs(self) -> list[InputParam]:
266
+ return [
267
+ InputParam.template("num_images_per_prompt"),
268
+ InputParam(
269
+ "prompt_embeds",
270
+ required=True,
271
+ type_hint=torch.Tensor,
272
+ description="Conditioned prompt embeddings generated by the Anima text conditioner.",
273
+ ),
274
+ InputParam(
275
+ "negative_prompt_embeds",
276
+ type_hint=torch.Tensor,
277
+ description="Conditioned negative prompt embeddings generated by the Anima text conditioner.",
278
+ ),
279
+ ]
280
+
281
+ @property
282
+ def intermediate_outputs(self) -> list[OutputParam]:
283
+ return [
284
+ OutputParam(
285
+ "prompt_embeds",
286
+ type_hint=torch.Tensor,
287
+ kwargs_type="denoiser_input_fields",
288
+ description="Prompt embeddings expanded to the final denoising batch.",
289
+ ),
290
+ OutputParam(
291
+ "negative_prompt_embeds",
292
+ type_hint=torch.Tensor,
293
+ kwargs_type="denoiser_input_fields",
294
+ description="Negative prompt embeddings expanded to the final denoising batch.",
295
+ ),
296
+ OutputParam(
297
+ "batch_size",
298
+ type_hint=int,
299
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
300
+ ),
301
+ OutputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
302
+ ]
303
+
304
+ @torch.no_grad()
305
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
306
+ block_state = self.get_block_state(state)
307
+
308
+ block_state.batch_size = block_state.prompt_embeds.shape[0]
309
+ block_state.dtype = components.transformer.dtype
310
+
311
+ _, seq_len, _ = block_state.prompt_embeds.shape
312
+ block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1)
313
+ block_state.prompt_embeds = block_state.prompt_embeds.view(
314
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
315
+ )
316
+
317
+ if block_state.negative_prompt_embeds is not None:
318
+ _, seq_len, _ = block_state.negative_prompt_embeds.shape
319
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
320
+ 1, block_state.num_images_per_prompt, 1
321
+ )
322
+ block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
323
+ block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1
324
+ )
325
+
326
+ self.set_block_state(state, block_state)
327
+ return components, state
328
+
329
+
330
+ class AnimaImageInputStep(ModularPipelineBlocks):
331
+ model_name = "anima"
332
+
333
+ @property
334
+ def description(self) -> str:
335
+ return (
336
+ "Input processing step that expands Anima image latents to the final denoising batch "
337
+ "and derives height/width from the latents when not provided."
338
+ )
339
+
340
+ @property
341
+ def inputs(self) -> list[InputParam]:
342
+ return [
343
+ InputParam.template("image_latents"),
344
+ InputParam(
345
+ "batch_size",
346
+ required=True,
347
+ type_hint=int,
348
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
349
+ ),
350
+ InputParam.template("num_images_per_prompt"),
351
+ InputParam.template("height"),
352
+ InputParam.template("width"),
353
+ ]
354
+
355
+ @property
356
+ def intermediate_outputs(self) -> list[OutputParam]:
357
+ return [
358
+ OutputParam(
359
+ "image_latents",
360
+ type_hint=torch.Tensor,
361
+ description="Image latents expanded to the final denoising batch.",
362
+ ),
363
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
364
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
365
+ ]
366
+
367
+ @torch.no_grad()
368
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
369
+ block_state = self.get_block_state(state)
370
+
371
+ latent_height, latent_width = block_state.image_latents.shape[-2:]
372
+ block_state.height = block_state.height or latent_height * components.vae_scale_factor
373
+ block_state.width = block_state.width or latent_width * components.vae_scale_factor
374
+
375
+ block_state.image_latents = repeat_tensor_to_batch_size(
376
+ input_name="image_latents",
377
+ input_tensor=block_state.image_latents,
378
+ batch_size=block_state.batch_size,
379
+ num_images_per_prompt=block_state.num_images_per_prompt,
380
+ )
381
+
382
+ self.set_block_state(state, block_state)
383
+ return components, state
384
+
385
+
386
+ class AnimaPrepareLatentsStep(ModularPipelineBlocks):
387
+ model_name = "anima"
388
+
389
+ @property
390
+ def description(self) -> str:
391
+ return "Prepare noisy image latents and padding mask for Anima denoising."
392
+
393
+ @property
394
+ def expected_components(self) -> list[ComponentSpec]:
395
+ return [ComponentSpec("transformer", CosmosTransformer3DModel)]
396
+
397
+ @property
398
+ def inputs(self) -> list[InputParam]:
399
+ return [
400
+ InputParam.template("height"),
401
+ InputParam.template("width"),
402
+ InputParam.template("latents"),
403
+ InputParam.template("num_images_per_prompt"),
404
+ InputParam.template("generator"),
405
+ InputParam(
406
+ "batch_size",
407
+ required=True,
408
+ type_hint=int,
409
+ description="Number of input prompts before `num_images_per_prompt` expansion.",
410
+ ),
411
+ InputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
412
+ ]
413
+
414
+ @property
415
+ def intermediate_outputs(self) -> list[OutputParam]:
416
+ return [
417
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
418
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
419
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy latents for the denoising process."),
420
+ OutputParam("padding_mask", type_hint=torch.Tensor, description="Cosmos padding mask for image latents."),
421
+ ]
422
+
423
+ def check_inputs(self, components: AnimaModularPipeline, block_state):
424
+ divisor = components.vae_scale_factor * 2
425
+ if block_state.height % divisor != 0 or block_state.width % divisor != 0:
426
+ raise ValueError(
427
+ f"`height` and `width` have to be divisible by {divisor} but are {block_state.height} and"
428
+ f" {block_state.width}."
429
+ )
430
+
431
+ @staticmethod
432
+ def prepare_latents(
433
+ batch_size: int,
434
+ num_channels_latents: int,
435
+ height: int,
436
+ width: int,
437
+ vae_scale_factor: int,
438
+ dtype: torch.dtype,
439
+ device: torch.device,
440
+ generator: torch.Generator | list[torch.Generator] | None,
441
+ latents: torch.Tensor | None = None,
442
+ ) -> torch.Tensor:
443
+ if latents is not None:
444
+ return latents.to(device=device, dtype=dtype)
445
+
446
+ latent_height = height // vae_scale_factor
447
+ latent_width = width // vae_scale_factor
448
+ shape = (batch_size, num_channels_latents, 1, latent_height, latent_width)
449
+
450
+ if isinstance(generator, list) and len(generator) != batch_size:
451
+ raise ValueError(
452
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
453
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
454
+ )
455
+
456
+ return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
457
+
458
+ @torch.no_grad()
459
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
460
+ block_state = self.get_block_state(state)
461
+
462
+ block_state.height = block_state.height or components.default_height
463
+ block_state.width = block_state.width or components.default_width
464
+ self.check_inputs(components, block_state)
465
+
466
+ device = components._execution_device
467
+ block_state.latents = self.prepare_latents(
468
+ batch_size=block_state.batch_size * block_state.num_images_per_prompt,
469
+ num_channels_latents=components.num_channels_latents,
470
+ height=block_state.height,
471
+ width=block_state.width,
472
+ vae_scale_factor=components.vae_scale_factor,
473
+ dtype=torch.float32,
474
+ device=device,
475
+ generator=block_state.generator,
476
+ latents=block_state.latents,
477
+ )
478
+ block_state.padding_mask = block_state.latents.new_zeros(
479
+ 1, 1, block_state.height, block_state.width, dtype=block_state.dtype
480
+ )
481
+
482
+ self.set_block_state(state, block_state)
483
+ return components, state
484
+
485
+
486
+ # Copied from diffusers.modular_pipelines.qwenimage.before_denoise.get_timesteps
487
+ def get_timesteps(scheduler, num_inference_steps, strength):
488
+ # get the original timestep using init_timestep
489
+ init_timestep = min(num_inference_steps * strength, num_inference_steps)
490
+
491
+ t_start = int(max(num_inference_steps - init_timestep, 0))
492
+ timesteps = scheduler.timesteps[t_start * scheduler.order :]
493
+ if hasattr(scheduler, "set_begin_index"):
494
+ scheduler.set_begin_index(t_start * scheduler.order)
495
+
496
+ return timesteps, num_inference_steps - t_start
497
+
498
+
499
+ class AnimaSetTimestepsStep(ModularPipelineBlocks):
500
+ model_name = "anima"
501
+
502
+ @property
503
+ def expected_components(self) -> list[ComponentSpec]:
504
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
505
+
506
+ @property
507
+ def description(self) -> str:
508
+ return "Set the scheduler timesteps for Anima inference."
509
+
510
+ @property
511
+ def inputs(self) -> list[InputParam]:
512
+ return [
513
+ InputParam.template("num_inference_steps"),
514
+ InputParam.template("sigmas"),
515
+ ]
516
+
517
+ @property
518
+ def intermediate_outputs(self) -> list[OutputParam]:
519
+ return [
520
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for the denoising loop."),
521
+ OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps."),
522
+ ]
523
+
524
+ @torch.no_grad()
525
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
526
+ block_state = self.get_block_state(state)
527
+ device = components._execution_device
528
+
529
+ sigmas = (
530
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
531
+ if block_state.sigmas is None
532
+ else block_state.sigmas
533
+ )
534
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
535
+ components.scheduler,
536
+ device=device,
537
+ sigmas=sigmas,
538
+ )
539
+ components.scheduler.set_begin_index(0)
540
+
541
+ self.set_block_state(state, block_state)
542
+ return components, state
543
+
544
+
545
+ class AnimaImg2ImgSetTimestepsStep(ModularPipelineBlocks):
546
+ """Set the scheduler timesteps for Anima image-to-image inference.
547
+
548
+ This step computes the full timestep schedule, then slices it based on ``strength`` via ``get_timesteps()``, which
549
+ also sets the scheduler's begin index.
550
+
551
+ Components:
552
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
553
+
554
+ Inputs:
555
+ num_inference_steps (`int`, *optional*, defaults to 50):
556
+ The number of denoising steps.
557
+ sigmas (`list`, *optional*):
558
+ Custom sigmas for the denoising process.
559
+ strength (`float`, *optional*, defaults to 0.9):
560
+ How much to transform the reference image.
561
+
562
+ Outputs:
563
+ timesteps (`Tensor`):
564
+ Timestep schedule sliced by ``strength``.
565
+ num_inference_steps (`int`):
566
+ Number of denoising steps after strength-based slicing.
567
+ """
568
+
569
+ model_name = "anima"
570
+
571
+ @property
572
+ def expected_components(self) -> list[ComponentSpec]:
573
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
574
+
575
+ @property
576
+ def description(self) -> str:
577
+ return "Set the scheduler timesteps for Anima image-to-image inference, sliced by strength."
578
+
579
+ @property
580
+ def inputs(self) -> list[InputParam]:
581
+ return [
582
+ InputParam.template("num_inference_steps"),
583
+ InputParam.template("sigmas"),
584
+ InputParam.template("strength"),
585
+ ]
586
+
587
+ @property
588
+ def intermediate_outputs(self) -> list[OutputParam]:
589
+ return [
590
+ OutputParam(
591
+ "timesteps",
592
+ type_hint=torch.Tensor,
593
+ description="Timestep schedule sliced by strength.",
594
+ ),
595
+ OutputParam(
596
+ "num_inference_steps",
597
+ type_hint=int,
598
+ description="Number of denoising steps after strength-based slicing.",
599
+ ),
600
+ ]
601
+
602
+ @torch.no_grad()
603
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
604
+ block_state = self.get_block_state(state)
605
+ device = components._execution_device
606
+
607
+ sigmas = (
608
+ np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
609
+ if block_state.sigmas is None
610
+ else block_state.sigmas
611
+ )
612
+ block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
613
+ components.scheduler,
614
+ device=device,
615
+ sigmas=sigmas,
616
+ )
617
+ block_state.timesteps, block_state.num_inference_steps = get_timesteps(
618
+ components.scheduler, block_state.num_inference_steps, block_state.strength
619
+ )
620
+
621
+ self.set_block_state(state, block_state)
622
+ return components, state
623
+
624
+
625
+ class AnimaImg2ImgPrepareLatentsStep(ModularPipelineBlocks):
626
+ """Prepares noisy latents for Anima image-to-image generation.
627
+
628
+ Generates noise and mixes it with the image latents via ``scheduler.scale_noise()`` at the first sliced timestep.
629
+ The image latents are expected to already be expanded to the final batch size by ``AnimaImageInputStep``.
630
+
631
+ Components:
632
+ scheduler (`FlowMatchEulerDiscreteScheduler`)
633
+
634
+ Inputs:
635
+ image_latents (`Tensor`):
636
+ Encoded image latents, expanded to the final denoising batch.
637
+ timesteps (`Tensor`):
638
+ Timestep schedule sliced by ``strength`` from ``AnimaImg2ImgSetTimestepsStep``.
639
+ generator (`Generator`, *optional*):
640
+ Torch generator for deterministic generation.
641
+ latents (`Tensor`, *optional*):
642
+ Pre-computed noise tensor. Generated randomly if ``None``.
643
+ dtype (`torch.dtype`):
644
+ Dtype used by the Anima denoiser.
645
+ height (`int`):
646
+ Image height.
647
+ width (`int`):
648
+ Image width.
649
+
650
+ Outputs:
651
+ latents (`Tensor`):
652
+ Noisy image latents for the denoising loop.
653
+ padding_mask (`Tensor`):
654
+ Cosmos padding mask for the image latents.
655
+ """
656
+
657
+ model_name = "anima"
658
+
659
+ @property
660
+ def expected_components(self) -> list[ComponentSpec]:
661
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
662
+
663
+ @property
664
+ def description(self) -> str:
665
+ return (
666
+ "Prepares noisy image-to-image latents for Anima by adding noise to the encoded "
667
+ "image latents via scheduler.scale_noise()."
668
+ )
669
+
670
+ @property
671
+ def inputs(self) -> list[InputParam]:
672
+ return [
673
+ InputParam.template("image_latents"),
674
+ InputParam.template("timesteps", required=True),
675
+ InputParam.template("generator"),
676
+ InputParam.template("latents"),
677
+ InputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
678
+ InputParam.template("height"),
679
+ InputParam.template("width"),
680
+ ]
681
+
682
+ @property
683
+ def intermediate_outputs(self) -> list[OutputParam]:
684
+ return [
685
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy latents for the denoising loop."),
686
+ OutputParam("padding_mask", type_hint=torch.Tensor, description="Cosmos padding mask for image latents."),
687
+ ]
688
+
689
+ @torch.no_grad()
690
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
691
+ block_state = self.get_block_state(state)
692
+
693
+ device = components._execution_device
694
+ image_latents = block_state.image_latents.to(device=device, dtype=torch.float32)
695
+
696
+ if block_state.latents is None:
697
+ noise = randn_tensor(
698
+ image_latents.shape,
699
+ generator=block_state.generator,
700
+ device=device,
701
+ dtype=torch.float32,
702
+ )
703
+ else:
704
+ noise = block_state.latents.to(device=device, dtype=torch.float32)
705
+
706
+ latent_timestep = block_state.timesteps[:1].repeat(image_latents.shape[0])
707
+ block_state.latents = components.scheduler.scale_noise(image_latents, latent_timestep, noise)
708
+
709
+ block_state.padding_mask = block_state.latents.new_zeros(
710
+ 1, 1, block_state.height, block_state.width, dtype=block_state.dtype
711
+ )
712
+
713
+ self.set_block_state(state, block_state)
714
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import numpy as np
16
+ import PIL
17
+ import torch
18
+
19
+ from ...configuration_utils import FrozenDict
20
+ from ...image_processor import VaeImageProcessor
21
+ from ...models import AutoencoderKLQwenImage
22
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
23
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
24
+ from .modular_pipeline import AnimaModularPipeline
25
+
26
+
27
+ class AnimaVaeDecoderStep(ModularPipelineBlocks):
28
+ model_name = "anima"
29
+
30
+ @property
31
+ def description(self) -> str:
32
+ return "Step that decodes Anima latents into image tensors."
33
+
34
+ @property
35
+ def expected_components(self) -> list[ComponentSpec]:
36
+ return [ComponentSpec("vae", AutoencoderKLQwenImage)]
37
+
38
+ @property
39
+ def inputs(self) -> list[InputParam]:
40
+ return [
41
+ InputParam("latents", required=True, type_hint=torch.Tensor, description="Denoised Anima latents."),
42
+ ]
43
+
44
+ @property
45
+ def intermediate_outputs(self) -> list[OutputParam]:
46
+ return [OutputParam.template("images", note="tensor output of the VAE decoder")]
47
+
48
+ @torch.no_grad()
49
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
50
+ block_state = self.get_block_state(state)
51
+
52
+ latents = block_state.latents.to(components.vae.dtype)
53
+ latents_mean = (
54
+ torch.tensor(components.vae.config.latents_mean)
55
+ .view(1, components.vae.config.z_dim, 1, 1, 1)
56
+ .to(latents.device, latents.dtype)
57
+ )
58
+ latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
59
+ 1, components.vae.config.z_dim, 1, 1, 1
60
+ ).to(latents.device, latents.dtype)
61
+ latents = latents / latents_std + latents_mean
62
+
63
+ block_state.images = components.vae.decode(latents, return_dict=False)[0][:, :, 0]
64
+
65
+ self.set_block_state(state, block_state)
66
+ return components, state
67
+
68
+
69
+ class AnimaProcessImagesOutputStep(ModularPipelineBlocks):
70
+ model_name = "anima"
71
+
72
+ @property
73
+ def description(self) -> str:
74
+ return "Postprocess decoded Anima image tensors."
75
+
76
+ @property
77
+ def expected_components(self) -> list[ComponentSpec]:
78
+ return [
79
+ ComponentSpec(
80
+ "image_processor",
81
+ VaeImageProcessor,
82
+ config=FrozenDict({"vae_scale_factor": 8}),
83
+ default_creation_method="from_config",
84
+ ),
85
+ ]
86
+
87
+ @property
88
+ def inputs(self) -> list[InputParam]:
89
+ return [
90
+ InputParam("images", required=True, type_hint=torch.Tensor, description="Decoded Anima image tensors."),
91
+ InputParam.template("output_type"),
92
+ ]
93
+
94
+ @property
95
+ def intermediate_outputs(self) -> list[OutputParam]:
96
+ return [
97
+ OutputParam(
98
+ "images",
99
+ type_hint=list[PIL.Image.Image] | np.ndarray | torch.Tensor,
100
+ description="Generated images.",
101
+ )
102
+ ]
103
+
104
+ @staticmethod
105
+ def check_inputs(output_type):
106
+ if output_type not in ["pil", "np", "pt"]:
107
+ raise ValueError(f"Invalid output_type: {output_type}")
108
+
109
+ @torch.no_grad()
110
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
111
+ block_state = self.get_block_state(state)
112
+ self.check_inputs(block_state.output_type)
113
+
114
+ block_state.images = components.image_processor.postprocess(
115
+ image=block_state.images,
116
+ output_type=block_state.output_type,
117
+ )
118
+
119
+ self.set_block_state(state, block_state)
120
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py ADDED
@@ -0,0 +1,211 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from typing import Any
16
+
17
+ import torch
18
+
19
+ from ...configuration_utils import FrozenDict
20
+ from ...guiders import ClassifierFreeGuidance
21
+ from ...models import CosmosTransformer3DModel
22
+ from ...schedulers import FlowMatchEulerDiscreteScheduler
23
+ from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
24
+ from ..modular_pipeline_utils import ComponentSpec, InputParam
25
+ from .modular_pipeline import AnimaModularPipeline
26
+
27
+
28
+ class AnimaLoopBeforeDenoiser(ModularPipelineBlocks):
29
+ model_name = "anima"
30
+
31
+ @property
32
+ def description(self) -> str:
33
+ return "Step within the denoising loop that prepares Anima latent and timestep inputs."
34
+
35
+ @property
36
+ def inputs(self) -> list[InputParam]:
37
+ return [
38
+ InputParam("latents", required=True, type_hint=torch.Tensor, description="Current Anima latents."),
39
+ InputParam("dtype", required=True, type_hint=torch.dtype, description="Dtype used by the Anima denoiser."),
40
+ ]
41
+
42
+ @torch.no_grad()
43
+ def __call__(self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
44
+ block_state.latent_model_input = block_state.latents.to(block_state.dtype)
45
+
46
+ timestep = t.expand(block_state.latents.shape[0]).to(block_state.dtype)
47
+ block_state.timestep = timestep / components.scheduler.config.num_train_timesteps
48
+ return components, block_state
49
+
50
+
51
+ class AnimaLoopDenoiser(ModularPipelineBlocks):
52
+ model_name = "anima"
53
+
54
+ def __init__(
55
+ self,
56
+ guider_input_fields: dict[str, Any] | None = None,
57
+ ):
58
+ if guider_input_fields is None:
59
+ guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}
60
+ if not isinstance(guider_input_fields, dict):
61
+ raise ValueError(f"`guider_input_fields` must be a dictionary but is {type(guider_input_fields)}")
62
+ self._guider_input_fields = guider_input_fields
63
+ super().__init__()
64
+
65
+ @property
66
+ def expected_components(self) -> list[ComponentSpec]:
67
+ return [
68
+ ComponentSpec(
69
+ "guider",
70
+ ClassifierFreeGuidance,
71
+ config=FrozenDict({"guidance_scale": 4.0}),
72
+ default_creation_method="from_config",
73
+ ),
74
+ ComponentSpec("transformer", CosmosTransformer3DModel),
75
+ ]
76
+
77
+ @property
78
+ def description(self) -> str:
79
+ return "Step within the denoising loop that predicts Anima noise with guidance."
80
+
81
+ @property
82
+ def inputs(self) -> list[InputParam]:
83
+ inputs = [
84
+ InputParam(
85
+ "num_inference_steps",
86
+ required=True,
87
+ type_hint=int,
88
+ description="Number of denoising steps.",
89
+ ),
90
+ InputParam(
91
+ "padding_mask",
92
+ required=True,
93
+ type_hint=torch.Tensor,
94
+ description="Cosmos padding mask for image latents.",
95
+ ),
96
+ InputParam(
97
+ kwargs_type="denoiser_input_fields",
98
+ description="The conditional model inputs for the Anima denoiser.",
99
+ ),
100
+ ]
101
+
102
+ guider_input_names = []
103
+ uncond_guider_input_names = []
104
+ for value in self._guider_input_fields.values():
105
+ if isinstance(value, tuple):
106
+ guider_input_names.append(value[0])
107
+ uncond_guider_input_names.append(value[1])
108
+ else:
109
+ guider_input_names.append(value)
110
+
111
+ for name in guider_input_names:
112
+ inputs.append(InputParam(name=name, required=True))
113
+ for name in uncond_guider_input_names:
114
+ inputs.append(InputParam(name=name))
115
+ return inputs
116
+
117
+ @torch.no_grad()
118
+ def __call__(
119
+ self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
120
+ ) -> PipelineState:
121
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
122
+ guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
123
+
124
+ for guider_state_batch in guider_state:
125
+ components.guider.prepare_models(components.transformer)
126
+ cond_kwargs = {
127
+ key: getattr(guider_state_batch, key).to(block_state.dtype) for key in self._guider_input_fields.keys()
128
+ }
129
+ guider_state_batch.noise_pred = components.transformer(
130
+ hidden_states=block_state.latent_model_input,
131
+ timestep=block_state.timestep,
132
+ padding_mask=block_state.padding_mask,
133
+ return_dict=False,
134
+ **cond_kwargs,
135
+ )[0]
136
+ components.guider.cleanup_models(components.transformer)
137
+
138
+ block_state.noise_pred = components.guider(guider_state)[0]
139
+ return components, block_state
140
+
141
+
142
+ class AnimaLoopAfterDenoiser(ModularPipelineBlocks):
143
+ model_name = "anima"
144
+
145
+ @property
146
+ def expected_components(self) -> list[ComponentSpec]:
147
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
148
+
149
+ @property
150
+ def description(self) -> str:
151
+ return "Step within the denoising loop that updates Anima latents."
152
+
153
+ @torch.no_grad()
154
+ def __call__(self, components: AnimaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
155
+ latents_dtype = block_state.latents.dtype
156
+ block_state.latents = components.scheduler.step(
157
+ block_state.noise_pred, t, block_state.latents, return_dict=False
158
+ )[0]
159
+ if block_state.latents.dtype != latents_dtype and torch.backends.mps.is_available():
160
+ block_state.latents = block_state.latents.to(latents_dtype)
161
+
162
+ return components, block_state
163
+
164
+
165
+ class AnimaDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
166
+ model_name = "anima"
167
+
168
+ @property
169
+ def description(self) -> str:
170
+ return "Pipeline block that iteratively denoises Anima latents over scheduler timesteps."
171
+
172
+ @property
173
+ def loop_expected_components(self) -> list[ComponentSpec]:
174
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
175
+
176
+ @property
177
+ def loop_inputs(self) -> list[InputParam]:
178
+ return [
179
+ InputParam("timesteps", required=True, type_hint=torch.Tensor, description="Timesteps to denoise over."),
180
+ InputParam("num_inference_steps", required=True, type_hint=int, description="Number of denoising steps."),
181
+ ]
182
+
183
+ @torch.no_grad()
184
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
185
+ block_state = self.get_block_state(state)
186
+
187
+ num_warmup_steps = len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
188
+
189
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
190
+ for i, t in enumerate(block_state.timesteps):
191
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
192
+ if i == len(block_state.timesteps) - 1 or (
193
+ (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0
194
+ ):
195
+ progress_bar.update()
196
+
197
+ self.set_block_state(state, block_state)
198
+ return components, state
199
+
200
+
201
+ class AnimaDenoiseStep(AnimaDenoiseLoopWrapper):
202
+ block_classes = [
203
+ AnimaLoopBeforeDenoiser,
204
+ AnimaLoopDenoiser(guider_input_fields={"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}),
205
+ AnimaLoopAfterDenoiser,
206
+ ]
207
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
208
+
209
+ @property
210
+ def description(self) -> str:
211
+ return "Denoise step that iteratively denoises image latents for Anima."
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py ADDED
@@ -0,0 +1,404 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import torch
16
+ from transformers import Qwen2Tokenizer, Qwen3Model, T5TokenizerFast
17
+
18
+ from ...configuration_utils import FrozenDict
19
+ from ...guiders import ClassifierFreeGuidance
20
+ from ...image_processor import VaeImageProcessor
21
+ from ...models import AutoencoderKLQwenImage
22
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
23
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
24
+ from .modular_pipeline import AnimaModularPipeline
25
+
26
+
27
+ class AnimaTextEncoderStep(ModularPipelineBlocks):
28
+ model_name = "anima"
29
+
30
+ @property
31
+ def description(self) -> str:
32
+ return "Text encoder step that encodes Anima prompts into Qwen states and T5 token ids."
33
+
34
+ @property
35
+ def expected_components(self) -> list[ComponentSpec]:
36
+ return [
37
+ ComponentSpec("text_encoder", Qwen3Model),
38
+ ComponentSpec("tokenizer", Qwen2Tokenizer),
39
+ ComponentSpec("t5_tokenizer", T5TokenizerFast),
40
+ ComponentSpec(
41
+ "guider",
42
+ ClassifierFreeGuidance,
43
+ config=FrozenDict({"guidance_scale": 4.0}),
44
+ default_creation_method="from_config",
45
+ ),
46
+ ]
47
+
48
+ @property
49
+ def inputs(self) -> list[InputParam]:
50
+ return [
51
+ InputParam.template("prompt"),
52
+ InputParam.template("negative_prompt"),
53
+ InputParam.template("max_sequence_length"),
54
+ ]
55
+
56
+ @property
57
+ def intermediate_outputs(self) -> list[OutputParam]:
58
+ return [
59
+ OutputParam(
60
+ "qwen_prompt_embeds",
61
+ type_hint=torch.Tensor,
62
+ description="Qwen prompt embeddings to be consumed by the Anima text conditioner.",
63
+ ),
64
+ OutputParam(
65
+ "qwen_attention_mask",
66
+ type_hint=torch.Tensor,
67
+ description="Qwen prompt attention mask to be consumed by the Anima text conditioner.",
68
+ ),
69
+ OutputParam(
70
+ "t5_input_ids",
71
+ type_hint=torch.Tensor,
72
+ description="T5 prompt token ids to be consumed by the Anima text conditioner.",
73
+ ),
74
+ OutputParam(
75
+ "t5_attention_mask",
76
+ type_hint=torch.Tensor,
77
+ description="T5 prompt attention mask to be consumed by the Anima text conditioner.",
78
+ ),
79
+ OutputParam(
80
+ "negative_qwen_prompt_embeds",
81
+ type_hint=torch.Tensor,
82
+ description="Negative Qwen prompt embeddings to be consumed by the Anima text conditioner.",
83
+ ),
84
+ OutputParam(
85
+ "negative_qwen_attention_mask",
86
+ type_hint=torch.Tensor,
87
+ description="Negative Qwen prompt attention mask to be consumed by the Anima text conditioner.",
88
+ ),
89
+ OutputParam(
90
+ "negative_t5_input_ids",
91
+ type_hint=torch.Tensor,
92
+ description="Negative T5 prompt token ids to be consumed by the Anima text conditioner.",
93
+ ),
94
+ OutputParam(
95
+ "negative_t5_attention_mask",
96
+ type_hint=torch.Tensor,
97
+ description="Negative T5 prompt attention mask to be consumed by the Anima text conditioner.",
98
+ ),
99
+ ]
100
+
101
+ @staticmethod
102
+ def check_inputs(block_state):
103
+ if not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list):
104
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
105
+ if block_state.max_sequence_length is not None and block_state.max_sequence_length > 4096:
106
+ raise ValueError(
107
+ f"`max_sequence_length` cannot be greater than 4096 but is {block_state.max_sequence_length}"
108
+ )
109
+
110
+ @staticmethod
111
+ def _get_qwen_prompt_embeds(
112
+ components: AnimaModularPipeline,
113
+ prompt: str | list[str],
114
+ max_sequence_length: int,
115
+ device: torch.device,
116
+ dtype: torch.dtype,
117
+ ) -> tuple[torch.Tensor, torch.Tensor]:
118
+ prompt = [prompt] if isinstance(prompt, str) else prompt
119
+
120
+ text_inputs = components.tokenizer(
121
+ prompt,
122
+ padding="longest",
123
+ max_length=max_sequence_length,
124
+ truncation=True,
125
+ return_tensors="pt",
126
+ )
127
+ text_input_ids = text_inputs.input_ids.to(device)
128
+ prompt_attention_mask = text_inputs.attention_mask.to(device)
129
+ if text_input_ids.shape[-1] == 0:
130
+ text_input_ids = text_input_ids.new_zeros((text_input_ids.shape[0], 1))
131
+ prompt_attention_mask = prompt_attention_mask.new_zeros((prompt_attention_mask.shape[0], 1))
132
+
133
+ prompt_embeds = components.text_encoder(
134
+ input_ids=text_input_ids,
135
+ attention_mask=prompt_attention_mask,
136
+ output_hidden_states=False,
137
+ ).last_hidden_state
138
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
139
+ prompt_embeds = prompt_embeds * prompt_attention_mask.to(prompt_embeds).unsqueeze(-1)
140
+
141
+ return prompt_embeds, prompt_attention_mask
142
+
143
+ @staticmethod
144
+ def _get_t5_prompt_ids(
145
+ components: AnimaModularPipeline,
146
+ prompt: str | list[str],
147
+ max_sequence_length: int,
148
+ device: torch.device,
149
+ ) -> tuple[torch.Tensor, torch.Tensor]:
150
+ prompt = [prompt] if isinstance(prompt, str) else prompt
151
+
152
+ text_inputs = components.t5_tokenizer(
153
+ prompt,
154
+ padding="longest",
155
+ max_length=max_sequence_length,
156
+ truncation=True,
157
+ return_tensors="pt",
158
+ )
159
+ return text_inputs.input_ids.to(device), text_inputs.attention_mask.to(device)
160
+
161
+ @classmethod
162
+ def encode_prompt(
163
+ cls,
164
+ components: AnimaModularPipeline,
165
+ prompt: str | list[str],
166
+ negative_prompt: str | list[str] | None = None,
167
+ prepare_unconditional_embeds: bool = True,
168
+ max_sequence_length: int = 512,
169
+ device: torch.device | None = None,
170
+ dtype: torch.dtype | None = None,
171
+ ) -> dict[str, torch.Tensor | None]:
172
+ device = device or components._execution_device
173
+ dtype = dtype or components.text_encoder.dtype
174
+
175
+ prompt = [prompt] if isinstance(prompt, str) else prompt
176
+ batch_size = len(prompt)
177
+
178
+ prompt_embeds, prompt_attention_mask = cls._get_qwen_prompt_embeds(
179
+ components=components,
180
+ prompt=prompt,
181
+ max_sequence_length=max_sequence_length,
182
+ device=device,
183
+ dtype=dtype,
184
+ )
185
+ t5_input_ids, t5_attention_mask = cls._get_t5_prompt_ids(
186
+ components=components,
187
+ prompt=prompt,
188
+ max_sequence_length=max_sequence_length,
189
+ device=device,
190
+ )
191
+
192
+ negative_prompt_embeds = None
193
+ negative_prompt_attention_mask = None
194
+ negative_t5_input_ids = None
195
+ negative_t5_attention_mask = None
196
+ if prepare_unconditional_embeds:
197
+ negative_prompt = negative_prompt if negative_prompt is not None else ""
198
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
199
+
200
+ if prompt is not None and type(prompt) is not type(negative_prompt):
201
+ raise TypeError(
202
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
203
+ f" {type(prompt)}."
204
+ )
205
+ if batch_size != len(negative_prompt):
206
+ raise ValueError(
207
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
208
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
209
+ " the batch size of `prompt`."
210
+ )
211
+
212
+ negative_prompt_embeds, negative_prompt_attention_mask = cls._get_qwen_prompt_embeds(
213
+ components=components,
214
+ prompt=negative_prompt,
215
+ max_sequence_length=max_sequence_length,
216
+ device=device,
217
+ dtype=dtype,
218
+ )
219
+ negative_t5_input_ids, negative_t5_attention_mask = cls._get_t5_prompt_ids(
220
+ components=components,
221
+ prompt=negative_prompt,
222
+ max_sequence_length=max_sequence_length,
223
+ device=device,
224
+ )
225
+
226
+ return {
227
+ "qwen_prompt_embeds": prompt_embeds,
228
+ "qwen_attention_mask": prompt_attention_mask,
229
+ "t5_input_ids": t5_input_ids,
230
+ "t5_attention_mask": t5_attention_mask,
231
+ "negative_qwen_prompt_embeds": negative_prompt_embeds,
232
+ "negative_qwen_attention_mask": negative_prompt_attention_mask,
233
+ "negative_t5_input_ids": negative_t5_input_ids,
234
+ "negative_t5_attention_mask": negative_t5_attention_mask,
235
+ }
236
+
237
+ @torch.no_grad()
238
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
239
+ block_state = self.get_block_state(state)
240
+ self.check_inputs(block_state)
241
+
242
+ prompt_outputs = self.encode_prompt(
243
+ components=components,
244
+ prompt=block_state.prompt,
245
+ negative_prompt=block_state.negative_prompt,
246
+ prepare_unconditional_embeds=components.guider.num_conditions > 1,
247
+ max_sequence_length=block_state.max_sequence_length,
248
+ device=components._execution_device,
249
+ dtype=components.text_encoder.dtype,
250
+ )
251
+ for name, value in prompt_outputs.items():
252
+ setattr(block_state, name, value)
253
+
254
+ self.set_block_state(state, block_state)
255
+ return components, state
256
+
257
+
258
+ # Copied from diffusers.modular_pipelines.qwenimage.encoders.retrieve_latents
259
+ def retrieve_latents(
260
+ encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
261
+ ):
262
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
263
+ return encoder_output.latent_dist.sample(generator)
264
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
265
+ return encoder_output.latent_dist.mode()
266
+ elif hasattr(encoder_output, "latents"):
267
+ return encoder_output.latents
268
+ else:
269
+ raise AttributeError("Could not access latents of provided encoder_output")
270
+
271
+
272
+ # Copied from diffusers.modular_pipelines.qwenimage.encoders.encode_vae_image
273
+ def encode_vae_image(
274
+ image: torch.Tensor,
275
+ vae: AutoencoderKLQwenImage,
276
+ generator: torch.Generator,
277
+ device: torch.device,
278
+ dtype: torch.dtype,
279
+ latent_channels: int = 16,
280
+ sample_mode: str = "argmax",
281
+ ):
282
+ if not isinstance(image, torch.Tensor):
283
+ raise ValueError(f"Expected image to be a tensor, got {type(image)}.")
284
+
285
+ # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width
286
+ if image.dim() == 4:
287
+ image = image.unsqueeze(2)
288
+ elif image.dim() != 5:
289
+ raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.")
290
+
291
+ image = image.to(device=device, dtype=dtype)
292
+
293
+ if isinstance(generator, list):
294
+ image_latents = [
295
+ retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode)
296
+ for i in range(image.shape[0])
297
+ ]
298
+ image_latents = torch.cat(image_latents, dim=0)
299
+ else:
300
+ image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
301
+ latents_mean = (
302
+ torch.tensor(vae.config.latents_mean)
303
+ .view(1, latent_channels, 1, 1, 1)
304
+ .to(image_latents.device, image_latents.dtype)
305
+ )
306
+ latents_std = (
307
+ torch.tensor(vae.config.latents_std)
308
+ .view(1, latent_channels, 1, 1, 1)
309
+ .to(image_latents.device, image_latents.dtype)
310
+ )
311
+ image_latents = (image_latents - latents_mean) / latents_std
312
+
313
+ return image_latents
314
+
315
+
316
+ class AnimaImg2ImgVaeEncoderStep(ModularPipelineBlocks):
317
+ """VAE Encoder step for Anima image-to-image generation.
318
+
319
+ Preprocesses the input image and encodes it with the VAE, producing ``image_latents``. Timestep slicing is handled
320
+ downstream by ``AnimaImg2ImgSetTimestepsStep`` and noise addition by ``AnimaImg2ImgPrepareLatentsStep``.
321
+
322
+ Components:
323
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
324
+
325
+ Inputs:
326
+ image (`PIL.Image.Image`):
327
+ Input image to encode.
328
+ height (`int`, *optional*):
329
+ Height of the output image. Defaults to pipeline default.
330
+ width (`int`, *optional*):
331
+ Width of the output image. Defaults to pipeline default.
332
+ generator (`Generator`, *optional*):
333
+ Torch generator for deterministic generation.
334
+
335
+ Outputs:
336
+ image_latents (`Tensor`):
337
+ Encoded image latents.
338
+ height (`int`):
339
+ Output image height.
340
+ width (`int`):
341
+ Output image width.
342
+ """
343
+
344
+ model_name = "anima"
345
+
346
+ @property
347
+ def expected_components(self) -> list[ComponentSpec]:
348
+ return [
349
+ ComponentSpec("vae", AutoencoderKLQwenImage),
350
+ ComponentSpec(
351
+ "image_processor",
352
+ VaeImageProcessor,
353
+ config=FrozenDict({"vae_scale_factor": 8}),
354
+ default_creation_method="from_config",
355
+ ),
356
+ ]
357
+
358
+ @property
359
+ def description(self) -> str:
360
+ return (
361
+ "VAE Encoder step for Anima image-to-image generation. Encodes the input image to produce image_latents."
362
+ )
363
+
364
+ @property
365
+ def inputs(self) -> list[InputParam]:
366
+ return [
367
+ InputParam.template("image"),
368
+ InputParam.template("height"),
369
+ InputParam.template("width"),
370
+ InputParam.template("generator"),
371
+ ]
372
+
373
+ @property
374
+ def intermediate_outputs(self) -> list[OutputParam]:
375
+ return [
376
+ OutputParam("image_latents", type_hint=torch.Tensor, description="Encoded image latents."),
377
+ OutputParam("height", type_hint=int, description="Image height used for generation."),
378
+ OutputParam("width", type_hint=int, description="Image width used for generation."),
379
+ ]
380
+
381
+ @torch.no_grad()
382
+ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState:
383
+ block_state = self.get_block_state(state)
384
+
385
+ device = components._execution_device
386
+
387
+ block_state.height = block_state.height or components.default_height
388
+ block_state.width = block_state.width or components.default_width
389
+
390
+ processed_image = components.image_processor.preprocess(
391
+ image=block_state.image, height=block_state.height, width=block_state.width
392
+ )
393
+
394
+ block_state.image_latents = encode_vae_image(
395
+ image=processed_image,
396
+ vae=components.vae,
397
+ generator=block_state.generator,
398
+ device=device,
399
+ dtype=components.vae.dtype,
400
+ latent_channels=components.num_channels_latents,
401
+ )
402
+
403
+ self.set_block_state(state, block_state)
404
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py ADDED
@@ -0,0 +1,381 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
16
+ from ..modular_pipeline_utils import OutputParam
17
+ from .before_denoise import (
18
+ AnimaImageInputStep,
19
+ AnimaImg2ImgPrepareLatentsStep,
20
+ AnimaImg2ImgSetTimestepsStep,
21
+ AnimaPrepareLatentsStep,
22
+ AnimaSetTimestepsStep,
23
+ AnimaTextConditioningStep,
24
+ AnimaTextInputStep,
25
+ )
26
+ from .decoders import AnimaProcessImagesOutputStep, AnimaVaeDecoderStep
27
+ from .denoise import AnimaDenoiseStep
28
+ from .encoders import AnimaImg2ImgVaeEncoderStep, AnimaTextEncoderStep
29
+
30
+
31
+ # auto_docstring
32
+ class AnimaCoreDenoiseStep(SequentialPipelineBlocks):
33
+ """
34
+ Denoise block that takes encoded Anima text inputs and runs the denoising process.
35
+
36
+ Components:
37
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
38
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
39
+
40
+ Inputs:
41
+ qwen_prompt_embeds (`Tensor`):
42
+ Qwen prompt embeddings generated by the text encoder step.
43
+ qwen_attention_mask (`Tensor`):
44
+ Qwen prompt attention mask generated by the text encoder step.
45
+ t5_input_ids (`Tensor`):
46
+ T5 prompt token ids generated by the text encoder step.
47
+ t5_attention_mask (`Tensor`):
48
+ T5 prompt attention mask generated by the text encoder step.
49
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
50
+ Negative Qwen prompt embeddings generated by the text encoder step.
51
+ negative_qwen_attention_mask (`Tensor`, *optional*):
52
+ Negative Qwen prompt attention mask generated by the text encoder step.
53
+ negative_t5_input_ids (`Tensor`, *optional*):
54
+ Negative T5 prompt token ids generated by the text encoder step.
55
+ negative_t5_attention_mask (`Tensor`, *optional*):
56
+ Negative T5 prompt attention mask generated by the text encoder step.
57
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
58
+ The number of images to generate per prompt.
59
+ height (`int`, *optional*):
60
+ The height in pixels of the generated image.
61
+ width (`int`, *optional*):
62
+ The width in pixels of the generated image.
63
+ latents (`Tensor`, *optional*):
64
+ Pre-generated noisy latents for image generation.
65
+ generator (`Generator`, *optional*):
66
+ Torch generator for deterministic generation.
67
+ num_inference_steps (`int`, *optional*, defaults to 50):
68
+ The number of denoising steps.
69
+ sigmas (`list`, *optional*):
70
+ Custom sigmas for the denoising process.
71
+ **denoiser_input_fields (`None`, *optional*):
72
+ The conditional model inputs for the Anima denoiser.
73
+
74
+ Outputs:
75
+ latents (`Tensor`):
76
+ Denoised latents.
77
+ """
78
+
79
+ block_classes = [
80
+ AnimaTextConditioningStep,
81
+ AnimaTextInputStep,
82
+ AnimaPrepareLatentsStep,
83
+ AnimaSetTimestepsStep,
84
+ AnimaDenoiseStep,
85
+ ]
86
+ block_names = ["text_conditioning", "input", "prepare_latents", "set_timesteps", "denoise"]
87
+
88
+ @property
89
+ def description(self) -> str:
90
+ return "Denoise block that takes encoded Anima text inputs and runs the denoising process."
91
+
92
+ @property
93
+ def outputs(self):
94
+ return [OutputParam.template("latents")]
95
+
96
+
97
+ # auto_docstring
98
+ class AnimaDecodeStep(SequentialPipelineBlocks):
99
+ """
100
+ Decode Anima latents into generated images.
101
+
102
+ Components:
103
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
104
+
105
+ Inputs:
106
+ latents (`Tensor`):
107
+ Denoised Anima latents.
108
+ output_type (`str`, *optional*, defaults to pil):
109
+ Output format: 'pil', 'np', 'pt'.
110
+
111
+ Outputs:
112
+ images (`list`):
113
+ Generated images.
114
+ """
115
+
116
+ block_classes = [AnimaVaeDecoderStep, AnimaProcessImagesOutputStep]
117
+ block_names = ["decode", "postprocess"]
118
+
119
+ @property
120
+ def description(self) -> str:
121
+ return "Decode Anima latents into generated images."
122
+
123
+ @property
124
+ def outputs(self):
125
+ return [OutputParam.template("images")]
126
+
127
+
128
+ # auto_docstring
129
+ class AnimaImg2ImgCoreDenoiseStep(SequentialPipelineBlocks):
130
+ """
131
+ Denoise block for Anima image-to-image generation. Uses image_latents already in state from
132
+ AnimaImg2ImgVaeEncoderStep.
133
+
134
+ Components:
135
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
136
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
137
+
138
+ Inputs:
139
+ qwen_prompt_embeds (`Tensor`):
140
+ Qwen prompt embeddings generated by the text encoder step.
141
+ qwen_attention_mask (`Tensor`):
142
+ Qwen prompt attention mask generated by the text encoder step.
143
+ t5_input_ids (`Tensor`):
144
+ T5 prompt token ids generated by the text encoder step.
145
+ t5_attention_mask (`Tensor`):
146
+ T5 prompt attention mask generated by the text encoder step.
147
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
148
+ Negative Qwen prompt embeddings generated by the text encoder step.
149
+ negative_qwen_attention_mask (`Tensor`, *optional*):
150
+ Negative Qwen prompt attention mask generated by the text encoder step.
151
+ negative_t5_input_ids (`Tensor`, *optional*):
152
+ Negative T5 prompt token ids generated by the text encoder step.
153
+ negative_t5_attention_mask (`Tensor`, *optional*):
154
+ Negative T5 prompt attention mask generated by the text encoder step.
155
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
156
+ The number of images to generate per prompt.
157
+ image_latents (`Tensor`):
158
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
159
+ height (`int`, *optional*):
160
+ The height in pixels of the generated image.
161
+ width (`int`, *optional*):
162
+ The width in pixels of the generated image.
163
+ num_inference_steps (`int`, *optional*, defaults to 50):
164
+ The number of denoising steps.
165
+ sigmas (`list`, *optional*):
166
+ Custom sigmas for the denoising process.
167
+ strength (`float`, *optional*, defaults to 0.9):
168
+ Strength for img2img/inpainting.
169
+ generator (`Generator`, *optional*):
170
+ Torch generator for deterministic generation.
171
+ latents (`Tensor`, *optional*):
172
+ Pre-generated noisy latents for image generation.
173
+ **denoiser_input_fields (`None`, *optional*):
174
+ The conditional model inputs for the Anima denoiser.
175
+
176
+ Outputs:
177
+ latents (`Tensor`):
178
+ Denoised latents.
179
+ """
180
+
181
+ block_classes = [
182
+ AnimaTextConditioningStep,
183
+ AnimaTextInputStep,
184
+ AnimaImageInputStep,
185
+ AnimaImg2ImgSetTimestepsStep,
186
+ AnimaImg2ImgPrepareLatentsStep,
187
+ AnimaDenoiseStep,
188
+ ]
189
+ block_names = ["text_conditioning", "input", "image_input", "set_timesteps", "prepare_latents", "denoise"]
190
+
191
+ @property
192
+ def description(self) -> str:
193
+ return (
194
+ "Denoise block for Anima image-to-image generation. "
195
+ "Uses image_latents already in state from AnimaImg2ImgVaeEncoderStep."
196
+ )
197
+
198
+ @property
199
+ def outputs(self):
200
+ return [OutputParam.template("latents")]
201
+
202
+
203
+ # auto_docstring
204
+ class AnimaAutoCoreDenoiseStep(AutoPipelineBlocks):
205
+ """
206
+ Denoise step that selects between text-to-image and image-to-image denoising based on whether image_latents is
207
+ present in state. - `AnimaCoreDenoiseStep` (text2image) is used when no image_latents are present. -
208
+ `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when image_latents are present.
209
+
210
+ Components:
211
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
212
+ (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
213
+
214
+ Inputs:
215
+ qwen_prompt_embeds (`Tensor`):
216
+ Qwen prompt embeddings generated by the text encoder step.
217
+ qwen_attention_mask (`Tensor`):
218
+ Qwen prompt attention mask generated by the text encoder step.
219
+ t5_input_ids (`Tensor`):
220
+ T5 prompt token ids generated by the text encoder step.
221
+ t5_attention_mask (`Tensor`):
222
+ T5 prompt attention mask generated by the text encoder step.
223
+ negative_qwen_prompt_embeds (`Tensor`, *optional*):
224
+ Negative Qwen prompt embeddings generated by the text encoder step.
225
+ negative_qwen_attention_mask (`Tensor`, *optional*):
226
+ Negative Qwen prompt attention mask generated by the text encoder step.
227
+ negative_t5_input_ids (`Tensor`, *optional*):
228
+ Negative T5 prompt token ids generated by the text encoder step.
229
+ negative_t5_attention_mask (`Tensor`, *optional*):
230
+ Negative T5 prompt attention mask generated by the text encoder step.
231
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
232
+ The number of images to generate per prompt.
233
+ image_latents (`Tensor`, *optional*):
234
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
235
+ height (`int`, *optional*):
236
+ The height in pixels of the generated image.
237
+ width (`int`, *optional*):
238
+ The width in pixels of the generated image.
239
+ num_inference_steps (`int`):
240
+ The number of denoising steps.
241
+ sigmas (`list`, *optional*):
242
+ Custom sigmas for the denoising process.
243
+ strength (`float`, *optional*, defaults to 0.9):
244
+ Strength for img2img/inpainting.
245
+ generator (`Generator`, *optional*):
246
+ Torch generator for deterministic generation.
247
+ latents (`Tensor`):
248
+ Pre-generated noisy latents for image generation.
249
+ **denoiser_input_fields (`None`, *optional*):
250
+ The conditional model inputs for the Anima denoiser.
251
+
252
+ Outputs:
253
+ latents (`Tensor`):
254
+ Denoised latents.
255
+ """
256
+
257
+ block_classes = [AnimaImg2ImgCoreDenoiseStep, AnimaCoreDenoiseStep]
258
+ block_names = ["img2img", "text2image"]
259
+ block_trigger_inputs = ["image_latents", None]
260
+
261
+ @property
262
+ def description(self) -> str:
263
+ return (
264
+ "Denoise step that selects between text-to-image and image-to-image denoising based on whether "
265
+ "image_latents is present in state."
266
+ " - `AnimaCoreDenoiseStep` (text2image) is used when no image_latents are present."
267
+ " - `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when image_latents are present."
268
+ )
269
+
270
+
271
+ # auto_docstring
272
+ class AnimaAutoVaeImageEncoderStep(AutoPipelineBlocks):
273
+ """
274
+ VAE Image Encoder step that encodes the input image to produce image_latents. Skipped when no image is provided
275
+ (text-to-image workflow).
276
+
277
+ Components:
278
+ vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
279
+
280
+ Inputs:
281
+ image (`Image | list`, *optional*):
282
+ Reference image(s) for denoising. Can be a single image or list of images.
283
+ height (`int`, *optional*):
284
+ The height in pixels of the generated image.
285
+ width (`int`, *optional*):
286
+ The width in pixels of the generated image.
287
+ generator (`Generator`, *optional*):
288
+ Torch generator for deterministic generation.
289
+
290
+ Outputs:
291
+ image_latents (`Tensor`):
292
+ Encoded image latents.
293
+ height (`int`):
294
+ Image height used for generation.
295
+ width (`int`):
296
+ Image width used for generation.
297
+ """
298
+
299
+ block_classes = [AnimaImg2ImgVaeEncoderStep]
300
+ block_names = ["vae_encoder"]
301
+ block_trigger_inputs = ["image"]
302
+
303
+ @property
304
+ def description(self) -> str:
305
+ return (
306
+ "VAE Image Encoder step that encodes the input image to produce image_latents. "
307
+ "Skipped when no image is provided (text-to-image workflow)."
308
+ )
309
+
310
+
311
+ # auto_docstring
312
+ class AnimaAutoBlocks(SequentialPipelineBlocks):
313
+ """
314
+ Auto Modular pipeline for text-to-image and image-to-image generation using Anima.
315
+
316
+ Supported workflows:
317
+ - `text2image`: requires `prompt`
318
+ - `img2img`: requires `image`, `prompt`
319
+
320
+ Components:
321
+ text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) t5_tokenizer (`T5Tokenizer`) guider
322
+ (`ClassifierFreeGuidance`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
323
+ text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler
324
+ (`FlowMatchEulerDiscreteScheduler`)
325
+
326
+ Inputs:
327
+ prompt (`str`):
328
+ The prompt or prompts to guide image generation.
329
+ negative_prompt (`str`, *optional*):
330
+ The prompt or prompts not to guide the image generation.
331
+ max_sequence_length (`int`, *optional*, defaults to 512):
332
+ Maximum sequence length for prompt encoding.
333
+ image (`Image | list`, *optional*):
334
+ Reference image(s) for denoising. Can be a single image or list of images.
335
+ height (`int`, *optional*):
336
+ The height in pixels of the generated image.
337
+ width (`int`, *optional*):
338
+ The width in pixels of the generated image.
339
+ generator (`Generator`, *optional*):
340
+ Torch generator for deterministic generation.
341
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
342
+ The number of images to generate per prompt.
343
+ image_latents (`Tensor`, *optional*):
344
+ image latents used to guide the image generation. Can be generated from vae_encoder step.
345
+ num_inference_steps (`int`):
346
+ The number of denoising steps.
347
+ sigmas (`list`, *optional*):
348
+ Custom sigmas for the denoising process.
349
+ strength (`float`, *optional*, defaults to 0.9):
350
+ Strength for img2img/inpainting.
351
+ latents (`Tensor`):
352
+ Pre-generated noisy latents for image generation.
353
+ **denoiser_input_fields (`None`, *optional*):
354
+ The conditional model inputs for the Anima denoiser.
355
+ output_type (`str`, *optional*, defaults to pil):
356
+ Output format: 'pil', 'np', 'pt'.
357
+
358
+ Outputs:
359
+ images (`list`):
360
+ Generated images.
361
+ """
362
+
363
+ block_classes = [
364
+ AnimaTextEncoderStep,
365
+ AnimaAutoVaeImageEncoderStep,
366
+ AnimaAutoCoreDenoiseStep,
367
+ AnimaDecodeStep,
368
+ ]
369
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
370
+ _workflow_map = {
371
+ "text2image": {"prompt": True},
372
+ "img2img": {"image": True, "prompt": True},
373
+ }
374
+
375
+ @property
376
+ def description(self) -> str:
377
+ return "Auto Modular pipeline for text-to-image and image-to-image generation using Anima."
378
+
379
+ @property
380
+ def outputs(self):
381
+ return [OutputParam.template("images")]
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from ...loaders import AnimaLoraLoaderMixin
16
+ from ..modular_pipeline import ModularPipeline
17
+
18
+
19
+ class AnimaModularPipeline(ModularPipeline, AnimaLoraLoaderMixin):
20
+ """
21
+ A ModularPipeline for Anima.
22
+
23
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
24
+ """
25
+
26
+ default_blocks_name = "AnimaAutoBlocks"
27
+
28
+ @property
29
+ def default_height(self):
30
+ return self.default_sample_size * self.vae_scale_factor
31
+
32
+ @property
33
+ def default_width(self):
34
+ return self.default_sample_size * self.vae_scale_factor
35
+
36
+ @property
37
+ def default_sample_size(self):
38
+ return 128
39
+
40
+ @property
41
+ def vae_scale_factor(self):
42
+ vae_scale_factor = 8
43
+ if self.vae is not None:
44
+ vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
45
+ return vae_scale_factor
46
+
47
+ @property
48
+ def num_channels_latents(self):
49
+ num_channels_latents = 16
50
+ if self.transformer is not None:
51
+ num_channels_latents = self.transformer.config.in_channels
52
+ return num_channels_latents
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py ADDED
@@ -0,0 +1,1109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from __future__ import annotations
16
+
17
+ import copy
18
+ import time
19
+ from collections import OrderedDict
20
+ from itertools import combinations
21
+ from typing import Any
22
+
23
+ import torch
24
+
25
+ from ..hooks import ModelHook
26
+ from ..utils import (
27
+ is_accelerate_available,
28
+ logging,
29
+ )
30
+ from ..utils.torch_utils import get_device
31
+
32
+
33
+ if is_accelerate_available():
34
+ from accelerate.hooks import add_hook_to_module, remove_hook_from_module
35
+ from accelerate.state import PartialState
36
+ from accelerate.utils import send_to_device
37
+ from accelerate.utils.memory import clear_device_cache
38
+ from accelerate.utils.modeling import convert_file_size_to_int
39
+
40
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
41
+
42
+
43
+ class CustomOffloadHook(ModelHook):
44
+ """
45
+ A hook that offloads a model on the CPU until its forward pass is called. It ensures the model and its inputs are
46
+ on the given device. Optionally offloads other models to the CPU before the forward pass is called.
47
+
48
+ Args:
49
+ execution_device(`str`, `int` or `torch.device`, *optional*):
50
+ The device on which the model should be executed. Will default to the MPS device if it's available, then
51
+ GPU 0 if there is a GPU, and finally to the CPU.
52
+ """
53
+
54
+ no_grad = False
55
+
56
+ def __init__(
57
+ self,
58
+ execution_device: str | int | torch.device | None = None,
59
+ other_hooks: list["UserCustomOffloadHook"] | None = None,
60
+ offload_strategy: "AutoOffloadStrategy" | None = None,
61
+ ):
62
+ self.execution_device = execution_device if execution_device is not None else PartialState().default_device
63
+ self.other_hooks = other_hooks
64
+ self.offload_strategy = offload_strategy
65
+ self.model_id = None
66
+
67
+ def set_strategy(self, offload_strategy: "AutoOffloadStrategy"):
68
+ self.offload_strategy = offload_strategy
69
+
70
+ def add_other_hook(self, hook: "UserCustomOffloadHook"):
71
+ """
72
+ Add a hook to the list of hooks to consider for offloading.
73
+ """
74
+ if self.other_hooks is None:
75
+ self.other_hooks = []
76
+ self.other_hooks.append(hook)
77
+
78
+ def init_hook(self, module):
79
+ return module.to("cpu")
80
+
81
+ def pre_forward(self, module, *args, **kwargs):
82
+ if module.device != self.execution_device:
83
+ if self.other_hooks is not None:
84
+ hooks_to_offload = [hook for hook in self.other_hooks if hook.model.device == self.execution_device]
85
+ # offload all other hooks
86
+ start_time = time.perf_counter()
87
+ if self.offload_strategy is not None:
88
+ hooks_to_offload = self.offload_strategy(
89
+ hooks=hooks_to_offload,
90
+ model_id=self.model_id,
91
+ model=module,
92
+ execution_device=self.execution_device,
93
+ )
94
+ end_time = time.perf_counter()
95
+ logger.info(
96
+ f" time taken to apply offload strategy for {self.model_id}: {(end_time - start_time):.2f} seconds"
97
+ )
98
+
99
+ for hook in hooks_to_offload:
100
+ logger.info(
101
+ f"moving {self.model_id} to {self.execution_device}, offloading {hook.model_id} to cpu"
102
+ )
103
+ hook.offload()
104
+
105
+ if hooks_to_offload:
106
+ clear_device_cache()
107
+ module.to(self.execution_device)
108
+ return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)
109
+
110
+
111
+ class UserCustomOffloadHook:
112
+ """
113
+ A simple hook grouping a model and a `CustomOffloadHook`, which provides easy APIs for to call the init method of
114
+ the hook or remove it entirely.
115
+ """
116
+
117
+ def __init__(self, model_id, model, hook):
118
+ self.model_id = model_id
119
+ self.model = model
120
+ self.hook = hook
121
+
122
+ def offload(self):
123
+ self.hook.init_hook(self.model)
124
+
125
+ def attach(self):
126
+ add_hook_to_module(self.model, self.hook)
127
+ self.hook.model_id = self.model_id
128
+
129
+ def remove(self):
130
+ remove_hook_from_module(self.model)
131
+ self.hook.model_id = None
132
+
133
+ def add_other_hook(self, hook: "UserCustomOffloadHook"):
134
+ self.hook.add_other_hook(hook)
135
+
136
+
137
+ def custom_offload_with_hook(
138
+ model_id: str,
139
+ model: torch.nn.Module,
140
+ execution_device: str | int | torch.device = None,
141
+ offload_strategy: "AutoOffloadStrategy" | None = None,
142
+ ):
143
+ hook = CustomOffloadHook(execution_device=execution_device, offload_strategy=offload_strategy)
144
+ user_hook = UserCustomOffloadHook(model_id=model_id, model=model, hook=hook)
145
+ user_hook.attach()
146
+ return user_hook
147
+
148
+
149
+ # this is the class that user can customize to implement their own offload strategy
150
+ class AutoOffloadStrategy:
151
+ """
152
+ Offload strategy that should be used with `CustomOffloadHook` to automatically offload models to the CPU based on
153
+ the available memory on the device.
154
+ """
155
+
156
+ # YiYi TODO: instead of memory_reserve_margin, we should let user set the maximum_total_models_size to keep on device
157
+ # the actual memory usage would be higher. But it's simpler this way, and can be tested
158
+ def __init__(self, memory_reserve_margin="3GB"):
159
+ self.memory_reserve_margin = convert_file_size_to_int(memory_reserve_margin)
160
+
161
+ def __call__(self, hooks, model_id, model, execution_device):
162
+ if len(hooks) == 0:
163
+ return []
164
+
165
+ try:
166
+ current_module_size = model.get_memory_footprint()
167
+ except AttributeError:
168
+ raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.")
169
+
170
+ device_type = execution_device.type
171
+ device_module = getattr(torch, device_type, torch.cuda)
172
+ try:
173
+ mem_on_device = device_module.mem_get_info(execution_device.index)[0]
174
+ except AttributeError:
175
+ raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.")
176
+
177
+ mem_on_device = mem_on_device - self.memory_reserve_margin
178
+ if current_module_size < mem_on_device:
179
+ return []
180
+
181
+ min_memory_offload = current_module_size - mem_on_device
182
+ logger.info(f" search for models to offload in order to free up {min_memory_offload / 1024**3:.2f} GB memory")
183
+
184
+ # exlucde models that's not currently loaded on the device
185
+ module_sizes = dict(
186
+ sorted(
187
+ {hook.model_id: hook.model.get_memory_footprint() for hook in hooks}.items(),
188
+ key=lambda x: x[1],
189
+ reverse=True,
190
+ )
191
+ )
192
+
193
+ # YiYi/Dhruv TODO: sort smallest to largest, and offload in that order we would tend to keep the larger models on GPU more often
194
+ def search_best_candidate(module_sizes, min_memory_offload):
195
+ """
196
+ search the optimal combination of models to offload to cpu, given a dictionary of module sizes and a
197
+ minimum memory offload size. the combination of models should add up to the smallest modulesize that is
198
+ larger than `min_memory_offload`
199
+ """
200
+ model_ids = list(module_sizes.keys())
201
+ best_candidate = None
202
+ best_size = float("inf")
203
+ for r in range(1, len(model_ids) + 1):
204
+ for candidate_model_ids in combinations(model_ids, r):
205
+ candidate_size = sum(
206
+ module_sizes[candidate_model_id] for candidate_model_id in candidate_model_ids
207
+ )
208
+ if candidate_size < min_memory_offload:
209
+ continue
210
+ else:
211
+ if best_candidate is None or candidate_size < best_size:
212
+ best_candidate = candidate_model_ids
213
+ best_size = candidate_size
214
+
215
+ return best_candidate
216
+
217
+ best_offload_model_ids = search_best_candidate(module_sizes, min_memory_offload)
218
+
219
+ if best_offload_model_ids is None:
220
+ # if no combination is found, meaning that we cannot meet the memory requirement, offload all models
221
+ logger.warning("no combination of models to offload to cpu is found, offloading all models")
222
+ hooks_to_offload = hooks
223
+ else:
224
+ hooks_to_offload = [hook for hook in hooks if hook.model_id in best_offload_model_ids]
225
+
226
+ return hooks_to_offload
227
+
228
+
229
+ # utils for display component info in a readable format
230
+ # TODO: move to a different file
231
+ def summarize_dict_by_value_and_parts(d: dict[str, Any]) -> dict[str, Any]:
232
+ """Summarizes a dictionary by finding common prefixes that share the same value.
233
+
234
+ For a dictionary with dot-separated keys like: {
235
+ 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor': [0.6],
236
+ 'down_blocks.1.attentions.1.transformer_blocks.1.attn2.processor': [0.6],
237
+ 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor': [0.3],
238
+ }
239
+
240
+ Returns a dictionary where keys are the shortest common prefixes and values are their shared values: {
241
+ 'down_blocks': [0.6], 'up_blocks': [0.3]
242
+ }
243
+ """
244
+ # First group by values - convert lists to tuples to make them hashable
245
+ value_to_keys = {}
246
+ for key, value in d.items():
247
+ value_tuple = tuple(value) if isinstance(value, list) else value
248
+ if value_tuple not in value_to_keys:
249
+ value_to_keys[value_tuple] = []
250
+ value_to_keys[value_tuple].append(key)
251
+
252
+ def find_common_prefix(keys: list[str]) -> str:
253
+ """Find the shortest common prefix among a list of dot-separated keys."""
254
+ if not keys:
255
+ return ""
256
+ if len(keys) == 1:
257
+ return keys[0]
258
+
259
+ # Split all keys into parts
260
+ key_parts = [k.split(".") for k in keys]
261
+
262
+ # Find how many initial parts are common
263
+ common_length = 0
264
+ for parts in zip(*key_parts):
265
+ if len(set(parts)) == 1: # All parts at this position are the same
266
+ common_length += 1
267
+ else:
268
+ break
269
+
270
+ if common_length == 0:
271
+ return ""
272
+
273
+ # Return the common prefix
274
+ return ".".join(key_parts[0][:common_length])
275
+
276
+ # Create summary by finding common prefixes for each value group
277
+ summary = {}
278
+ for value_tuple, keys in value_to_keys.items():
279
+ prefix = find_common_prefix(keys)
280
+ if prefix: # Only add if we found a common prefix
281
+ # Convert tuple back to list if it was originally a list
282
+ value = list(value_tuple) if isinstance(d[keys[0]], list) else value_tuple
283
+ summary[prefix] = value
284
+ else:
285
+ summary[""] = value # Use empty string if no common prefix
286
+
287
+ return summary
288
+
289
+
290
+ class ComponentsManager:
291
+ """
292
+ A central registry and management system for model components across multiple pipelines.
293
+
294
+ [`ComponentsManager`] provides a unified way to register, track, and reuse model components (like UNet, VAE, text
295
+ encoders, etc.) across different modular pipelines. It includes features for duplicate detection, memory
296
+ management, and component organization.
297
+
298
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
299
+
300
+ Example:
301
+ ```python
302
+ from diffusers import ComponentsManager
303
+
304
+ # Create a components manager
305
+ cm = ComponentsManager()
306
+
307
+ # Add components
308
+ cm.add("unet", unet_model, collection="sdxl")
309
+ cm.add("vae", vae_model, collection="sdxl")
310
+
311
+ # Enable auto offloading
312
+ cm.enable_auto_cpu_offload()
313
+
314
+ # Retrieve components
315
+ unet = cm.get_one(name="unet", collection="sdxl")
316
+ ```
317
+ """
318
+
319
+ _available_info_fields = [
320
+ "model_id",
321
+ "added_time",
322
+ "collection",
323
+ "class_name",
324
+ "size_gb",
325
+ "adapters",
326
+ "has_hook",
327
+ "execution_device",
328
+ "ip_adapter",
329
+ "quantization",
330
+ ]
331
+
332
+ def __init__(self):
333
+ self.components = OrderedDict()
334
+ # YiYi TODO: can remove once confirm we don't need this in mellon
335
+ self.added_time = OrderedDict() # Store when components were added
336
+ self.collections = OrderedDict() # collection_name -> set of component_names
337
+ self.model_hooks = None
338
+ self._auto_offload_enabled = False
339
+
340
+ def _lookup_ids(
341
+ self,
342
+ name: str | None = None,
343
+ collection: str | None = None,
344
+ load_id: str | None = None,
345
+ components: OrderedDict | None = None,
346
+ ):
347
+ """
348
+ Lookup component_ids by name, collection, or load_id. Does not support pattern matching. Returns a set of
349
+ component_ids
350
+ """
351
+ if components is None:
352
+ components = self.components
353
+
354
+ if name:
355
+ ids_by_name = set()
356
+ for component_id, component in components.items():
357
+ comp_name = self._id_to_name(component_id)
358
+ if comp_name == name:
359
+ ids_by_name.add(component_id)
360
+ else:
361
+ ids_by_name = set(components.keys())
362
+ if collection and collection not in self.collections:
363
+ return set()
364
+ elif collection and collection in self.collections:
365
+ ids_by_collection = set()
366
+ for component_id, component in components.items():
367
+ if component_id in self.collections[collection]:
368
+ ids_by_collection.add(component_id)
369
+ else:
370
+ ids_by_collection = set(components.keys())
371
+ if load_id:
372
+ ids_by_load_id = set()
373
+ for name, component in components.items():
374
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id == load_id:
375
+ ids_by_load_id.add(name)
376
+ else:
377
+ ids_by_load_id = set(components.keys())
378
+
379
+ ids = ids_by_name.intersection(ids_by_collection).intersection(ids_by_load_id)
380
+ return ids
381
+
382
+ @staticmethod
383
+ def _id_to_name(component_id: str):
384
+ return "_".join(component_id.split("_")[:-1])
385
+
386
+ def add(self, name: str, component: Any, collection: str | None = None):
387
+ """
388
+ Add a component to the ComponentsManager.
389
+
390
+ Args:
391
+ name (str): The name of the component
392
+ component (Any): The component to add
393
+ collection (str | None): The collection to add the component to
394
+
395
+ Returns:
396
+ str: The unique component ID, which is generated as "{name}_{id(component)}" where
397
+ id(component) is Python's built-in unique identifier for the object
398
+ """
399
+ component_id = f"{name}_{id(component)}"
400
+ is_new_component = True
401
+
402
+ # check for duplicated components
403
+ for comp_id, comp in self.components.items():
404
+ if comp == component:
405
+ comp_name = self._id_to_name(comp_id)
406
+ if comp_name == name:
407
+ logger.warning(f"ComponentsManager: component '{name}' already exists as '{comp_id}'")
408
+ component_id = comp_id
409
+ is_new_component = False
410
+ break
411
+ else:
412
+ logger.warning(
413
+ f"ComponentsManager: adding component '{name}' as '{component_id}', but it is duplicate of '{comp_id}'"
414
+ f"To remove a duplicate, call `components_manager.remove('<component_id>')`."
415
+ )
416
+
417
+ # check for duplicated load_id and warn (we do not delete for you)
418
+ if hasattr(component, "_diffusers_load_id") and component._diffusers_load_id != "null":
419
+ components_with_same_load_id = self._lookup_ids(load_id=component._diffusers_load_id)
420
+ components_with_same_load_id = [id for id in components_with_same_load_id if id != component_id]
421
+
422
+ if components_with_same_load_id:
423
+ existing = ", ".join(components_with_same_load_id)
424
+ logger.warning(
425
+ f"ComponentsManager: adding component '{component_id}', but it has duplicate load_id '{component._diffusers_load_id}' with existing components: {existing}. "
426
+ f"To remove a duplicate, call `components_manager.remove('<component_id>')`."
427
+ )
428
+
429
+ # add component to components manager
430
+ self.components[component_id] = component
431
+ if is_new_component:
432
+ self.added_time[component_id] = time.time()
433
+
434
+ if collection:
435
+ if collection not in self.collections:
436
+ self.collections[collection] = set()
437
+ if component_id not in self.collections[collection]:
438
+ comp_ids_in_collection = self._lookup_ids(name=name, collection=collection)
439
+ for comp_id in comp_ids_in_collection:
440
+ logger.warning(
441
+ f"ComponentsManager: removing existing {name} from collection '{collection}': {comp_id}"
442
+ )
443
+ # remove existing component from this collection (if it is not in any other collection, will be removed from ComponentsManager)
444
+ self.remove_from_collection(comp_id, collection)
445
+
446
+ self.collections[collection].add(component_id)
447
+ logger.info(
448
+ f"ComponentsManager: added component '{name}' in collection '{collection}': {component_id}"
449
+ )
450
+ else:
451
+ logger.info(f"ComponentsManager: added component '{name}' as '{component_id}'")
452
+
453
+ if self._auto_offload_enabled and is_new_component:
454
+ self.enable_auto_cpu_offload(self._auto_offload_device)
455
+
456
+ return component_id
457
+
458
+ def remove_from_collection(self, component_id: str, collection: str):
459
+ """
460
+ Remove a component from a collection.
461
+ """
462
+ if collection not in self.collections:
463
+ logger.warning(f"Collection '{collection}' not found in ComponentsManager")
464
+ return
465
+ if component_id not in self.collections[collection]:
466
+ logger.warning(f"Component '{component_id}' not found in collection '{collection}'")
467
+ return
468
+ # remove from the collection
469
+ self.collections[collection].remove(component_id)
470
+ # check if this component is in any other collection
471
+ comp_colls = [coll for coll, comps in self.collections.items() if component_id in comps]
472
+ if not comp_colls: # only if no other collection contains this component, remove it
473
+ logger.warning(f"ComponentsManager: removing component '{component_id}' from ComponentsManager")
474
+ self.remove(component_id)
475
+
476
+ def remove(self, component_id: str = None):
477
+ """
478
+ Remove a component from the ComponentsManager.
479
+
480
+ Args:
481
+ component_id (str): The ID of the component to remove
482
+ """
483
+ if component_id not in self.components:
484
+ logger.warning(f"Component '{component_id}' not found in ComponentsManager")
485
+ return
486
+
487
+ component = self.components.pop(component_id)
488
+ self.added_time.pop(component_id)
489
+
490
+ for collection in self.collections:
491
+ if component_id in self.collections[collection]:
492
+ self.collections[collection].remove(component_id)
493
+
494
+ if self._auto_offload_enabled:
495
+ self.enable_auto_cpu_offload(self._auto_offload_device)
496
+ else:
497
+ if isinstance(component, torch.nn.Module):
498
+ component.to("cpu")
499
+ del component
500
+ import gc
501
+
502
+ gc.collect()
503
+ if torch.cuda.is_available():
504
+ torch.cuda.empty_cache()
505
+ if torch.xpu.is_available():
506
+ torch.xpu.empty_cache()
507
+
508
+ # YiYi TODO: rename to search_components for now, may remove this method
509
+ def search_components(
510
+ self,
511
+ names: str | None = None,
512
+ collection: str | None = None,
513
+ load_id: str | None = None,
514
+ return_dict_with_names: bool = True,
515
+ ):
516
+ """
517
+ Search components by name with simple pattern matching. Optionally filter by collection or load_id.
518
+
519
+ Args:
520
+ names: Component name(s) or pattern(s)
521
+ Patterns:
522
+ - "unet" : match any component with base name "unet" (e.g., unet_123abc)
523
+ - "!unet" : everything except components with base name "unet"
524
+ - "unet*" : anything with base name starting with "unet"
525
+ - "!unet*" : anything with base name NOT starting with "unet"
526
+ - "*unet*" : anything with base name containing "unet"
527
+ - "!*unet*" : anything with base name NOT containing "unet"
528
+ - "refiner|vae|unet" : anything with base name exactly matching "refiner", "vae", or "unet"
529
+ - "!refiner|vae|unet" : anything with base name NOT exactly matching "refiner", "vae", or "unet"
530
+ - "unet*|vae*" : anything with base name starting with "unet" OR starting with "vae"
531
+ collection: Optional collection to filter by
532
+ load_id: Optional load_id to filter by
533
+ return_dict_with_names:
534
+ If True, returns a dictionary with component names as keys, throw an error if
535
+ multiple components with the same name are found If False, returns a dictionary
536
+ with component IDs as keys
537
+
538
+ Returns:
539
+ Dictionary mapping component names to components if return_dict_with_names=True, or a dictionary mapping
540
+ component IDs to components if return_dict_with_names=False
541
+ """
542
+
543
+ # select components based on collection and load_id filters
544
+ selected_ids = self._lookup_ids(collection=collection, load_id=load_id)
545
+ components = {k: self.components[k] for k in selected_ids}
546
+
547
+ def get_return_dict(components, return_dict_with_names):
548
+ """
549
+ Create a dictionary mapping component names to components if return_dict_with_names=True, or a dictionary
550
+ mapping component IDs to components if return_dict_with_names=False, throw an error if duplicate component
551
+ names are found when return_dict_with_names=True
552
+ """
553
+ if return_dict_with_names:
554
+ dict_to_return = {}
555
+ for comp_id, comp in components.items():
556
+ comp_name = self._id_to_name(comp_id)
557
+ if comp_name in dict_to_return:
558
+ raise ValueError(
559
+ 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"
560
+ )
561
+ dict_to_return[comp_name] = comp
562
+ return dict_to_return
563
+ else:
564
+ return components
565
+
566
+ # if no names are provided, return the filtered components as it is
567
+ if names is None:
568
+ return get_return_dict(components, return_dict_with_names)
569
+
570
+ # if names is not a string, raise an error
571
+ elif not isinstance(names, str):
572
+ raise ValueError(f"Invalid type for `names: {type(names)}, only support string")
573
+
574
+ # Create mapping from component_id to base_name for components to be used for pattern matching
575
+ base_names = {comp_id: self._id_to_name(comp_id) for comp_id in components.keys()}
576
+
577
+ # Helper function to check if a component matches a pattern based on its base name
578
+ def matches_pattern(component_id, pattern, exact_match=False):
579
+ """
580
+ Helper function to check if a component matches a pattern based on its base name.
581
+
582
+ Args:
583
+ component_id: The component ID to check
584
+ pattern: The pattern to match against
585
+ exact_match: If True, only exact matches to base_name are considered
586
+ """
587
+ base_name = base_names[component_id]
588
+
589
+ # Exact match with base name
590
+ if exact_match:
591
+ return pattern == base_name
592
+
593
+ # Prefix match (ends with *)
594
+ elif pattern.endswith("*"):
595
+ prefix = pattern[:-1]
596
+ return base_name.startswith(prefix)
597
+
598
+ # Contains match (starts with *)
599
+ elif pattern.startswith("*"):
600
+ search = pattern[1:-1] if pattern.endswith("*") else pattern[1:]
601
+ return search in base_name
602
+
603
+ # Exact match (no wildcards)
604
+ else:
605
+ return pattern == base_name
606
+
607
+ # Check if this is a "not" pattern
608
+ is_not_pattern = names.startswith("!")
609
+ if is_not_pattern:
610
+ names = names[1:] # Remove the ! prefix
611
+
612
+ # Handle OR patterns (containing |)
613
+ if "|" in names:
614
+ terms = names.split("|")
615
+ matches = {}
616
+
617
+ for comp_id, comp in components.items():
618
+ # For OR patterns with exact names (no wildcards), we do exact matching on base names
619
+ exact_match = all(not (term.startswith("*") or term.endswith("*")) for term in terms)
620
+
621
+ # Check if any of the terms match this component
622
+ should_include = any(matches_pattern(comp_id, term, exact_match) for term in terms)
623
+
624
+ # Flip the decision if this is a NOT pattern
625
+ if is_not_pattern:
626
+ should_include = not should_include
627
+
628
+ if should_include:
629
+ matches[comp_id] = comp
630
+
631
+ log_msg = "NOT " if is_not_pattern else ""
632
+ match_type = "exactly matching" if exact_match else "matching any of patterns"
633
+ logger.info(f"Getting components {log_msg}{match_type} {terms}: {list(matches.keys())}")
634
+
635
+ # Try exact match with a base name
636
+ elif any(names == base_name for base_name in base_names.values()):
637
+ # Find all components with this base name
638
+ matches = {
639
+ comp_id: comp
640
+ for comp_id, comp in components.items()
641
+ if (base_names[comp_id] == names) != is_not_pattern
642
+ }
643
+
644
+ if is_not_pattern:
645
+ logger.info(f"Getting all components except those with base name '{names}': {list(matches.keys())}")
646
+ else:
647
+ logger.info(f"Getting components with base name '{names}': {list(matches.keys())}")
648
+
649
+ # Prefix match (ends with *)
650
+ elif names.endswith("*"):
651
+ prefix = names[:-1]
652
+ matches = {
653
+ comp_id: comp
654
+ for comp_id, comp in components.items()
655
+ if base_names[comp_id].startswith(prefix) != is_not_pattern
656
+ }
657
+ if is_not_pattern:
658
+ logger.info(f"Getting components NOT starting with '{prefix}': {list(matches.keys())}")
659
+ else:
660
+ logger.info(f"Getting components starting with '{prefix}': {list(matches.keys())}")
661
+
662
+ # Contains match (starts with *)
663
+ elif names.startswith("*"):
664
+ search = names[1:-1] if names.endswith("*") else names[1:]
665
+ matches = {
666
+ comp_id: comp
667
+ for comp_id, comp in components.items()
668
+ if (search in base_names[comp_id]) != is_not_pattern
669
+ }
670
+ if is_not_pattern:
671
+ logger.info(f"Getting components NOT containing '{search}': {list(matches.keys())}")
672
+ else:
673
+ logger.info(f"Getting components containing '{search}': {list(matches.keys())}")
674
+
675
+ # Substring match (no wildcards, but not an exact component name)
676
+ elif any(names in base_name for base_name in base_names.values()):
677
+ matches = {
678
+ comp_id: comp
679
+ for comp_id, comp in components.items()
680
+ if (names in base_names[comp_id]) != is_not_pattern
681
+ }
682
+ if is_not_pattern:
683
+ logger.info(f"Getting components NOT containing '{names}': {list(matches.keys())}")
684
+ else:
685
+ logger.info(f"Getting components containing '{names}': {list(matches.keys())}")
686
+
687
+ else:
688
+ raise ValueError(f"Component or pattern '{names}' not found in ComponentsManager")
689
+
690
+ if not matches:
691
+ raise ValueError(f"No components found matching pattern '{names}'")
692
+
693
+ return get_return_dict(matches, return_dict_with_names)
694
+
695
+ def enable_auto_cpu_offload(self, device: str | int | torch.device = None, memory_reserve_margin="3GB"):
696
+ """
697
+ Enable automatic CPU offloading for all components.
698
+
699
+ The algorithm works as follows:
700
+ 1. All models start on CPU by default
701
+ 2. When a model's forward pass is called, it's moved to the execution device
702
+ 3. If there's insufficient memory, other models on the device are moved back to CPU
703
+ 4. The system tries to offload the smallest combination of models that frees enough memory
704
+ 5. Models stay on the execution device until another model needs memory and forces them off
705
+
706
+ Args:
707
+ device (str | int | torch.device): The execution device where models are moved for forward passes
708
+ memory_reserve_margin (str): The memory reserve margin to use, default is 3GB. This is the amount of
709
+ memory to keep free on the device to avoid running out of memory during model
710
+ execution (e.g., for intermediate activations, gradients, etc.)
711
+ """
712
+ if not is_accelerate_available():
713
+ raise ImportError("Make sure to install accelerate to use auto_cpu_offload")
714
+
715
+ if device is None:
716
+ device = get_device()
717
+ if not isinstance(device, torch.device):
718
+ device = torch.device(device)
719
+
720
+ device_type = device.type
721
+ device_module = getattr(torch, device_type, torch.cuda)
722
+ if not hasattr(device_module, "mem_get_info"):
723
+ raise NotImplementedError(
724
+ f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}."
725
+ )
726
+
727
+ if device.index is None:
728
+ device = torch.device(f"{device.type}:{0}")
729
+
730
+ for name, component in self.components.items():
731
+ if isinstance(component, torch.nn.Module) and hasattr(component, "_hf_hook"):
732
+ remove_hook_from_module(component, recurse=True)
733
+
734
+ self.disable_auto_cpu_offload()
735
+ offload_strategy = AutoOffloadStrategy(memory_reserve_margin=memory_reserve_margin)
736
+
737
+ all_hooks = []
738
+ for name, component in self.components.items():
739
+ if isinstance(component, torch.nn.Module):
740
+ hook = custom_offload_with_hook(name, component, device, offload_strategy=offload_strategy)
741
+ all_hooks.append(hook)
742
+
743
+ for hook in all_hooks:
744
+ other_hooks = [h for h in all_hooks if h is not hook]
745
+ for other_hook in other_hooks:
746
+ if other_hook.hook.execution_device == hook.hook.execution_device:
747
+ hook.add_other_hook(other_hook)
748
+
749
+ self.model_hooks = all_hooks
750
+ self._auto_offload_enabled = True
751
+ self._auto_offload_device = device
752
+
753
+ def disable_auto_cpu_offload(self):
754
+ """
755
+ Disable automatic CPU offloading for all components.
756
+ """
757
+ if self.model_hooks is None:
758
+ self._auto_offload_enabled = False
759
+ return
760
+
761
+ for hook in self.model_hooks:
762
+ hook.offload()
763
+ hook.remove()
764
+ if self.model_hooks:
765
+ clear_device_cache()
766
+ self.model_hooks = None
767
+ self._auto_offload_enabled = False
768
+
769
+ def get_model_info(
770
+ self,
771
+ component_id: str,
772
+ fields: str | list[str] | None = None,
773
+ ) -> dict[str, Any] | None:
774
+ """Get comprehensive information about a component.
775
+
776
+ Args:
777
+ component_id (str): Name of the component to get info for
778
+ fields (str | list[str] | None):
779
+ Field(s) to return. Can be a string for single field or list of fields. If None, uses the
780
+ available_info_fields setting.
781
+
782
+ Returns:
783
+ Dictionary containing requested component metadata. If fields is specified, returns only those fields.
784
+ Otherwise, returns all fields.
785
+ """
786
+ if component_id not in self.components:
787
+ raise ValueError(f"Component '{component_id}' not found in ComponentsManager")
788
+
789
+ component = self.components[component_id]
790
+
791
+ # Validate fields if specified
792
+ if fields is not None:
793
+ if isinstance(fields, str):
794
+ fields = [fields]
795
+ for field in fields:
796
+ if field not in self._available_info_fields:
797
+ raise ValueError(f"Field '{field}' not found in available_info_fields")
798
+
799
+ # Build complete info dict first
800
+ info = {
801
+ "model_id": component_id,
802
+ "added_time": self.added_time[component_id],
803
+ "collection": ", ".join([coll for coll, comps in self.collections.items() if component_id in comps])
804
+ or None,
805
+ }
806
+
807
+ # Additional info for torch.nn.Module components
808
+ if isinstance(component, torch.nn.Module):
809
+ # Check for hook information
810
+ has_hook = hasattr(component, "_hf_hook")
811
+ execution_device = None
812
+ if has_hook and hasattr(component._hf_hook, "execution_device"):
813
+ execution_device = component._hf_hook.execution_device
814
+
815
+ info.update(
816
+ {
817
+ "class_name": component.__class__.__name__,
818
+ "size_gb": component.get_memory_footprint() / (1024**3),
819
+ "adapters": None, # Default to None
820
+ "has_hook": has_hook,
821
+ "execution_device": execution_device,
822
+ }
823
+ )
824
+
825
+ # Get adapters if applicable
826
+ if hasattr(component, "peft_config"):
827
+ info["adapters"] = list(component.peft_config.keys())
828
+
829
+ # Check for IP-Adapter scales
830
+ if hasattr(component, "_load_ip_adapter_weights") and hasattr(component, "attn_processors"):
831
+ processors = copy.deepcopy(component.attn_processors)
832
+ # First check if any processor is an IP-Adapter
833
+ processor_types = [v.__class__.__name__ for v in processors.values()]
834
+ if any("IPAdapter" in ptype for ptype in processor_types):
835
+ # Then get scales only from IP-Adapter processors
836
+ scales = {
837
+ k: v.scale
838
+ for k, v in processors.items()
839
+ if hasattr(v, "scale") and "IPAdapter" in v.__class__.__name__
840
+ }
841
+ if scales:
842
+ info["ip_adapter"] = summarize_dict_by_value_and_parts(scales)
843
+
844
+ # Check for quantization
845
+ hf_quantizer = getattr(component, "hf_quantizer", None)
846
+ if hf_quantizer is not None:
847
+ quant_config = hf_quantizer.quantization_config
848
+ if hasattr(quant_config, "to_diff_dict"):
849
+ info["quantization"] = quant_config.to_diff_dict()
850
+ else:
851
+ info["quantization"] = quant_config.to_dict()
852
+ else:
853
+ info["quantization"] = None
854
+
855
+ # If fields specified, filter info
856
+ if fields is not None:
857
+ return {k: v for k, v in info.items() if k in fields}
858
+ else:
859
+ return info
860
+
861
+ # YiYi TODO: (1) add display fields, allow user to set which fields to display in the comnponents table
862
+ def __repr__(self):
863
+ # Handle empty components case
864
+ if not self.components:
865
+ return "Components:\n" + "=" * 50 + "\nNo components registered.\n" + "=" * 50
866
+
867
+ # Extract load_id if available
868
+ def get_load_id(component):
869
+ if hasattr(component, "_diffusers_load_id"):
870
+ return component._diffusers_load_id
871
+ return "N/A"
872
+
873
+ # Format device info compactly
874
+ def format_device(component, info):
875
+ if not info["has_hook"]:
876
+ return str(getattr(component, "device", "N/A"))
877
+ else:
878
+ device = str(getattr(component, "device", "N/A"))
879
+ exec_device = str(info["execution_device"] or "N/A")
880
+ return f"{device}({exec_device})"
881
+
882
+ # Get max length of load_ids for models
883
+ load_ids = [
884
+ get_load_id(component)
885
+ for component in self.components.values()
886
+ if isinstance(component, torch.nn.Module) and hasattr(component, "_diffusers_load_id")
887
+ ]
888
+ max_load_id_len = max([15] + [len(str(lid)) for lid in load_ids]) if load_ids else 15
889
+
890
+ # Get all collections for each component
891
+ component_collections = {}
892
+ for name in self.components.keys():
893
+ component_collections[name] = []
894
+ for coll, comps in self.collections.items():
895
+ if name in comps:
896
+ component_collections[name].append(coll)
897
+ if not component_collections[name]:
898
+ component_collections[name] = ["N/A"]
899
+
900
+ # Find the maximum collection name length
901
+ all_collections = [coll for colls in component_collections.values() for coll in colls]
902
+ max_collection_len = max(10, max(len(str(c)) for c in all_collections)) if all_collections else 10
903
+
904
+ col_widths = {
905
+ "id": max(15, max(len(name) for name in self.components.keys())),
906
+ "class": max(25, max(len(component.__class__.__name__) for component in self.components.values())),
907
+ "device": 20,
908
+ "dtype": 15,
909
+ "size": 10,
910
+ "load_id": max_load_id_len,
911
+ "collection": max_collection_len,
912
+ }
913
+
914
+ # Create the header lines
915
+ sep_line = "=" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n"
916
+ dash_line = "-" * (sum(col_widths.values()) + len(col_widths) * 3 - 1) + "\n"
917
+
918
+ output = "Components:\n" + sep_line
919
+
920
+ # Separate components into models and others
921
+ models = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)}
922
+ others = {k: v for k, v in self.components.items() if not isinstance(v, torch.nn.Module)}
923
+
924
+ # Models section
925
+ if models:
926
+ output += "Models:\n" + dash_line
927
+ # Column headers
928
+ output += f"{'Name_ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | "
929
+ output += f"{'Device: act(exec)':<{col_widths['device']}} | {'Dtype':<{col_widths['dtype']}} | "
930
+ output += f"{'Size (GB)':<{col_widths['size']}} | {'Load ID':<{col_widths['load_id']}} | Collection\n"
931
+ output += dash_line
932
+
933
+ # Model entries
934
+ for name, component in models.items():
935
+ info = self.get_model_info(name)
936
+ device_str = format_device(component, info)
937
+ dtype = str(component.dtype) if hasattr(component, "dtype") else "N/A"
938
+ load_id = get_load_id(component)
939
+
940
+ # Print first collection on the main line
941
+ first_collection = component_collections[name][0] if component_collections[name] else "N/A"
942
+
943
+ output += f"{name:<{col_widths['id']}} | {info['class_name']:<{col_widths['class']}} | "
944
+ output += f"{device_str:<{col_widths['device']}} | {dtype:<{col_widths['dtype']}} | "
945
+ output += f"{info['size_gb']:<{col_widths['size']}.2f} | {load_id:<{col_widths['load_id']}} | {first_collection}\n"
946
+
947
+ # Print additional collections on separate lines if they exist
948
+ for i in range(1, len(component_collections[name])):
949
+ collection = component_collections[name][i]
950
+ output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | "
951
+ output += f"{'':<{col_widths['device']}} | {'':<{col_widths['dtype']}} | "
952
+ output += f"{'':<{col_widths['size']}} | {'':<{col_widths['load_id']}} | {collection}\n"
953
+
954
+ output += dash_line
955
+
956
+ # Other components section
957
+ if others:
958
+ if models: # Add extra newline if we had models section
959
+ output += "\n"
960
+ output += "Other Components:\n" + dash_line
961
+ # Column headers for other components
962
+ output += f"{'ID':<{col_widths['id']}} | {'Class':<{col_widths['class']}} | Collection\n"
963
+ output += dash_line
964
+
965
+ # Other component entries
966
+ for name, component in others.items():
967
+ info = self.get_model_info(name)
968
+
969
+ # Print first collection on the main line
970
+ first_collection = component_collections[name][0] if component_collections[name] else "N/A"
971
+
972
+ output += f"{name:<{col_widths['id']}} | {component.__class__.__name__:<{col_widths['class']}} | {first_collection}\n"
973
+
974
+ # Print additional collections on separate lines if they exist
975
+ for i in range(1, len(component_collections[name])):
976
+ collection = component_collections[name][i]
977
+ output += f"{'':<{col_widths['id']}} | {'':<{col_widths['class']}} | {collection}\n"
978
+
979
+ output += dash_line
980
+
981
+ # Add additional component info
982
+ output += "\nAdditional Component Info:\n" + "=" * 50 + "\n"
983
+ for name in self.components:
984
+ info = self.get_model_info(name)
985
+ if info is not None and (
986
+ info.get("adapters") is not None or info.get("ip_adapter") or info.get("quantization")
987
+ ):
988
+ output += f"\n{name}:\n"
989
+ if info.get("adapters") is not None:
990
+ output += f" Adapters: {info['adapters']}\n"
991
+ if info.get("ip_adapter"):
992
+ output += " IP-Adapter: Enabled\n"
993
+ if info.get("quantization"):
994
+ output += f" Quantization: {info['quantization']}\n"
995
+
996
+ return output
997
+
998
+ def get_one(
999
+ self,
1000
+ component_id: str | None = None,
1001
+ name: str | None = None,
1002
+ collection: str | None = None,
1003
+ load_id: str | None = None,
1004
+ ) -> Any:
1005
+ """
1006
+ Get a single component by either:
1007
+ - searching name (pattern matching), collection, or load_id.
1008
+ - passing in a component_id
1009
+ Raises an error if multiple components match or none are found.
1010
+
1011
+ Args:
1012
+ component_id (str | None): Optional component ID to get
1013
+ name (str | None): Component name or pattern
1014
+ collection (str | None): Optional collection to filter by
1015
+ load_id (str | None): Optional load_id to filter by
1016
+
1017
+ Returns:
1018
+ A single component
1019
+
1020
+ Raises:
1021
+ ValueError: If no components match or multiple components match
1022
+ """
1023
+
1024
+ if component_id is not None and (name is not None or collection is not None or load_id is not None):
1025
+ raise ValueError("If searching by component_id, do not pass name, collection, or load_id")
1026
+
1027
+ # search by component_id
1028
+ if component_id is not None:
1029
+ if component_id not in self.components:
1030
+ raise ValueError(f"Component '{component_id}' not found in ComponentsManager")
1031
+ return self.components[component_id]
1032
+ # search with name/collection/load_id
1033
+ results = self.search_components(name, collection, load_id)
1034
+
1035
+ if not results:
1036
+ raise ValueError(f"No components found matching '{name}'")
1037
+
1038
+ if len(results) > 1:
1039
+ raise ValueError(f"Multiple components found matching '{name}': {list(results.keys())}")
1040
+
1041
+ return next(iter(results.values()))
1042
+
1043
+ def get_ids(self, names: str | list[str] = None, collection: str | None = None):
1044
+ """
1045
+ Get component IDs by a list of names, optionally filtered by collection.
1046
+
1047
+ Args:
1048
+ names (str | list[str]): list of component names
1049
+ collection (str | None): Optional collection to filter by
1050
+
1051
+ Returns:
1052
+ list[str]: list of component IDs
1053
+ """
1054
+ ids = set()
1055
+ if not isinstance(names, list):
1056
+ names = [names]
1057
+ for name in names:
1058
+ ids.update(self._lookup_ids(name=name, collection=collection))
1059
+ return list(ids)
1060
+
1061
+ def get_components_by_ids(self, ids: list[str], return_dict_with_names: bool | None = True):
1062
+ """
1063
+ Get components by a list of IDs.
1064
+
1065
+ Args:
1066
+ ids (list[str]):
1067
+ list of component IDs
1068
+ return_dict_with_names (bool | None):
1069
+ Whether to return a dictionary with component names as keys:
1070
+
1071
+ Returns:
1072
+ dict[str, Any]: Dictionary of components.
1073
+ - If return_dict_with_names=True, keys are component names.
1074
+ - If return_dict_with_names=False, keys are component IDs.
1075
+
1076
+ Raises:
1077
+ ValueError: If duplicate component names are found in the search results when return_dict_with_names=True
1078
+ """
1079
+ components = {id: self.components[id] for id in ids}
1080
+
1081
+ if return_dict_with_names:
1082
+ dict_to_return = {}
1083
+ for comp_id, comp in components.items():
1084
+ comp_name = self._id_to_name(comp_id)
1085
+ if comp_name in dict_to_return:
1086
+ raise ValueError(
1087
+ 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"
1088
+ )
1089
+ dict_to_return[comp_name] = comp
1090
+ return dict_to_return
1091
+ else:
1092
+ return components
1093
+
1094
+ def get_components_by_names(self, names: list[str], collection: str | None = None):
1095
+ """
1096
+ Get components by a list of names, optionally filtered by collection.
1097
+
1098
+ Args:
1099
+ names (list[str]): list of component names
1100
+ collection (str | None): Optional collection to filter by
1101
+
1102
+ Returns:
1103
+ dict[str, Any]: Dictionary of components with component names as keys
1104
+
1105
+ Raises:
1106
+ ValueError: If duplicate component names are found in the search results
1107
+ """
1108
+ ids = self.get_ids(names, collection)
1109
+ return self.get_components_by_ids(ids)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ...utils import (
4
+ DIFFUSERS_SLOW_IMPORT,
5
+ OptionalDependencyNotAvailable,
6
+ _LazyModule,
7
+ get_objects_from_module,
8
+ is_torch_available,
9
+ is_transformers_available,
10
+ )
11
+
12
+
13
+ _dummy_objects = {}
14
+ _import_structure = {}
15
+
16
+ try:
17
+ if not (is_transformers_available() and is_torch_available()):
18
+ raise OptionalDependencyNotAvailable()
19
+ except OptionalDependencyNotAvailable:
20
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
21
+
22
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
23
+ else:
24
+ _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"]
25
+ _import_structure["modular_blocks_cosmos3_distilled"] = ["Cosmos3DistilledBlocks"]
26
+ _import_structure["modular_pipeline"] = ["Cosmos3DistilledModularPipeline", "Cosmos3OmniModularPipeline"]
27
+
28
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
29
+ try:
30
+ if not (is_transformers_available() and is_torch_available()):
31
+ raise OptionalDependencyNotAvailable()
32
+ except OptionalDependencyNotAvailable:
33
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
34
+ else:
35
+ from .modular_blocks_cosmos3 import Cosmos3OmniBlocks
36
+ from .modular_blocks_cosmos3_distilled import Cosmos3DistilledBlocks
37
+ from .modular_pipeline import Cosmos3DistilledModularPipeline, Cosmos3OmniModularPipeline
38
+ else:
39
+ import sys
40
+
41
+ sys.modules[__name__] = _LazyModule(
42
+ __name__,
43
+ globals()["__file__"],
44
+ _import_structure,
45
+ module_spec=__spec__,
46
+ )
47
+
48
+ for name, value in _dummy_objects.items():
49
+ setattr(sys.modules[__name__], name, value)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ...utils import encode_video, export_to_video
4
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
5
+ from ..modular_pipeline_utils import InputParam, OutputParam
6
+ from .modular_pipeline import Cosmos3OmniModularPipeline
7
+
8
+
9
+ class Cosmos3ActionOutputStep(ModularPipelineBlocks):
10
+ model_name = "cosmos3-omni"
11
+
12
+ @property
13
+ def description(self) -> str:
14
+ return "Post-processes action latents into action outputs."
15
+
16
+ @property
17
+ def inputs(self) -> list[InputParam]:
18
+ return [
19
+ InputParam(
20
+ name="action_latents",
21
+ type_hint=torch.Tensor,
22
+ default=None,
23
+ description="Denoised action latents.",
24
+ ),
25
+ InputParam(
26
+ name="action_mode", type_hint=str, default=None, description="Requested action-generation mode."
27
+ ),
28
+ InputParam(
29
+ name="raw_action_dim_resolved",
30
+ type_hint=int,
31
+ default=None,
32
+ description="Unpadded action-vector dimension.",
33
+ ),
34
+ ]
35
+
36
+ @property
37
+ def intermediate_outputs(self) -> list[OutputParam]:
38
+ return [OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors.")]
39
+
40
+ @torch.no_grad()
41
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
42
+ block_state = self.get_block_state(state)
43
+ action_output = None
44
+ if block_state.action_mode in {"inverse_dynamics", "policy"} and block_state.action_latents is not None:
45
+ action_output = block_state.action_latents
46
+ if block_state.raw_action_dim_resolved is not None:
47
+ action_output = action_output[:, : block_state.raw_action_dim_resolved]
48
+ action_output = [action_output.detach().cpu()]
49
+ block_state.action = action_output
50
+ self.set_block_state(state, block_state)
51
+ return components, state
52
+
53
+
54
+ class Cosmos3ExportStep(ModularPipelineBlocks):
55
+ model_name = "cosmos3-omni"
56
+
57
+ @property
58
+ def description(self) -> str:
59
+ return (
60
+ "Optional export block that writes decoded outputs to disk. Writes `videos` to `output_path` via "
61
+ "`export_to_video`, or muxes `videos` with `sound` via `encode_video` when a waveform is present. "
62
+ "Not wired into the default blocks; add it explicitly when you want the pipeline to produce a file."
63
+ )
64
+
65
+ @property
66
+ def inputs(self) -> list[InputParam]:
67
+ return [
68
+ InputParam(name="videos", required=True, description="Generated video frames to export."),
69
+ InputParam(
70
+ name="output_path",
71
+ type_hint=str,
72
+ required=True,
73
+ description="Destination path for the exported video.",
74
+ ),
75
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the exported video."),
76
+ InputParam(
77
+ name="sound",
78
+ type_hint=torch.Tensor,
79
+ default=None,
80
+ description="Generated waveform to mux into the video.",
81
+ ),
82
+ InputParam(
83
+ name="sampling_rate",
84
+ type_hint=int,
85
+ default=None,
86
+ description="Sample rate of the generated waveform in Hz.",
87
+ ),
88
+ ]
89
+
90
+ @property
91
+ def intermediate_outputs(self) -> list[OutputParam]:
92
+ return [OutputParam("output_path", type_hint=str, description="Path of the exported video file.")]
93
+
94
+ @torch.no_grad()
95
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
96
+ block_state = self.get_block_state(state)
97
+ output_path = str(block_state.output_path)
98
+ fps = int(round(block_state.fps))
99
+ if block_state.sound is not None:
100
+ if block_state.sampling_rate is None:
101
+ raise ValueError("`sampling_rate` is required to export a video with sound.")
102
+ encode_video(
103
+ block_state.videos,
104
+ fps=fps,
105
+ audio=block_state.sound,
106
+ audio_sample_rate=int(block_state.sampling_rate),
107
+ output_path=output_path,
108
+ )
109
+ else:
110
+ export_to_video(block_state.videos, output_path, fps=fps, macro_block_size=1)
111
+ block_state.output_path = output_path
112
+ self.set_block_state(state, block_state)
113
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py ADDED
@@ -0,0 +1,1331 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+
3
+ import numpy as np
4
+ import torch
5
+
6
+ from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
7
+ from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
8
+ from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
9
+ from ...utils.torch_utils import randn_tensor
10
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
11
+ from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
12
+ from .modular_pipeline import Cosmos3OmniModularPipeline
13
+
14
+
15
+ class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks):
16
+ model_name = "cosmos3-omni"
17
+
18
+ @property
19
+ def description(self) -> str:
20
+ return "Builds cond/uncond text segments before denoising."
21
+
22
+ @property
23
+ def expected_components(self) -> list[ComponentSpec]:
24
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
25
+
26
+ @property
27
+ def inputs(self) -> list[InputParam]:
28
+ return [
29
+ InputParam(name="cond_input_ids", required=True, description="Token IDs for the conditional prompt."),
30
+ InputParam(name="uncond_input_ids", required=True, description="Token IDs for the unconditional prompt."),
31
+ ]
32
+
33
+ @property
34
+ def intermediate_outputs(self) -> list[OutputParam]:
35
+ return [
36
+ OutputParam(
37
+ "cond_text_segment",
38
+ type_hint=dict,
39
+ kwargs_type="denoiser_input_fields",
40
+ description="Conditional text segment for the denoiser.",
41
+ ),
42
+ OutputParam(
43
+ "uncond_text_segment",
44
+ type_hint=dict,
45
+ kwargs_type="denoiser_input_fields",
46
+ description="Unconditional text segment for the denoiser.",
47
+ ),
48
+ ]
49
+
50
+ @torch.no_grad()
51
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
52
+ block_state = self.get_block_state(state)
53
+ device = components._execution_device
54
+ block_state.cond_text_segment = components._prepare_text_segment(block_state.cond_input_ids, device=device)
55
+ block_state.uncond_text_segment = components._prepare_text_segment(block_state.uncond_input_ids, device=device)
56
+ self.set_block_state(state, block_state)
57
+ return components, state
58
+
59
+
60
+ class Cosmos3VisionPrepareLatentsStep(ModularPipelineBlocks):
61
+ model_name = "cosmos3-omni"
62
+
63
+ @property
64
+ def description(self) -> str:
65
+ return "Prepares noisy vision latents and the vision conditioning mask."
66
+
67
+ @property
68
+ def expected_components(self) -> list[ComponentSpec]:
69
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
70
+
71
+ @property
72
+ def inputs(self) -> list[InputParam]:
73
+ return [
74
+ InputParam(
75
+ name="x0_tokens_vision",
76
+ type_hint=torch.Tensor,
77
+ default=None,
78
+ description="Vision latents encoded from the conditioning image or video.",
79
+ ),
80
+ InputParam(
81
+ name="vision_condition_frames",
82
+ type_hint=list[int],
83
+ default=None,
84
+ description="Latent-frame indexes fixed by visual conditioning.",
85
+ ),
86
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
87
+ InputParam(
88
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
89
+ ),
90
+ InputParam(
91
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
92
+ ),
93
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
94
+ InputParam(
95
+ name="latents",
96
+ type_hint=torch.Tensor,
97
+ default=None,
98
+ description="Pre-generated noisy vision latents.",
99
+ ),
100
+ InputParam.template("generator"),
101
+ ]
102
+
103
+ @property
104
+ def intermediate_outputs(self) -> list[OutputParam]:
105
+ return [
106
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy vision latents for denoising."),
107
+ OutputParam("fps_vision", type_hint=float, description="Frame rate used to pack vision latents."),
108
+ OutputParam(
109
+ "vision_condition_mask",
110
+ type_hint=torch.Tensor,
111
+ kwargs_type="denoiser_input_fields",
112
+ description="Mask marking conditioned vision latent frames.",
113
+ ),
114
+ OutputParam(
115
+ "vision_condition_indexes_for_pack",
116
+ type_hint=list[int],
117
+ description="Indexes of conditioned vision latent frames.",
118
+ ),
119
+ OutputParam(
120
+ "vision_conditioning_latents",
121
+ type_hint=torch.Tensor,
122
+ description="Clean encoded vision latents used to re-anchor image conditioning each step.",
123
+ ),
124
+ ]
125
+
126
+ @torch.no_grad()
127
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
128
+ block_state = self.get_block_state(state)
129
+ device = components._execution_device
130
+ dtype = components.transformer.dtype
131
+
132
+ x0_tokens_vision = block_state.x0_tokens_vision
133
+ if x0_tokens_vision is None:
134
+ if block_state.num_frames < 1:
135
+ raise ValueError(f"num_frames must be >= 1, got {block_state.num_frames}.")
136
+ sf_spatial = components.vae_scale_factor_spatial
137
+ if block_state.height % sf_spatial != 0 or block_state.width % sf_spatial != 0:
138
+ raise ValueError(
139
+ f"height and width must be multiples of {sf_spatial}, got ({block_state.height}, {block_state.width})."
140
+ )
141
+ latent_shape = (
142
+ 1,
143
+ components.num_channels_latents,
144
+ (block_state.num_frames - 1) // components.vae_scale_factor_temporal + 1,
145
+ block_state.height // sf_spatial,
146
+ block_state.width // sf_spatial,
147
+ )
148
+ x0_tokens_vision = torch.zeros(latent_shape, device=device, dtype=torch.float32)
149
+ else:
150
+ x0_tokens_vision = x0_tokens_vision.to(device=device, dtype=torch.float32)
151
+
152
+ block_state.fps_vision = float(block_state.fps)
153
+ condition_frames = block_state.vision_condition_frames or []
154
+ block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype)
155
+ for frame_idx in condition_frames:
156
+ if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]:
157
+ block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0
158
+
159
+ if block_state.latents is None:
160
+ pure_noise = randn_tensor(
161
+ tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype
162
+ )
163
+ block_state.latents = (
164
+ block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype)
165
+ + (1.0 - block_state.vision_condition_mask) * pure_noise
166
+ )
167
+ else:
168
+ block_state.latents = block_state.latents.to(device=device, dtype=dtype)
169
+
170
+ vision_condition_indexes = torch.nonzero(
171
+ block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False
172
+ ).flatten()
173
+ block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes]
174
+ block_state.vision_conditioning_latents = x0_tokens_vision
175
+
176
+ self.set_block_state(state, block_state)
177
+ return components, state
178
+
179
+
180
+ class Cosmos3SoundPrepareLatentsStep(ModularPipelineBlocks):
181
+ model_name = "cosmos3-omni"
182
+
183
+ @property
184
+ def description(self) -> str:
185
+ return "Prepares noisy sound latents and the sound conditioning mask."
186
+
187
+ @property
188
+ def expected_components(self) -> list[ComponentSpec]:
189
+ return [
190
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
191
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
192
+ ]
193
+
194
+ @property
195
+ def inputs(self) -> list[InputParam]:
196
+ return [
197
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
198
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
199
+ InputParam(
200
+ name="sound_latents",
201
+ type_hint=torch.Tensor,
202
+ default=None,
203
+ description="Pre-generated noisy sound latents.",
204
+ ),
205
+ InputParam.template("generator"),
206
+ ]
207
+
208
+ @property
209
+ def intermediate_outputs(self) -> list[OutputParam]:
210
+ return [
211
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Noisy sound latents for denoising."),
212
+ OutputParam("fps_sound", type_hint=float, description="Frame rate of the sound latent sequence."),
213
+ OutputParam(
214
+ "sound_condition_mask",
215
+ type_hint=torch.Tensor,
216
+ kwargs_type="denoiser_input_fields",
217
+ description="Mask marking conditioned sound latent frames.",
218
+ ),
219
+ OutputParam("sound_scheduler", description="Scheduler used to update sound latents."),
220
+ ]
221
+
222
+ @torch.no_grad()
223
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
224
+ block_state = self.get_block_state(state)
225
+ device = components._execution_device
226
+ dtype = components.transformer.dtype
227
+
228
+ if not components.transformer.config.sound_gen:
229
+ raise ValueError("Sound generation requires a transformer trained with sound_gen=True.")
230
+
231
+ sound_dim = components.transformer.config.sound_dim
232
+ block_state.fps_sound = float(components.transformer.config.sound_latent_fps)
233
+ n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate)
234
+ hop_size = components.sound_hop_size
235
+ t_sound = (n_audio_samples + hop_size - 1) // hop_size
236
+ x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype)
237
+ block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype)
238
+
239
+ if block_state.sound_latents is None:
240
+ pure_noise = randn_tensor(
241
+ tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype
242
+ )
243
+ block_state.sound_latents = (
244
+ block_state.sound_condition_mask.T * x0_tokens_sound
245
+ + (1.0 - block_state.sound_condition_mask.T) * pure_noise
246
+ )
247
+ else:
248
+ block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype)
249
+
250
+ block_state.sound_scheduler = copy.deepcopy(components.scheduler)
251
+
252
+ self.set_block_state(state, block_state)
253
+ return components, state
254
+
255
+
256
+ class Cosmos3ActionPrepareLatentsStep(ModularPipelineBlocks):
257
+ model_name = "cosmos3-omni"
258
+
259
+ @property
260
+ def description(self) -> str:
261
+ return "Prepares noisy action latents and the action conditioning mask."
262
+
263
+ @property
264
+ def expected_components(self) -> list[ComponentSpec]:
265
+ return [
266
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
267
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
268
+ ]
269
+
270
+ @property
271
+ def inputs(self) -> list[InputParam]:
272
+ return [
273
+ InputParam(
274
+ name="action",
275
+ type_hint=CosmosActionCondition,
276
+ required=True,
277
+ description="Action-conditioning metadata.",
278
+ ),
279
+ InputParam(
280
+ name="action_condition_frame_indexes",
281
+ type_hint=list[int],
282
+ default=None,
283
+ description="Action-frame indexes fixed by action conditioning.",
284
+ ),
285
+ InputParam(
286
+ name="action_latents",
287
+ type_hint=torch.Tensor,
288
+ default=None,
289
+ description="Pre-generated noisy action latents.",
290
+ ),
291
+ InputParam.template("generator"),
292
+ ]
293
+
294
+ @property
295
+ def intermediate_outputs(self) -> list[OutputParam]:
296
+ return [
297
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Noisy action latents for denoising."),
298
+ OutputParam(
299
+ "action_condition_mask",
300
+ type_hint=torch.Tensor,
301
+ kwargs_type="denoiser_input_fields",
302
+ description="Mask marking conditioned action latent frames.",
303
+ ),
304
+ OutputParam(
305
+ "action_domain_ids",
306
+ type_hint=list[torch.Tensor],
307
+ kwargs_type="denoiser_input_fields",
308
+ description="Embodiment domain IDs for action conditioning.",
309
+ ),
310
+ OutputParam(
311
+ "raw_action_dim_resolved",
312
+ type_hint=int,
313
+ kwargs_type="denoiser_input_fields",
314
+ description="Unpadded action-vector dimension.",
315
+ ),
316
+ OutputParam("action_scheduler", description="Scheduler used to update action latents."),
317
+ ]
318
+
319
+ @torch.no_grad()
320
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
321
+ block_state = self.get_block_state(state)
322
+ device = components._execution_device
323
+ dtype = components.transformer.dtype
324
+ action = block_state.action
325
+
326
+ if not components.transformer.config.action_gen:
327
+ raise ValueError("action requires a transformer trained with action_gen=True.")
328
+
329
+ block_state.raw_action_dim_resolved = int(action.raw_action_dim) if action.raw_action_dim is not None else None
330
+ if (
331
+ block_state.raw_action_dim_resolved is not None
332
+ and block_state.raw_action_dim_resolved > components.transformer.config.action_dim
333
+ ):
334
+ raise ValueError(
335
+ f"raw_action_dim={block_state.raw_action_dim_resolved} exceeds the model action_dim="
336
+ f"{components.transformer.config.action_dim}."
337
+ )
338
+
339
+ action_chunk_size = action.chunk_size
340
+ action_dim = components.transformer.action_dim
341
+ if action.mode == "forward_dynamics":
342
+ raw_actions = action.raw_actions
343
+ if raw_actions is None:
344
+ raise ValueError("action_mode='forward_dynamics' requires an action tensor.")
345
+ raw_actions = raw_actions.to(device=device, dtype=dtype)
346
+ if raw_actions.shape[-1] > action_dim:
347
+ raise ValueError(
348
+ f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}."
349
+ )
350
+ if raw_actions.shape[0] < action_chunk_size:
351
+ raw_actions = torch.cat(
352
+ [raw_actions, raw_actions[-1:].expand(action_chunk_size - raw_actions.shape[0], -1)],
353
+ dim=0,
354
+ )
355
+ raw_actions = raw_actions[:action_chunk_size]
356
+ if raw_actions.shape[-1] < action_dim:
357
+ action_padding = torch.zeros(
358
+ raw_actions.shape[0],
359
+ action_dim - raw_actions.shape[-1],
360
+ dtype=raw_actions.dtype,
361
+ device=raw_actions.device,
362
+ )
363
+ raw_actions = torch.cat([raw_actions, action_padding], dim=-1)
364
+ x0_tokens_action = raw_actions
365
+ else:
366
+ x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype)
367
+
368
+ if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID:
369
+ raise ValueError(
370
+ f"Unknown Cosmos3 action domain_name={action.domain_name!r}; expected one of {sorted(_EMBODIMENT_TO_DOMAIN_ID)}."
371
+ )
372
+ block_state.action_domain_ids = [
373
+ torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device)
374
+ ]
375
+ condition_frames = block_state.action_condition_frame_indexes or []
376
+ block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype)
377
+ for frame_idx in condition_frames:
378
+ if 0 <= frame_idx < block_state.action_condition_mask.shape[0]:
379
+ block_state.action_condition_mask[frame_idx, 0] = 1.0
380
+
381
+ if block_state.action_latents is None:
382
+ pure_noise = randn_tensor(
383
+ tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype
384
+ )
385
+ block_state.action_latents = (
386
+ block_state.action_condition_mask * x0_tokens_action
387
+ + (1.0 - block_state.action_condition_mask) * pure_noise
388
+ )
389
+ if block_state.raw_action_dim_resolved is not None:
390
+ block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0
391
+ else:
392
+ block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype)
393
+
394
+ block_state.action_scheduler = copy.deepcopy(components.scheduler)
395
+
396
+ self.set_block_state(state, block_state)
397
+ return components, state
398
+
399
+
400
+ class Cosmos3VisionPackSequenceStep(ModularPipelineBlocks):
401
+ model_name = "cosmos3-omni"
402
+
403
+ @property
404
+ def description(self) -> str:
405
+ return "Builds separate cond/uncond vision sequence segments."
406
+
407
+ @property
408
+ def expected_components(self) -> list[ComponentSpec]:
409
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
410
+
411
+ @property
412
+ def inputs(self) -> list[InputParam]:
413
+ return [
414
+ InputParam(
415
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
416
+ ),
417
+ InputParam(
418
+ name="uncond_text_segment",
419
+ type_hint=dict,
420
+ required=True,
421
+ description="Unconditional text segment.",
422
+ ),
423
+ InputParam(
424
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy vision latents to pack."
425
+ ),
426
+ InputParam(
427
+ name="fps_vision",
428
+ type_hint=float,
429
+ required=True,
430
+ description="Frame rate used to pack vision latents.",
431
+ ),
432
+ InputParam(
433
+ name="vision_condition_indexes_for_pack",
434
+ type_hint=list[int],
435
+ required=True,
436
+ description="Indexes of conditioned vision latent frames.",
437
+ ),
438
+ ]
439
+
440
+ @property
441
+ def intermediate_outputs(self) -> list[OutputParam]:
442
+ return [
443
+ OutputParam(
444
+ "cond_vision_segment",
445
+ type_hint=dict,
446
+ kwargs_type="denoiser_input_fields",
447
+ description="Conditional vision segment for the denoiser.",
448
+ ),
449
+ OutputParam(
450
+ "uncond_vision_segment",
451
+ type_hint=dict,
452
+ kwargs_type="denoiser_input_fields",
453
+ description="Unconditional vision segment for the denoiser.",
454
+ ),
455
+ ]
456
+
457
+ @torch.no_grad()
458
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
459
+ block_state = self.get_block_state(state)
460
+ device = components._execution_device
461
+ has_image_condition = bool(block_state.vision_condition_indexes_for_pack)
462
+
463
+ block_state.cond_vision_segment = components._prepare_vision_segment(
464
+ input_vision_tokens=block_state.latents,
465
+ has_image_condition=has_image_condition,
466
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
467
+ vision_fps=block_state.fps_vision,
468
+ curr=block_state.cond_text_segment["und_len"],
469
+ device=device,
470
+ condition_frame_indexes=block_state.vision_condition_indexes_for_pack,
471
+ )
472
+ block_state.uncond_vision_segment = components._prepare_vision_segment(
473
+ input_vision_tokens=block_state.latents,
474
+ has_image_condition=has_image_condition,
475
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
476
+ vision_fps=block_state.fps_vision,
477
+ curr=block_state.uncond_text_segment["und_len"],
478
+ device=device,
479
+ condition_frame_indexes=block_state.vision_condition_indexes_for_pack,
480
+ )
481
+
482
+ self.set_block_state(state, block_state)
483
+ return components, state
484
+
485
+
486
+ class Cosmos3SoundPackSequenceStep(ModularPipelineBlocks):
487
+ model_name = "cosmos3-omni"
488
+
489
+ @property
490
+ def description(self) -> str:
491
+ return "Builds separate cond/uncond sound sequence segments."
492
+
493
+ @property
494
+ def expected_components(self) -> list[ComponentSpec]:
495
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
496
+
497
+ @property
498
+ def inputs(self) -> list[InputParam]:
499
+ return [
500
+ InputParam(
501
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
502
+ ),
503
+ InputParam(
504
+ name="uncond_text_segment",
505
+ type_hint=dict,
506
+ required=True,
507
+ description="Unconditional text segment.",
508
+ ),
509
+ InputParam(
510
+ name="cond_sequence_length",
511
+ type_hint=int,
512
+ required=True,
513
+ description="Conditional multimodal sequence length.",
514
+ ),
515
+ InputParam(
516
+ name="uncond_sequence_length",
517
+ type_hint=int,
518
+ required=True,
519
+ description="Unconditional multimodal sequence length.",
520
+ ),
521
+ InputParam(
522
+ name="sound_latents", type_hint=torch.Tensor, required=True, description="Noisy sound latents to pack."
523
+ ),
524
+ InputParam(
525
+ name="fps_sound",
526
+ type_hint=float,
527
+ required=True,
528
+ description="Frame rate of the sound latent sequence.",
529
+ ),
530
+ ]
531
+
532
+ @property
533
+ def intermediate_outputs(self) -> list[OutputParam]:
534
+ return [
535
+ OutputParam(
536
+ "cond_sound_segment",
537
+ type_hint=dict,
538
+ kwargs_type="denoiser_input_fields",
539
+ description="Conditional sound segment for the denoiser.",
540
+ ),
541
+ OutputParam(
542
+ "uncond_sound_segment",
543
+ type_hint=dict,
544
+ kwargs_type="denoiser_input_fields",
545
+ description="Unconditional sound segment for the denoiser.",
546
+ ),
547
+ ]
548
+
549
+ @torch.no_grad()
550
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
551
+ block_state = self.get_block_state(state)
552
+ device = components._execution_device
553
+
554
+ block_state.cond_sound_segment = components._prepare_sound_segment(
555
+ input_sound_tokens=block_state.sound_latents,
556
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
557
+ sound_fps=block_state.fps_sound,
558
+ curr=block_state.cond_sequence_length,
559
+ device=device,
560
+ )
561
+ block_state.uncond_sound_segment = components._prepare_sound_segment(
562
+ input_sound_tokens=block_state.sound_latents,
563
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
564
+ sound_fps=block_state.fps_sound,
565
+ curr=block_state.uncond_sequence_length,
566
+ device=device,
567
+ )
568
+
569
+ self.set_block_state(state, block_state)
570
+ return components, state
571
+
572
+
573
+ class Cosmos3ActionPackSequenceStep(ModularPipelineBlocks):
574
+ model_name = "cosmos3-omni"
575
+
576
+ @property
577
+ def description(self) -> str:
578
+ return "Builds separate cond/uncond action sequence segments."
579
+
580
+ @property
581
+ def expected_components(self) -> list[ComponentSpec]:
582
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
583
+
584
+ @property
585
+ def inputs(self) -> list[InputParam]:
586
+ return [
587
+ InputParam(
588
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
589
+ ),
590
+ InputParam(
591
+ name="uncond_text_segment",
592
+ type_hint=dict,
593
+ required=True,
594
+ description="Unconditional text segment.",
595
+ ),
596
+ InputParam(
597
+ name="cond_sequence_length",
598
+ type_hint=int,
599
+ required=True,
600
+ description="Conditional multimodal sequence length.",
601
+ ),
602
+ InputParam(
603
+ name="uncond_sequence_length",
604
+ type_hint=int,
605
+ required=True,
606
+ description="Unconditional multimodal sequence length.",
607
+ ),
608
+ InputParam(
609
+ name="action_latents",
610
+ type_hint=torch.Tensor,
611
+ required=True,
612
+ description="Noisy action latents to pack.",
613
+ ),
614
+ InputParam(
615
+ name="action_condition_frame_indexes",
616
+ type_hint=list[int],
617
+ default=None,
618
+ description="Action-frame indexes fixed by action conditioning.",
619
+ ),
620
+ InputParam(
621
+ name="fps_vision",
622
+ type_hint=float,
623
+ required=True,
624
+ description="Frame rate used to pack vision latents.",
625
+ ),
626
+ ]
627
+
628
+ @property
629
+ def intermediate_outputs(self) -> list[OutputParam]:
630
+ return [
631
+ OutputParam(
632
+ "cond_action_segment",
633
+ type_hint=dict,
634
+ kwargs_type="denoiser_input_fields",
635
+ description="Conditional action segment for the denoiser.",
636
+ ),
637
+ OutputParam(
638
+ "uncond_action_segment",
639
+ type_hint=dict,
640
+ kwargs_type="denoiser_input_fields",
641
+ description="Unconditional action segment for the denoiser.",
642
+ ),
643
+ ]
644
+
645
+ @torch.no_grad()
646
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
647
+ block_state = self.get_block_state(state)
648
+ device = components._execution_device
649
+
650
+ block_state.cond_action_segment = components._prepare_action_segment(
651
+ input_action_tokens=block_state.action_latents,
652
+ condition_frame_indexes=block_state.action_condition_frame_indexes,
653
+ mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"],
654
+ action_fps=block_state.fps_vision,
655
+ curr=block_state.cond_sequence_length,
656
+ device=device,
657
+ )
658
+ block_state.uncond_action_segment = components._prepare_action_segment(
659
+ input_action_tokens=block_state.action_latents,
660
+ condition_frame_indexes=block_state.action_condition_frame_indexes,
661
+ mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"],
662
+ action_fps=block_state.fps_vision,
663
+ curr=block_state.uncond_sequence_length,
664
+ device=device,
665
+ )
666
+
667
+ self.set_block_state(state, block_state)
668
+ return components, state
669
+
670
+
671
+ class Cosmos3VisionDenoiseInputStep(ModularPipelineBlocks):
672
+ model_name = "cosmos3-omni"
673
+
674
+ @property
675
+ def description(self) -> str:
676
+ return "Assembles text and vision sequence metadata for the denoising loop."
677
+
678
+ @property
679
+ def inputs(self) -> list[InputParam]:
680
+ return [
681
+ InputParam(
682
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
683
+ ),
684
+ InputParam(
685
+ name="uncond_text_segment",
686
+ type_hint=dict,
687
+ required=True,
688
+ description="Unconditional text segment.",
689
+ ),
690
+ InputParam(
691
+ name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment."
692
+ ),
693
+ InputParam(
694
+ name="uncond_vision_segment",
695
+ type_hint=dict,
696
+ required=True,
697
+ description="Unconditional vision segment.",
698
+ ),
699
+ ]
700
+
701
+ @property
702
+ def intermediate_outputs(self) -> list[OutputParam]:
703
+ return [
704
+ OutputParam(
705
+ "cond_position_ids",
706
+ type_hint=torch.Tensor,
707
+ kwargs_type="denoiser_input_fields",
708
+ description="Conditional multimodal RoPE position IDs.",
709
+ ),
710
+ OutputParam(
711
+ "uncond_position_ids",
712
+ type_hint=torch.Tensor,
713
+ kwargs_type="denoiser_input_fields",
714
+ description="Unconditional multimodal RoPE position IDs.",
715
+ ),
716
+ OutputParam(
717
+ "cond_sequence_length",
718
+ type_hint=int,
719
+ kwargs_type="denoiser_input_fields",
720
+ description="Conditional multimodal sequence length.",
721
+ ),
722
+ OutputParam(
723
+ "uncond_sequence_length",
724
+ type_hint=int,
725
+ kwargs_type="denoiser_input_fields",
726
+ description="Unconditional multimodal sequence length.",
727
+ ),
728
+ ]
729
+
730
+ @torch.no_grad()
731
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
732
+ block_state = self.get_block_state(state)
733
+ block_state.cond_position_ids = torch.cat(
734
+ [
735
+ block_state.cond_text_segment["text_mrope_ids"],
736
+ block_state.cond_vision_segment["vision_mrope_ids"],
737
+ ],
738
+ dim=1,
739
+ )
740
+ block_state.uncond_position_ids = torch.cat(
741
+ [
742
+ block_state.uncond_text_segment["text_mrope_ids"],
743
+ block_state.uncond_vision_segment["vision_mrope_ids"],
744
+ ],
745
+ dim=1,
746
+ )
747
+ block_state.cond_sequence_length = (
748
+ block_state.cond_text_segment["und_len"] + block_state.cond_vision_segment["num_vision_tokens"]
749
+ )
750
+ block_state.uncond_sequence_length = (
751
+ block_state.uncond_text_segment["und_len"] + block_state.uncond_vision_segment["num_vision_tokens"]
752
+ )
753
+ self.set_block_state(state, block_state)
754
+ return components, state
755
+
756
+
757
+ class Cosmos3SoundDenoiseInputStep(ModularPipelineBlocks):
758
+ model_name = "cosmos3-omni"
759
+
760
+ @property
761
+ def description(self) -> str:
762
+ return "Appends sound sequence metadata to the denoising-loop inputs."
763
+
764
+ @property
765
+ def inputs(self) -> list[InputParam]:
766
+ return [
767
+ InputParam(
768
+ name="cond_position_ids",
769
+ type_hint=torch.Tensor,
770
+ required=True,
771
+ description="Conditional multimodal RoPE position IDs.",
772
+ ),
773
+ InputParam(
774
+ name="uncond_position_ids",
775
+ type_hint=torch.Tensor,
776
+ required=True,
777
+ description="Unconditional multimodal RoPE position IDs.",
778
+ ),
779
+ InputParam(
780
+ name="cond_sequence_length",
781
+ type_hint=int,
782
+ required=True,
783
+ description="Conditional multimodal sequence length.",
784
+ ),
785
+ InputParam(
786
+ name="uncond_sequence_length",
787
+ type_hint=int,
788
+ required=True,
789
+ description="Unconditional multimodal sequence length.",
790
+ ),
791
+ InputParam(
792
+ name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment."
793
+ ),
794
+ InputParam(
795
+ name="uncond_sound_segment",
796
+ type_hint=dict,
797
+ required=True,
798
+ description="Unconditional sound segment.",
799
+ ),
800
+ ]
801
+
802
+ @property
803
+ def intermediate_outputs(self) -> list[OutputParam]:
804
+ return [
805
+ OutputParam(
806
+ "cond_position_ids",
807
+ type_hint=torch.Tensor,
808
+ kwargs_type="denoiser_input_fields",
809
+ description="Conditional multimodal RoPE position IDs.",
810
+ ),
811
+ OutputParam(
812
+ "uncond_position_ids",
813
+ type_hint=torch.Tensor,
814
+ kwargs_type="denoiser_input_fields",
815
+ description="Unconditional multimodal RoPE position IDs.",
816
+ ),
817
+ OutputParam(
818
+ "cond_sequence_length",
819
+ type_hint=int,
820
+ kwargs_type="denoiser_input_fields",
821
+ description="Conditional multimodal sequence length.",
822
+ ),
823
+ OutputParam(
824
+ "uncond_sequence_length",
825
+ type_hint=int,
826
+ kwargs_type="denoiser_input_fields",
827
+ description="Unconditional multimodal sequence length.",
828
+ ),
829
+ ]
830
+
831
+ @torch.no_grad()
832
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
833
+ block_state = self.get_block_state(state)
834
+ block_state.cond_position_ids = torch.cat(
835
+ [block_state.cond_position_ids, block_state.cond_sound_segment["sound_mrope_ids"]], dim=1
836
+ )
837
+ block_state.uncond_position_ids = torch.cat(
838
+ [block_state.uncond_position_ids, block_state.uncond_sound_segment["sound_mrope_ids"]], dim=1
839
+ )
840
+ block_state.cond_sequence_length += block_state.cond_sound_segment["sound_len"]
841
+ block_state.uncond_sequence_length += block_state.uncond_sound_segment["sound_len"]
842
+ self.set_block_state(state, block_state)
843
+ return components, state
844
+
845
+
846
+ class Cosmos3ActionDenoiseInputStep(ModularPipelineBlocks):
847
+ model_name = "cosmos3-omni"
848
+
849
+ @property
850
+ def description(self) -> str:
851
+ return "Appends action sequence metadata to the denoising-loop inputs."
852
+
853
+ @property
854
+ def inputs(self) -> list[InputParam]:
855
+ return [
856
+ InputParam(
857
+ name="cond_position_ids",
858
+ type_hint=torch.Tensor,
859
+ required=True,
860
+ description="Conditional multimodal RoPE position IDs.",
861
+ ),
862
+ InputParam(
863
+ name="uncond_position_ids",
864
+ type_hint=torch.Tensor,
865
+ required=True,
866
+ description="Unconditional multimodal RoPE position IDs.",
867
+ ),
868
+ InputParam(
869
+ name="cond_sequence_length",
870
+ type_hint=int,
871
+ required=True,
872
+ description="Conditional multimodal sequence length.",
873
+ ),
874
+ InputParam(
875
+ name="uncond_sequence_length",
876
+ type_hint=int,
877
+ required=True,
878
+ description="Unconditional multimodal sequence length.",
879
+ ),
880
+ InputParam(
881
+ name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment."
882
+ ),
883
+ InputParam(
884
+ name="uncond_action_segment",
885
+ type_hint=dict,
886
+ required=True,
887
+ description="Unconditional action segment.",
888
+ ),
889
+ ]
890
+
891
+ @property
892
+ def intermediate_outputs(self) -> list[OutputParam]:
893
+ return [
894
+ OutputParam(
895
+ "cond_position_ids",
896
+ type_hint=torch.Tensor,
897
+ kwargs_type="denoiser_input_fields",
898
+ description="Conditional multimodal RoPE position IDs.",
899
+ ),
900
+ OutputParam(
901
+ "uncond_position_ids",
902
+ type_hint=torch.Tensor,
903
+ kwargs_type="denoiser_input_fields",
904
+ description="Unconditional multimodal RoPE position IDs.",
905
+ ),
906
+ OutputParam(
907
+ "cond_sequence_length",
908
+ type_hint=int,
909
+ kwargs_type="denoiser_input_fields",
910
+ description="Conditional multimodal sequence length.",
911
+ ),
912
+ OutputParam(
913
+ "uncond_sequence_length",
914
+ type_hint=int,
915
+ kwargs_type="denoiser_input_fields",
916
+ description="Unconditional multimodal sequence length.",
917
+ ),
918
+ ]
919
+
920
+ @torch.no_grad()
921
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
922
+ block_state = self.get_block_state(state)
923
+ block_state.cond_position_ids = torch.cat(
924
+ [block_state.cond_position_ids, block_state.cond_action_segment["action_mrope_ids"]], dim=1
925
+ )
926
+ block_state.uncond_position_ids = torch.cat(
927
+ [block_state.uncond_position_ids, block_state.uncond_action_segment["action_mrope_ids"]], dim=1
928
+ )
929
+ block_state.cond_sequence_length += block_state.cond_action_segment["action_len"]
930
+ block_state.uncond_sequence_length += block_state.uncond_action_segment["action_len"]
931
+ self.set_block_state(state, block_state)
932
+ return components, state
933
+
934
+
935
+ class Cosmos3SetTimestepsStep(ModularPipelineBlocks):
936
+ model_name = "cosmos3-omni"
937
+
938
+ @property
939
+ def description(self) -> str:
940
+ return "Initializes scheduler timesteps."
941
+
942
+ @property
943
+ def expected_components(self) -> list[ComponentSpec]:
944
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
945
+
946
+ @property
947
+ def expected_configs(self) -> list[ConfigSpec]:
948
+ return [ConfigSpec(name="use_native_flow_schedule", default=False)]
949
+
950
+ @property
951
+ def inputs(self) -> list[InputParam]:
952
+ return [
953
+ InputParam.template("num_inference_steps", required=True),
954
+ ]
955
+
956
+ @property
957
+ def intermediate_outputs(self) -> list[OutputParam]:
958
+ return [
959
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."),
960
+ OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."),
961
+ ]
962
+
963
+ @torch.no_grad()
964
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
965
+ block_state = self.get_block_state(state)
966
+ device = components._execution_device
967
+ if components.config.use_native_flow_schedule:
968
+ sigmas = np.linspace(
969
+ 1.0 - 1.0 / components.scheduler.config.num_train_timesteps,
970
+ 0.0,
971
+ block_state.num_inference_steps + 1,
972
+ )[:-1]
973
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device, sigmas=sigmas)
974
+ else:
975
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
976
+ block_state.timesteps = components.scheduler.timesteps
977
+ block_state.num_warmup_steps = (
978
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
979
+ )
980
+ self.set_block_state(state, block_state)
981
+ return components, state
982
+
983
+
984
+ class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks):
985
+ model_name = "cosmos3-omni"
986
+
987
+ @property
988
+ def description(self) -> str:
989
+ return (
990
+ "Per-chunk transfer latent prep: takes the clean target latents encoded by "
991
+ "Cosmos3TransferChunkVaeEncoderStep and builds the noisy target latents, velocity mask, condition latents "
992
+ "and conditioned-frame indexes for this chunk."
993
+ )
994
+
995
+ @property
996
+ def expected_components(self) -> list[ComponentSpec]:
997
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
998
+
999
+ @property
1000
+ def inputs(self) -> list[InputParam]:
1001
+ return [
1002
+ InputParam(
1003
+ name="x0_tokens_vision",
1004
+ type_hint=torch.Tensor,
1005
+ required=True,
1006
+ description="Clean target vision latents encoded from the seeded target frames.",
1007
+ ),
1008
+ InputParam(
1009
+ name="current_conditional_frames",
1010
+ type_hint=int,
1011
+ required=True,
1012
+ description="Number of pixel frames used to seed this chunk's target.",
1013
+ ),
1014
+ InputParam.template("generator"),
1015
+ ]
1016
+
1017
+ @property
1018
+ def intermediate_outputs(self) -> list[OutputParam]:
1019
+ return [
1020
+ OutputParam("latents", type_hint=torch.Tensor, description="Noisy target latents for this chunk."),
1021
+ OutputParam(
1022
+ "velocity_mask",
1023
+ type_hint=torch.Tensor,
1024
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
1025
+ ),
1026
+ OutputParam(
1027
+ "condition_latents",
1028
+ type_hint=torch.Tensor,
1029
+ description="Clean target latents on the conditioned frames (the autoregressive seed).",
1030
+ ),
1031
+ OutputParam(
1032
+ "target_condition_indexes",
1033
+ type_hint=list[int],
1034
+ description="Latent-frame indexes fixed by the chunk's conditioning.",
1035
+ ),
1036
+ ]
1037
+
1038
+ @torch.no_grad()
1039
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1040
+ block_state = self.get_block_state(state)
1041
+ device = components._execution_device
1042
+ dtype = components.transformer.dtype
1043
+ tcf = components.vae_scale_factor_temporal
1044
+
1045
+ target_x0 = block_state.x0_tokens_vision.to(device=device)
1046
+ current_conditional_frames = block_state.current_conditional_frames
1047
+
1048
+ # Build the noisy target latents + conditioning mask from the clean target latents.
1049
+ latent_t = target_x0.shape[2]
1050
+ condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype)
1051
+ latent_condition_frames = 0
1052
+ if current_conditional_frames > 0:
1053
+ latent_condition_frames = (current_conditional_frames - 1) // tcf + 1
1054
+ condition_mask[:latent_condition_frames] = 1.0
1055
+ noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype)
1056
+ block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise
1057
+ block_state.velocity_mask = 1.0 - condition_mask
1058
+ block_state.condition_latents = condition_mask * target_x0
1059
+ block_state.target_condition_indexes = list(range(latent_condition_frames))
1060
+
1061
+ self.set_block_state(state, block_state)
1062
+ return components, state
1063
+
1064
+
1065
+ class Cosmos3TransferPackSequenceStep(ModularPipelineBlocks):
1066
+ model_name = "cosmos3-omni"
1067
+
1068
+ @property
1069
+ def description(self) -> str:
1070
+ return (
1071
+ "Pre-packs the three transfer CFG sequence variants: cond_full / uncond_full carry every control item, "
1072
+ "the no-control branch drops them (only [text, target]) so the control axis can be amplified."
1073
+ )
1074
+
1075
+ @property
1076
+ def inputs(self) -> list[InputParam]:
1077
+ return [
1078
+ InputParam(
1079
+ name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."
1080
+ ),
1081
+ InputParam(
1082
+ name="uncond_text_segment", type_hint=dict, required=True, description="Unconditional text segment."
1083
+ ),
1084
+ InputParam(
1085
+ name="control_latents",
1086
+ type_hint=list[torch.Tensor],
1087
+ required=True,
1088
+ description="Clean control latents for this chunk, one per hint in canonical order.",
1089
+ ),
1090
+ InputParam(
1091
+ name="latents",
1092
+ type_hint=torch.Tensor,
1093
+ required=True,
1094
+ description="Noisy target latents for this chunk.",
1095
+ ),
1096
+ InputParam(
1097
+ name="target_condition_indexes",
1098
+ type_hint=list[int],
1099
+ required=True,
1100
+ description="Latent-frame indexes fixed by the chunk's conditioning.",
1101
+ ),
1102
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
1103
+ ]
1104
+
1105
+ @property
1106
+ def intermediate_outputs(self) -> list[OutputParam]:
1107
+ return [
1108
+ OutputParam(
1109
+ "cond_full_static",
1110
+ type_hint=dict,
1111
+ kwargs_type="denoiser_input_fields",
1112
+ description="Conditional [control..., target] transfer sequence carrying every control item.",
1113
+ ),
1114
+ OutputParam(
1115
+ "cond_no_control_static",
1116
+ type_hint=dict,
1117
+ kwargs_type="denoiser_input_fields",
1118
+ description="Conditional [target] transfer sequence with the control items dropped.",
1119
+ ),
1120
+ OutputParam(
1121
+ "uncond_full_static",
1122
+ type_hint=dict,
1123
+ kwargs_type="denoiser_input_fields",
1124
+ description="Unconditional [control..., target] transfer sequence for text CFG.",
1125
+ ),
1126
+ OutputParam(
1127
+ "num_noisy_vision_tokens",
1128
+ type_hint=int,
1129
+ description="Number of noisy target vision tokens denoised each step.",
1130
+ ),
1131
+ ]
1132
+
1133
+ @torch.no_grad()
1134
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1135
+ block_state = self.get_block_state(state)
1136
+ device = components._execution_device
1137
+ num_hints = len(block_state.control_latents)
1138
+
1139
+ def _vision_pack(text_segment: dict, include_controls: bool) -> dict:
1140
+ if include_controls:
1141
+ vision_items = [*block_state.control_latents, block_state.latents]
1142
+ condition_indexes = [None] * num_hints + [block_state.target_condition_indexes]
1143
+ clean_flags = [True] * num_hints + [False]
1144
+ else:
1145
+ vision_items = [block_state.latents]
1146
+ condition_indexes = [block_state.target_condition_indexes]
1147
+ clean_flags = [False]
1148
+
1149
+ # Transfer packs [ctrl_1, ..., ctrl_N, target] into one vision segment
1150
+ mrope_offset = text_segment["vision_start_temporal_offset"]
1151
+ item_curr = text_segment["und_len"]
1152
+ token_shapes = []
1153
+ sequence_index_parts = []
1154
+ mse_loss_index_parts = []
1155
+ noisy_frame_indexes_per_item = []
1156
+ mrope_id_parts = []
1157
+ num_vision_tokens = 0
1158
+ num_noisy_vision_tokens = 0
1159
+ for item, item_condition, is_clean in zip(vision_items, condition_indexes, clean_flags):
1160
+ latent_t = item.shape[2]
1161
+ if is_clean:
1162
+ frame_condition = list(range(latent_t))
1163
+ else:
1164
+ frame_condition = item_condition if item_condition is not None else []
1165
+ item_segment = components._prepare_vision_segment(
1166
+ input_vision_tokens=item,
1167
+ has_image_condition=False,
1168
+ mrope_offset=mrope_offset,
1169
+ vision_fps=block_state.fps,
1170
+ curr=item_curr,
1171
+ device=device,
1172
+ condition_frame_indexes=frame_condition,
1173
+ )
1174
+ token_shapes.extend(item_segment["vision_token_shapes"])
1175
+ sequence_index_parts.append(item_segment["vision_sequence_indexes"])
1176
+ mse_loss_index_parts.append(item_segment["vision_mse_loss_indexes"])
1177
+ noisy_frame_indexes_per_item.extend(item_segment["vision_noisy_frame_indexes"])
1178
+ mrope_id_parts.append(item_segment["vision_mrope_ids"])
1179
+ num_vision_tokens += item_segment["num_vision_tokens"]
1180
+ num_noisy_vision_tokens += item_segment["num_noisy_vision_tokens"]
1181
+ item_curr += item_segment["num_vision_tokens"]
1182
+
1183
+ vision_segment = {
1184
+ "vision_token_shapes": token_shapes,
1185
+ "vision_sequence_indexes": torch.cat(sequence_index_parts, dim=0),
1186
+ "vision_mse_loss_indexes": torch.cat(mse_loss_index_parts, dim=0),
1187
+ "vision_noisy_frame_indexes": noisy_frame_indexes_per_item,
1188
+ "vision_mrope_ids": torch.cat(mrope_id_parts, dim=1),
1189
+ "num_vision_tokens": num_vision_tokens,
1190
+ "num_noisy_vision_tokens": num_noisy_vision_tokens,
1191
+ }
1192
+ return {
1193
+ **text_segment,
1194
+ **vision_segment,
1195
+ "position_ids": torch.cat([text_segment["text_mrope_ids"], vision_segment["vision_mrope_ids"]], dim=1),
1196
+ "sequence_length": text_segment["und_len"] + vision_segment["num_vision_tokens"],
1197
+ }
1198
+
1199
+ block_state.cond_full_static = _vision_pack(block_state.cond_text_segment, include_controls=True)
1200
+ block_state.cond_no_control_static = _vision_pack(block_state.cond_text_segment, include_controls=False)
1201
+ block_state.uncond_full_static = _vision_pack(block_state.uncond_text_segment, include_controls=True)
1202
+ block_state.num_noisy_vision_tokens = block_state.cond_full_static["num_noisy_vision_tokens"]
1203
+
1204
+ self.set_block_state(state, block_state)
1205
+ return components, state
1206
+
1207
+
1208
+ class Cosmos3TransferSetTimestepsStep(ModularPipelineBlocks):
1209
+ model_name = "cosmos3-omni"
1210
+
1211
+ @property
1212
+ def description(self) -> str:
1213
+ return (
1214
+ "Resets the scheduler and computes timesteps for a single transfer chunk. UniPCMultistepScheduler keeps "
1215
+ "per-step state on the instance, so it is reset per chunk (each autoregressive chunk is a full denoise)."
1216
+ )
1217
+
1218
+ @property
1219
+ def expected_components(self) -> list[ComponentSpec]:
1220
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
1221
+
1222
+ @property
1223
+ def inputs(self) -> list[InputParam]:
1224
+ return [InputParam.template("num_inference_steps", required=True)]
1225
+
1226
+ @property
1227
+ def intermediate_outputs(self) -> list[OutputParam]:
1228
+ return [
1229
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for this chunk."),
1230
+ OutputParam(
1231
+ "num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps for this chunk."
1232
+ ),
1233
+ ]
1234
+
1235
+ @torch.no_grad()
1236
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1237
+ block_state = self.get_block_state(state)
1238
+ device = components._execution_device
1239
+ components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
1240
+ block_state.timesteps = components.scheduler.timesteps
1241
+ block_state.num_warmup_steps = (
1242
+ len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
1243
+ )
1244
+ self.set_block_state(state, block_state)
1245
+ return components, state
1246
+
1247
+
1248
+ class Cosmos3DistilledSetTimestepsStep(ModularPipelineBlocks):
1249
+ model_name = "cosmos3-omni"
1250
+
1251
+ @property
1252
+ def description(self) -> str:
1253
+ return "Initializes the fixed distilled sampling schedule from the pipeline's `distilled_sigmas` config."
1254
+
1255
+ @property
1256
+ def expected_components(self) -> list[ComponentSpec]:
1257
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
1258
+
1259
+ @property
1260
+ def expected_configs(self) -> list[ConfigSpec]:
1261
+ return [
1262
+ ConfigSpec(name="is_distilled", default=True),
1263
+ ConfigSpec(name="distilled_sigmas", default=None),
1264
+ ]
1265
+
1266
+ @property
1267
+ def inputs(self) -> list[InputParam]:
1268
+ return [
1269
+ InputParam.template("num_inference_steps", required=False, default=None),
1270
+ InputParam(
1271
+ name="guidance_scale",
1272
+ type_hint=float,
1273
+ default=None,
1274
+ description=(
1275
+ "Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the "
1276
+ "scale is forced to 1.0. Passing a value other than 1.0 raises an error."
1277
+ ),
1278
+ ),
1279
+ ]
1280
+
1281
+ @property
1282
+ def intermediate_outputs(self) -> list[OutputParam]:
1283
+ return [
1284
+ OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."),
1285
+ OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."),
1286
+ OutputParam(
1287
+ "num_inference_steps",
1288
+ type_hint=int,
1289
+ description="Resolved number of denoising steps (fixed by the distilled schedule).",
1290
+ ),
1291
+ OutputParam(
1292
+ name="guidance_scale",
1293
+ type_hint=float,
1294
+ description="Resolved classifier-free guidance scale (always 1.0 for distilled checkpoints).",
1295
+ ),
1296
+ ]
1297
+
1298
+ @torch.no_grad()
1299
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1300
+ block_state = self.get_block_state(state)
1301
+ device = components._execution_device
1302
+
1303
+ sigmas = components.config.distilled_sigmas
1304
+ if not sigmas:
1305
+ raise ValueError(
1306
+ "Cosmos3DistilledSetTimestepsStep requires the pipeline config `distilled_sigmas` to be set "
1307
+ "(populated from the distilled checkpoint's `modular_model_index.json`). Load a distilled Cosmos3 "
1308
+ "checkpoint or use `Cosmos3OmniModularPipeline` for base checkpoints."
1309
+ )
1310
+ sigmas = [float(s) for s in sigmas]
1311
+ distilled_steps = len(sigmas)
1312
+
1313
+ if block_state.num_inference_steps is not None and block_state.num_inference_steps != distilled_steps:
1314
+ raise ValueError(
1315
+ "This is a distilled checkpoint; the step count is fixed by the pipeline's "
1316
+ f"`distilled_sigmas` config ({distilled_steps} steps). "
1317
+ f"`num_inference_steps` must be {distilled_steps} or left unset (got {block_state.num_inference_steps})."
1318
+ )
1319
+ if block_state.guidance_scale is not None and block_state.guidance_scale != 1.0:
1320
+ raise ValueError(
1321
+ "This is a distilled checkpoint; classifier-free guidance is baked into the weights. "
1322
+ f"`guidance_scale` must be 1.0 or left unset (got {block_state.guidance_scale})."
1323
+ )
1324
+
1325
+ components.scheduler.set_timesteps(sigmas=sigmas, device=device)
1326
+ block_state.num_inference_steps = distilled_steps
1327
+ block_state.guidance_scale = 1.0
1328
+ block_state.timesteps = components.scheduler.timesteps
1329
+ block_state.num_warmup_steps = len(block_state.timesteps) - distilled_steps * components.scheduler.order
1330
+ self.set_block_state(state, block_state)
1331
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+
5
+ from ...configuration_utils import FrozenDict
6
+ from ...video_processor import VideoProcessor
7
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
8
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
9
+ from .modular_pipeline import Cosmos3OmniModularPipeline
10
+
11
+
12
+ class Cosmos3TransferSetupStep(ModularPipelineBlocks):
13
+ model_name = "cosmos3-omni"
14
+
15
+ @property
16
+ def description(self) -> str:
17
+ return (
18
+ "Preprocesses the transfer control videos and resolves the autoregressive chunk geometry "
19
+ "(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop."
20
+ )
21
+
22
+ @property
23
+ def expected_components(self) -> list[ComponentSpec]:
24
+ return [
25
+ ComponentSpec(
26
+ "video_processor",
27
+ VideoProcessor,
28
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
29
+ default_creation_method="from_config",
30
+ ),
31
+ ]
32
+
33
+ @property
34
+ def inputs(self) -> list[InputParam]:
35
+ return [
36
+ InputParam(
37
+ name="control_videos",
38
+ type_hint=dict,
39
+ required=True,
40
+ description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.",
41
+ ),
42
+ InputParam(
43
+ name="height", type_hint=int, default=None, description="Height of the generated video in pixels."
44
+ ),
45
+ InputParam(
46
+ name="width", type_hint=int, default=None, description="Width of the generated video in pixels."
47
+ ),
48
+ InputParam(
49
+ name="num_frames",
50
+ type_hint=int,
51
+ default=None,
52
+ description="Optional cap on the number of output frames (defaults to the control video length).",
53
+ ),
54
+ InputParam(
55
+ name="num_video_frames_per_chunk",
56
+ type_hint=int,
57
+ default=None,
58
+ description="Number of pixel frames generated per autoregressive chunk.",
59
+ ),
60
+ InputParam(
61
+ name="num_conditional_frames",
62
+ type_hint=int,
63
+ default=1,
64
+ description="Number of frames each chunk reuses from the previous chunk's tail.",
65
+ ),
66
+ ]
67
+
68
+ @property
69
+ def intermediate_outputs(self) -> list[OutputParam]:
70
+ return [
71
+ OutputParam("height", type_hint=int, description="Resolved output height in pixels."),
72
+ OutputParam("width", type_hint=int, description="Resolved output width in pixels."),
73
+ OutputParam(
74
+ "control_frames",
75
+ type_hint=dict,
76
+ description="Preprocessed, time-padded control maps in canonical hint order.",
77
+ ),
78
+ OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."),
79
+ OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."),
80
+ OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."),
81
+ OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."),
82
+ ]
83
+
84
+ @torch.no_grad()
85
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
86
+ block_state = self.get_block_state(state)
87
+ device = components._execution_device
88
+ dtype = components.transformer.dtype
89
+
90
+ if block_state.height is None:
91
+ block_state.height = 720
92
+ if block_state.width is None:
93
+ block_state.width = 1280
94
+
95
+ # Canonical hint order used both to validate and to order the preprocessed control maps.
96
+ hint_order = ["edge", "blur", "depth", "seg", "wsm"]
97
+ control_videos = block_state.control_videos
98
+ if not isinstance(control_videos, dict) or not control_videos:
99
+ raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.")
100
+ unknown = [k for k in control_videos if k not in hint_order]
101
+ if unknown:
102
+ raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.")
103
+ if any(v is None for v in control_videos.values()):
104
+ raise ValueError("`control_videos` entries must be loaded videos, not None.")
105
+
106
+ tcf = components.vae_scale_factor_temporal
107
+ sf = components.vae_scale_factor_spatial
108
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
109
+ raise ValueError(
110
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
111
+ )
112
+
113
+ # Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order.
114
+ # The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys).
115
+ hint_keys = [k for k in hint_order if k in control_videos]
116
+ control_frames = {
117
+ key: components.video_processor.preprocess_video(
118
+ control_videos[key], height=block_state.height, width=block_state.width
119
+ ).to(device=device, dtype=dtype)
120
+ for key in hint_keys
121
+ }
122
+
123
+ # Output frame count / chunking come from the (first) control video, optionally capped by num_frames.
124
+ total_frames = next(iter(control_frames.values())).shape[2]
125
+ if block_state.num_frames is not None:
126
+ total_frames = min(total_frames, block_state.num_frames)
127
+ total_frames = max(1, total_frames)
128
+
129
+ per_chunk = (
130
+ block_state.num_video_frames_per_chunk
131
+ if block_state.num_video_frames_per_chunk is not None
132
+ else total_frames
133
+ )
134
+ chunk_frames = 1 if total_frames == 1 else per_chunk
135
+ chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1
136
+
137
+ if total_frames <= chunk_frames:
138
+ num_chunks, stride = 1, chunk_frames
139
+ else:
140
+ stride = chunk_frames - block_state.num_conditional_frames
141
+ if stride <= 0:
142
+ raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.")
143
+ remaining = total_frames - chunk_frames
144
+ num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0))
145
+
146
+ # Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to
147
+ # keep reflecting). No truncation here; per-chunk slicing happens later.
148
+ padded = max(total_frames, chunk_frames)
149
+ control_frames_padded = {}
150
+ for key, frames in control_frames.items():
151
+ while frames.shape[2] < padded:
152
+ pad_len = min(frames.shape[2] - 1, padded - frames.shape[2])
153
+ if pad_len <= 0:
154
+ pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1)
155
+ frames = torch.cat([frames, pad_frame], dim=2)
156
+ break
157
+ frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
158
+ control_frames_padded[key] = frames
159
+ block_state.control_frames = control_frames_padded
160
+ block_state.total_frames = total_frames
161
+ block_state.chunk_frames = chunk_frames
162
+ block_state.num_chunks = num_chunks
163
+ block_state.stride = stride
164
+
165
+ self.set_block_state(state, block_state)
166
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py ADDED
@@ -0,0 +1,259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ...configuration_utils import FrozenDict
4
+ from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
5
+ from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
6
+ from ...utils import logging
7
+ from ...video_processor import VideoProcessor
8
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
9
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
10
+ from .modular_pipeline import Cosmos3OmniModularPipeline
11
+
12
+
13
+ logger = logging.get_logger(__name__)
14
+
15
+
16
+ class Cosmos3VideoDecodeStep(ModularPipelineBlocks):
17
+ model_name = "cosmos3-omni"
18
+
19
+ @property
20
+ def description(self) -> str:
21
+ return "Decodes denoised vision latents into video outputs."
22
+
23
+ @property
24
+ def expected_components(self) -> list[ComponentSpec]:
25
+ return [
26
+ ComponentSpec("vae", AutoencoderKLWan),
27
+ ComponentSpec(
28
+ "video_processor",
29
+ VideoProcessor,
30
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
31
+ default_creation_method="from_config",
32
+ ),
33
+ ]
34
+
35
+ @property
36
+ def inputs(self) -> list[InputParam]:
37
+ return [
38
+ InputParam.template("latents", required=True, description="Denoised vision latents to decode."),
39
+ InputParam.template("output_type", default="pil"),
40
+ ]
41
+
42
+ @property
43
+ def intermediate_outputs(self) -> list[OutputParam]:
44
+ return [OutputParam.template("videos")]
45
+
46
+ @torch.no_grad()
47
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
48
+ block_state = self.get_block_state(state)
49
+
50
+ device = components._execution_device
51
+
52
+ if block_state.output_type == "latent":
53
+ block_state.videos = block_state.latents
54
+ else:
55
+ in_dtype = block_state.latents.dtype
56
+ vae_dtype = components.vae.dtype
57
+ mean = components._vae_latents_mean.to(device=block_state.latents.device, dtype=vae_dtype)
58
+ inv_std = components._vae_latents_inv_std.to(device=block_state.latents.device, dtype=vae_dtype)
59
+ z_raw = block_state.latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1)
60
+ decoded = components.vae.decode(z_raw).sample.to(in_dtype)
61
+ block_state.videos = components.video_processor.postprocess_video(
62
+ decoded, output_type=block_state.output_type
63
+ )[0]
64
+
65
+ if components.requires_safety_checker and block_state.output_type != "latent":
66
+ if getattr(components, "safety_checker", None) is None:
67
+ raise ValueError(
68
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
69
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
70
+ )
71
+ block_state.videos = components._apply_video_safety_check(
72
+ block_state.videos, output_type=block_state.output_type, device=device
73
+ )
74
+
75
+ self.set_block_state(state, block_state)
76
+ return components, state
77
+
78
+
79
+ class Cosmos3SoundDecodeStep(ModularPipelineBlocks):
80
+ model_name = "cosmos3-omni"
81
+
82
+ @property
83
+ def description(self) -> str:
84
+ return "Decodes sound latents into waveform output."
85
+
86
+ @property
87
+ def expected_components(self) -> list[ComponentSpec]:
88
+ return [ComponentSpec("sound_tokenizer", Cosmos3AVAEAudioTokenizer)]
89
+
90
+ @property
91
+ def inputs(self) -> list[InputParam]:
92
+ return [
93
+ InputParam(
94
+ name="sound_latents",
95
+ type_hint=torch.Tensor,
96
+ required=True,
97
+ description="Denoised sound latents to decode.",
98
+ )
99
+ ]
100
+
101
+ @property
102
+ def intermediate_outputs(self) -> list[OutputParam]:
103
+ return [
104
+ OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."),
105
+ OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."),
106
+ ]
107
+
108
+ @torch.no_grad()
109
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
110
+ block_state = self.get_block_state(state)
111
+ if components.sound_tokenizer is None:
112
+ raise ValueError("Sound decoding requires a sound-capable checkpoint with a sound_tokenizer.")
113
+ block_state.sound = components.decode_sound(block_state.sound_latents)
114
+ block_state.sampling_rate = int(components.sound_tokenizer.config.sampling_rate)
115
+ self.set_block_state(state, block_state)
116
+ return components, state
117
+
118
+
119
+ class Cosmos3TransferDecodeChunkStep(ModularPipelineBlocks):
120
+ model_name = "cosmos3-omni"
121
+
122
+ @property
123
+ def description(self) -> str:
124
+ return (
125
+ "Decodes one transfer chunk's latents to pixels (float32, clamped to [-1, 1]), records it as the "
126
+ "autoregressive seed for the next chunk, and appends it to output_chunks (dropping the overlap that "
127
+ "later chunks share with the previous chunk's conditioning frames)."
128
+ )
129
+
130
+ @property
131
+ def expected_components(self) -> list[ComponentSpec]:
132
+ return [ComponentSpec("vae", AutoencoderKLWan)]
133
+
134
+ @property
135
+ def inputs(self) -> list[InputParam]:
136
+ return [
137
+ InputParam(
138
+ name="latents",
139
+ type_hint=torch.Tensor,
140
+ required=True,
141
+ description="Denoised target latents for this chunk.",
142
+ ),
143
+ InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."),
144
+ InputParam(
145
+ name="current_conditional_frames",
146
+ type_hint=int,
147
+ required=True,
148
+ description="Number of pixel frames this chunk reused from the previous chunk.",
149
+ ),
150
+ InputParam(
151
+ name="output_chunks",
152
+ type_hint=list[torch.Tensor],
153
+ required=True,
154
+ description="Decoded pixel chunks accumulated so far.",
155
+ ),
156
+ ]
157
+
158
+ @property
159
+ def intermediate_outputs(self) -> list[OutputParam]:
160
+ return [
161
+ OutputParam(
162
+ "previous_output",
163
+ type_hint=torch.Tensor,
164
+ description="Decoded pixels of this chunk, used to seed the next chunk.",
165
+ ),
166
+ OutputParam(
167
+ "output_chunks",
168
+ type_hint=list[torch.Tensor],
169
+ description="Decoded pixel chunks accumulated so far (with this chunk appended).",
170
+ ),
171
+ ]
172
+
173
+ @torch.no_grad()
174
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
175
+ block_state = self.get_block_state(state)
176
+ latents = block_state.latents
177
+ vae_dtype = components.vae.dtype
178
+ mean = components._vae_latents_mean.to(device=latents.device, dtype=vae_dtype)
179
+ inv_std = components._vae_latents_inv_std.to(device=latents.device, dtype=vae_dtype)
180
+ z_raw = latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1)
181
+ output_video = components.vae.decode(z_raw).sample.to(torch.float32).clamp(-1, 1)
182
+ block_state.previous_output = output_video
183
+ chunk = (
184
+ output_video if block_state.chunk_id == 0 else output_video[:, :, block_state.current_conditional_frames :]
185
+ )
186
+ block_state.output_chunks = [*block_state.output_chunks, chunk]
187
+ self.set_block_state(state, block_state)
188
+ return components, state
189
+
190
+
191
+ class Cosmos3TransferStitchStep(ModularPipelineBlocks):
192
+ model_name = "cosmos3-omni"
193
+
194
+ @property
195
+ def description(self) -> str:
196
+ return (
197
+ "Concatenates the decoded transfer chunks along time, truncates to total_frames, and post-processes to "
198
+ "the requested output type. Transfer produces no audio, so sound / sampling_rate are None."
199
+ )
200
+
201
+ @property
202
+ def expected_components(self) -> list[ComponentSpec]:
203
+ return [
204
+ ComponentSpec("vae", AutoencoderKLWan),
205
+ ComponentSpec(
206
+ "video_processor",
207
+ VideoProcessor,
208
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
209
+ default_creation_method="from_config",
210
+ ),
211
+ ]
212
+
213
+ @property
214
+ def inputs(self) -> list[InputParam]:
215
+ return [
216
+ InputParam(
217
+ name="output_chunks",
218
+ type_hint=list[torch.Tensor],
219
+ required=True,
220
+ description="Decoded pixel chunks to stitch together.",
221
+ ),
222
+ InputParam(
223
+ name="total_frames", type_hint=int, required=True, description="Total number of output frames to keep."
224
+ ),
225
+ InputParam.template("output_type", default="pil"),
226
+ ]
227
+
228
+ @property
229
+ def intermediate_outputs(self) -> list[OutputParam]:
230
+ return [
231
+ OutputParam("videos", description="The generated transfer video."),
232
+ OutputParam("sound", description="Always None for transfer (no audio)."),
233
+ OutputParam("sampling_rate", description="Always None for transfer (no audio)."),
234
+ ]
235
+
236
+ @torch.no_grad()
237
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
238
+ block_state = self.get_block_state(state)
239
+ device = components._execution_device
240
+
241
+ decoded = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames]
242
+ block_state.videos = components.video_processor.postprocess_video(
243
+ decoded, output_type=block_state.output_type
244
+ )[0]
245
+
246
+ if components.requires_safety_checker and block_state.output_type != "latent":
247
+ if getattr(components, "safety_checker", None) is None:
248
+ raise ValueError(
249
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
250
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
251
+ )
252
+ block_state.videos = components._apply_video_safety_check(
253
+ block_state.videos, output_type=block_state.output_type, device=device
254
+ )
255
+
256
+ block_state.sound = None
257
+ block_state.sampling_rate = None
258
+ self.set_block_state(state, block_state)
259
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py ADDED
@@ -0,0 +1,889 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import inspect
2
+
3
+ import torch
4
+
5
+ from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
6
+ from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
7
+ from ..modular_pipeline import (
8
+ BlockState,
9
+ LoopSequentialPipelineBlocks,
10
+ ModularPipelineBlocks,
11
+ PipelineState,
12
+ )
13
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
14
+ from .modular_pipeline import Cosmos3OmniModularPipeline
15
+
16
+
17
+ class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks):
18
+ model_name = "cosmos3-omni"
19
+
20
+ @property
21
+ def description(self) -> str:
22
+ return "Prepares vision tokens and timesteps for one denoising iteration."
23
+
24
+ @property
25
+ def expected_components(self) -> list[ComponentSpec]:
26
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
27
+
28
+ @property
29
+ def inputs(self) -> list[InputParam]:
30
+ return [
31
+ InputParam.template("latents", required=True, description="Noisy vision latents to denoise."),
32
+ InputParam(
33
+ name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment."
34
+ ),
35
+ ]
36
+
37
+ @property
38
+ def intermediate_outputs(self) -> list[OutputParam]:
39
+ return [
40
+ OutputParam(
41
+ "vision_tokens",
42
+ type_hint=list[torch.Tensor],
43
+ description="Vision tokens for the transformer denoiser.",
44
+ ),
45
+ OutputParam("vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the vision tokens."),
46
+ ]
47
+
48
+ @torch.no_grad()
49
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
50
+ device = components._execution_device
51
+ block_state.vision_tokens = [block_state.latents.to(device=device, dtype=components.transformer.dtype)]
52
+ block_state.vision_timesteps = torch.full(
53
+ (block_state.cond_vision_segment["num_noisy_vision_tokens"],), t.item(), device=device
54
+ )
55
+ return components, block_state
56
+
57
+
58
+ class Cosmos3SoundLoopPrepareStep(ModularPipelineBlocks):
59
+ model_name = "cosmos3-omni"
60
+
61
+ @property
62
+ def description(self) -> str:
63
+ return "Prepares sound tokens and timesteps for one denoising iteration."
64
+
65
+ @property
66
+ def expected_components(self) -> list[ComponentSpec]:
67
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
68
+
69
+ @property
70
+ def inputs(self) -> list[InputParam]:
71
+ return [
72
+ InputParam(
73
+ name="sound_latents",
74
+ type_hint=torch.Tensor,
75
+ required=True,
76
+ description="Noisy sound latents to denoise.",
77
+ ),
78
+ InputParam(
79
+ name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment."
80
+ ),
81
+ ]
82
+
83
+ @property
84
+ def intermediate_outputs(self) -> list[OutputParam]:
85
+ return [
86
+ OutputParam(
87
+ "sound_tokens", type_hint=list[torch.Tensor], description="Sound tokens for the transformer denoiser."
88
+ ),
89
+ OutputParam("sound_timesteps", type_hint=torch.Tensor, description="Timesteps for the sound tokens."),
90
+ ]
91
+
92
+ @torch.no_grad()
93
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
94
+ device = components._execution_device
95
+ block_state.sound_tokens = [block_state.sound_latents.to(device=device, dtype=components.transformer.dtype)]
96
+ block_state.sound_timesteps = torch.full(
97
+ (block_state.cond_sound_segment["sound_len"],), t.item(), device=device
98
+ )
99
+ return components, block_state
100
+
101
+
102
+ class Cosmos3ActionLoopPrepareStep(ModularPipelineBlocks):
103
+ model_name = "cosmos3-omni"
104
+
105
+ @property
106
+ def description(self) -> str:
107
+ return "Prepares action tokens and timesteps for one denoising iteration."
108
+
109
+ @property
110
+ def expected_components(self) -> list[ComponentSpec]:
111
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
112
+
113
+ @property
114
+ def inputs(self) -> list[InputParam]:
115
+ return [
116
+ InputParam(
117
+ name="action_latents",
118
+ type_hint=torch.Tensor,
119
+ required=True,
120
+ description="Noisy action latents to denoise.",
121
+ ),
122
+ InputParam(
123
+ name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment."
124
+ ),
125
+ ]
126
+
127
+ @property
128
+ def intermediate_outputs(self) -> list[OutputParam]:
129
+ return [
130
+ OutputParam(
131
+ "action_tokens",
132
+ type_hint=list[torch.Tensor],
133
+ description="Action tokens for the transformer denoiser.",
134
+ ),
135
+ OutputParam("action_timesteps", type_hint=torch.Tensor, description="Timesteps for the action tokens."),
136
+ ]
137
+
138
+ @torch.no_grad()
139
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
140
+ device = components._execution_device
141
+ block_state.action_tokens = [block_state.action_latents.to(device=device, dtype=components.transformer.dtype)]
142
+ block_state.action_timesteps = torch.full(
143
+ (block_state.cond_action_segment["num_noisy_action_tokens"],), t.item(), device=device
144
+ )
145
+ return components, block_state
146
+
147
+
148
+ class Cosmos3LoopDenoiser(ModularPipelineBlocks):
149
+ model_name = "cosmos3-omni"
150
+
151
+ @property
152
+ def description(self) -> str:
153
+ return "Predicts available Cosmos3 modality velocities for one denoising iteration."
154
+
155
+ @property
156
+ def expected_components(self) -> list[ComponentSpec]:
157
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
158
+
159
+ @property
160
+ def inputs(self) -> list[InputParam]:
161
+ return [
162
+ InputParam.template("denoiser_input_fields"),
163
+ InputParam(
164
+ name="guidance_scale",
165
+ type_hint=float,
166
+ default=6.0,
167
+ description="Scale for classifier-free guidance.",
168
+ ),
169
+ ]
170
+
171
+ @property
172
+ def intermediate_outputs(self) -> list[OutputParam]:
173
+ return [
174
+ OutputParam(
175
+ "velocity_vision", type_hint=torch.Tensor, description="Predicted velocity for vision latents."
176
+ ),
177
+ OutputParam("velocity_sound", type_hint=torch.Tensor, description="Predicted velocity for sound latents."),
178
+ OutputParam(
179
+ "velocity_action", type_hint=torch.Tensor, description="Predicted velocity for action latents."
180
+ ),
181
+ ]
182
+
183
+ @torch.no_grad()
184
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
185
+ denoiser_input_fields = block_state.denoiser_input_fields
186
+ loop_input_fields = block_state.as_dict()
187
+ has_sound = "sound_tokens" in loop_input_fields
188
+ has_action = "action_tokens" in loop_input_fields
189
+ do_cfg = block_state.guidance_scale != 1.0
190
+ transformer_args = set(inspect.signature(components.transformer.forward).parameters)
191
+
192
+ prediction_passes = ["cond"]
193
+ if do_cfg:
194
+ prediction_passes.append("uncond")
195
+
196
+ velocities = {}
197
+ for pass_name in prediction_passes:
198
+ transformer_kwargs = {}
199
+ for field_name, field_value in denoiser_input_fields.items():
200
+ if field_name.startswith(f"{pass_name}_"):
201
+ transformer_field_name = field_name.removeprefix(f"{pass_name}_")
202
+ if transformer_field_name.endswith("_segment"):
203
+ transformer_kwargs.update(field_value)
204
+ else:
205
+ transformer_kwargs[transformer_field_name] = field_value
206
+ elif field_name in transformer_args:
207
+ transformer_kwargs[field_name] = field_value
208
+ transformer_kwargs.update(
209
+ {
210
+ field_name: field_value
211
+ for field_name, field_value in loop_input_fields.items()
212
+ if field_name in transformer_args
213
+ }
214
+ )
215
+ transformer_kwargs = {
216
+ name: value for name, value in transformer_kwargs.items() if name in transformer_args
217
+ }
218
+ preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs, return_dict=False)
219
+ velocities[pass_name] = components._mask_velocity_predictions(
220
+ preds_vision,
221
+ preds_sound,
222
+ vision_condition_mask=[loop_input_fields["vision_condition_mask"]],
223
+ sound_condition_mask=[loop_input_fields["sound_condition_mask"]] if has_sound else None,
224
+ preds_action=preds_action,
225
+ action_condition_mask=[loop_input_fields["action_condition_mask"]] if has_action else None,
226
+ raw_action_dim=loop_input_fields.get("raw_action_dim_resolved"),
227
+ )
228
+
229
+ cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"]
230
+ if do_cfg:
231
+ uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"]
232
+ block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * (
233
+ cond_velocity_vision - uncond_velocity_vision
234
+ )
235
+ block_state.velocity_sound = (
236
+ uncond_velocity_sound + block_state.guidance_scale * (cond_velocity_sound - uncond_velocity_sound)
237
+ if has_sound
238
+ else None
239
+ )
240
+ block_state.velocity_action = (
241
+ uncond_velocity_action + block_state.guidance_scale * (cond_velocity_action - uncond_velocity_action)
242
+ if has_action
243
+ else None
244
+ )
245
+ else:
246
+ block_state.velocity_vision = cond_velocity_vision
247
+ block_state.velocity_sound = cond_velocity_sound if has_sound else None
248
+ block_state.velocity_action = cond_velocity_action if has_action else None
249
+
250
+ return components, block_state
251
+
252
+
253
+ class Cosmos3VisionLoopSchedulerStep(ModularPipelineBlocks):
254
+ model_name = "cosmos3-omni"
255
+
256
+ @property
257
+ def description(self) -> str:
258
+ return "Updates vision latents after one denoising iteration."
259
+
260
+ @property
261
+ def expected_components(self) -> list[ComponentSpec]:
262
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
263
+
264
+ @property
265
+ def inputs(self) -> list[InputParam]:
266
+ return [
267
+ InputParam.template("latents", required=True, description="Noisy vision latents to update."),
268
+ InputParam(
269
+ name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity."
270
+ ),
271
+ ]
272
+
273
+ @property
274
+ def intermediate_outputs(self) -> list[OutputParam]:
275
+ return [OutputParam.template("latents")]
276
+
277
+ @torch.no_grad()
278
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
279
+ block_state.latents = components.scheduler.step(
280
+ block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False
281
+ )[0].squeeze(0)
282
+ return components, block_state
283
+
284
+
285
+ class Cosmos3DistilledVisionLoopSchedulerStep(ModularPipelineBlocks):
286
+ model_name = "cosmos3-omni"
287
+
288
+ @property
289
+ def description(self) -> str:
290
+ return "Updates vision latents after one distilled denoising iteration, re-anchoring conditioned frames."
291
+
292
+ @property
293
+ def expected_components(self) -> list[ComponentSpec]:
294
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
295
+
296
+ @property
297
+ def inputs(self) -> list[InputParam]:
298
+ return [
299
+ InputParam.template("latents", required=True, description="Noisy vision latents to update."),
300
+ InputParam(
301
+ name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity."
302
+ ),
303
+ InputParam(
304
+ name="vision_condition_mask",
305
+ type_hint=torch.Tensor,
306
+ required=True,
307
+ description="Mask marking conditioned vision latent frames.",
308
+ ),
309
+ InputParam(
310
+ name="vision_conditioning_latents",
311
+ type_hint=torch.Tensor,
312
+ default=None,
313
+ description="Clean encoded vision latents for re-anchoring conditioned frames.",
314
+ ),
315
+ InputParam(
316
+ name="vision_condition_indexes_for_pack",
317
+ type_hint=list,
318
+ default=None,
319
+ description="Indexes of conditioned vision latent frames; non-empty for image-to-video.",
320
+ ),
321
+ InputParam.template("generator"),
322
+ ]
323
+
324
+ @property
325
+ def intermediate_outputs(self) -> list[OutputParam]:
326
+ return [OutputParam.template("latents")]
327
+
328
+ @torch.no_grad()
329
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
330
+ # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible.
331
+ block_state.latents = components.scheduler.step(
332
+ block_state.velocity_vision.unsqueeze(0),
333
+ t,
334
+ block_state.latents.unsqueeze(0),
335
+ generator=block_state.generator,
336
+ return_dict=False,
337
+ )[0].squeeze(0)
338
+
339
+ # Distilled checkpoints use stochastic (SDE) scheduler steps that re-noise every position.
340
+ # Re-anchor conditioned frames to the clean encoded reference after each step.
341
+ has_image_condition = bool(block_state.vision_condition_indexes_for_pack)
342
+ if has_image_condition and block_state.vision_conditioning_latents is not None:
343
+ mask = block_state.vision_condition_mask
344
+ reference = block_state.vision_conditioning_latents.to(block_state.latents.dtype)
345
+ block_state.latents = mask * reference + (1.0 - mask) * block_state.latents
346
+
347
+ return components, block_state
348
+
349
+
350
+ class Cosmos3SoundLoopSchedulerStep(ModularPipelineBlocks):
351
+ model_name = "cosmos3-omni"
352
+
353
+ @property
354
+ def description(self) -> str:
355
+ return "Updates sound latents after one denoising iteration."
356
+
357
+ @property
358
+ def inputs(self) -> list[InputParam]:
359
+ return [
360
+ InputParam(
361
+ name="sound_latents",
362
+ type_hint=torch.Tensor,
363
+ required=True,
364
+ description="Noisy sound latents to update.",
365
+ ),
366
+ InputParam(
367
+ name="sound_scheduler",
368
+ type_hint=UniPCMultistepScheduler,
369
+ required=True,
370
+ description="Scheduler used to update sound latents.",
371
+ ),
372
+ InputParam(
373
+ name="velocity_sound", type_hint=torch.Tensor, required=True, description="Predicted sound velocity."
374
+ ),
375
+ ]
376
+
377
+ @property
378
+ def intermediate_outputs(self) -> list[OutputParam]:
379
+ return [OutputParam("sound_latents", type_hint=torch.Tensor, description="Updated sound latents.")]
380
+
381
+ @torch.no_grad()
382
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
383
+ block_state.sound_latents = block_state.sound_scheduler.step(
384
+ block_state.velocity_sound.unsqueeze(0), t, block_state.sound_latents.unsqueeze(0), return_dict=False
385
+ )[0].squeeze(0)
386
+ return components, block_state
387
+
388
+
389
+ class Cosmos3ActionLoopSchedulerStep(ModularPipelineBlocks):
390
+ model_name = "cosmos3-omni"
391
+
392
+ @property
393
+ def description(self) -> str:
394
+ return "Updates action latents after one denoising iteration."
395
+
396
+ @property
397
+ def inputs(self) -> list[InputParam]:
398
+ return [
399
+ InputParam(
400
+ name="action_latents",
401
+ type_hint=torch.Tensor,
402
+ required=True,
403
+ description="Noisy action latents to update.",
404
+ ),
405
+ InputParam(
406
+ name="action_scheduler",
407
+ type_hint=UniPCMultistepScheduler,
408
+ required=True,
409
+ description="Scheduler used to update action latents.",
410
+ ),
411
+ InputParam(
412
+ name="velocity_action", type_hint=torch.Tensor, required=True, description="Predicted action velocity."
413
+ ),
414
+ InputParam(
415
+ name="action_condition_mask",
416
+ type_hint=torch.Tensor,
417
+ required=True,
418
+ description="Mask marking conditioned action latent frames.",
419
+ ),
420
+ InputParam(
421
+ name="raw_action_dim_resolved",
422
+ type_hint=int,
423
+ default=None,
424
+ description="Unpadded action-vector dimension.",
425
+ ),
426
+ ]
427
+
428
+ @property
429
+ def intermediate_outputs(self) -> list[OutputParam]:
430
+ return [OutputParam("action_latents", type_hint=torch.Tensor, description="Updated action latents.")]
431
+
432
+ @torch.no_grad()
433
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
434
+ has_noisy_action = block_state.action_condition_mask.sum() < block_state.action_condition_mask.numel()
435
+ if has_noisy_action:
436
+ block_state.action_latents = block_state.action_scheduler.step(
437
+ block_state.velocity_action.unsqueeze(0), t, block_state.action_latents.unsqueeze(0), return_dict=False
438
+ )[0].squeeze(0)
439
+ if block_state.raw_action_dim_resolved is not None:
440
+ block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0
441
+ return components, block_state
442
+
443
+
444
+ class Cosmos3DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
445
+ model_name = "cosmos3-omni"
446
+
447
+ @property
448
+ def description(self) -> str:
449
+ return "Iteratively denoises Cosmos3 latents over scheduler timesteps."
450
+
451
+ @property
452
+ def loop_expected_components(self) -> list[ComponentSpec]:
453
+ return [
454
+ ComponentSpec("scheduler", UniPCMultistepScheduler),
455
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
456
+ ]
457
+
458
+ @property
459
+ def loop_inputs(self) -> list[InputParam]:
460
+ return [
461
+ InputParam.template("timesteps", required=True),
462
+ InputParam.template("num_inference_steps", required=True),
463
+ InputParam(
464
+ name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps."
465
+ ),
466
+ ]
467
+
468
+ @torch.no_grad()
469
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
470
+ block_state = self.get_block_state(state)
471
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
472
+ for i, t in enumerate(block_state.timesteps):
473
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
474
+ if i == len(block_state.timesteps) - 1 or (
475
+ (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
476
+ ):
477
+ progress_bar.update()
478
+ self.set_block_state(state, block_state)
479
+ return components, state
480
+
481
+
482
+ class Cosmos3VisionDenoiseStep(Cosmos3DenoiseLoopWrapper):
483
+ block_classes = [
484
+ Cosmos3VisionLoopPrepareStep,
485
+ Cosmos3LoopDenoiser,
486
+ Cosmos3VisionLoopSchedulerStep,
487
+ ]
488
+ block_names = ["prepare_vision", "denoiser", "update_vision"]
489
+
490
+ @property
491
+ def description(self) -> str:
492
+ return "Runs the vision-only Cosmos3 denoising loop."
493
+
494
+
495
+ class Cosmos3DistilledVisionDenoiseStep(Cosmos3DenoiseLoopWrapper):
496
+ model_name = "cosmos3-omni"
497
+ block_classes = [
498
+ Cosmos3VisionLoopPrepareStep,
499
+ Cosmos3LoopDenoiser,
500
+ Cosmos3DistilledVisionLoopSchedulerStep,
501
+ ]
502
+ block_names = ["prepare_vision", "denoiser", "update_vision"]
503
+
504
+ @property
505
+ def description(self) -> str:
506
+ return "Runs the vision-only distilled Cosmos3 denoising loop."
507
+
508
+ @property
509
+ def loop_expected_components(self) -> list[ComponentSpec]:
510
+ return [
511
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
512
+ ComponentSpec("transformer", Cosmos3OmniTransformer),
513
+ ]
514
+
515
+
516
+ class Cosmos3VisionSoundDenoiseStep(Cosmos3DenoiseLoopWrapper):
517
+ block_classes = [
518
+ Cosmos3VisionLoopPrepareStep,
519
+ Cosmos3SoundLoopPrepareStep,
520
+ Cosmos3LoopDenoiser,
521
+ Cosmos3VisionLoopSchedulerStep,
522
+ Cosmos3SoundLoopSchedulerStep,
523
+ ]
524
+ block_names = ["prepare_vision", "prepare_sound", "denoiser", "update_vision", "update_sound"]
525
+
526
+ @property
527
+ def description(self) -> str:
528
+ return "Runs the vision-and-sound Cosmos3 denoising loop."
529
+
530
+
531
+ class Cosmos3VisionActionDenoiseStep(Cosmos3DenoiseLoopWrapper):
532
+ block_classes = [
533
+ Cosmos3VisionLoopPrepareStep,
534
+ Cosmos3ActionLoopPrepareStep,
535
+ Cosmos3LoopDenoiser,
536
+ Cosmos3VisionLoopSchedulerStep,
537
+ Cosmos3ActionLoopSchedulerStep,
538
+ ]
539
+ block_names = ["prepare_vision", "prepare_action", "denoiser", "update_vision", "update_action"]
540
+
541
+ @property
542
+ def description(self) -> str:
543
+ return "Runs the vision-and-action Cosmos3 denoising loop."
544
+
545
+
546
+ class Cosmos3VisionSoundActionDenoiseStep(Cosmos3DenoiseLoopWrapper):
547
+ block_classes = [
548
+ Cosmos3VisionLoopPrepareStep,
549
+ Cosmos3SoundLoopPrepareStep,
550
+ Cosmos3ActionLoopPrepareStep,
551
+ Cosmos3LoopDenoiser,
552
+ Cosmos3VisionLoopSchedulerStep,
553
+ Cosmos3SoundLoopSchedulerStep,
554
+ Cosmos3ActionLoopSchedulerStep,
555
+ ]
556
+ block_names = [
557
+ "prepare_vision",
558
+ "prepare_sound",
559
+ "prepare_action",
560
+ "denoiser",
561
+ "update_vision",
562
+ "update_sound",
563
+ "update_action",
564
+ ]
565
+
566
+ @property
567
+ def description(self) -> str:
568
+ return "Runs the vision, sound, and action Cosmos3 denoising loop."
569
+
570
+
571
+ class Cosmos3TransferLoopPrepareStep(ModularPipelineBlocks):
572
+ model_name = "cosmos3-omni"
573
+
574
+ @property
575
+ def description(self) -> str:
576
+ return "Prepares the full [control..., target] and target-only vision token lists plus timesteps for one transfer iteration."
577
+
578
+ @property
579
+ def expected_components(self) -> list[ComponentSpec]:
580
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
581
+
582
+ @property
583
+ def inputs(self) -> list[InputParam]:
584
+ return [
585
+ InputParam(
586
+ name="control_latents",
587
+ type_hint=list[torch.Tensor],
588
+ required=True,
589
+ description="Clean control latents for this chunk, one per hint in canonical order.",
590
+ ),
591
+ InputParam(
592
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to denoise."
593
+ ),
594
+ InputParam(
595
+ name="num_noisy_vision_tokens",
596
+ type_hint=int,
597
+ required=True,
598
+ description="Number of noisy target vision tokens denoised each step.",
599
+ ),
600
+ ]
601
+
602
+ @property
603
+ def intermediate_outputs(self) -> list[OutputParam]:
604
+ return [
605
+ OutputParam(
606
+ "vision_tokens_full",
607
+ type_hint=list[torch.Tensor],
608
+ description="Token list for the [control..., target] forward passes.",
609
+ ),
610
+ OutputParam(
611
+ "vision_tokens_target",
612
+ type_hint=list[torch.Tensor],
613
+ description="Token list for the target-only (no-control) forward pass.",
614
+ ),
615
+ OutputParam(
616
+ "vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the noisy target tokens."
617
+ ),
618
+ ]
619
+
620
+ @torch.no_grad()
621
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
622
+ device = components._execution_device
623
+ dtype = components.transformer.dtype
624
+ block_state.vision_tokens_full = [c.to(device=device, dtype=dtype) for c in block_state.control_latents] + [
625
+ block_state.latents.to(device=device, dtype=dtype)
626
+ ]
627
+ block_state.vision_tokens_target = [block_state.latents.to(device=device, dtype=dtype)]
628
+ block_state.vision_timesteps = torch.full((block_state.num_noisy_vision_tokens,), t.item(), device=device)
629
+ return components, block_state
630
+
631
+
632
+ class Cosmos3TransferLoopDenoiser(ModularPipelineBlocks):
633
+ # Dedicated (not Cosmos3LoopDenoiser): transfer runs up to 3 passes over different token sequences with nested
634
+ # control/text CFG and interval gating, which the generic cond/uncond denoiser cannot express.
635
+ model_name = "cosmos3-omni"
636
+
637
+ @property
638
+ def description(self) -> str:
639
+ return (
640
+ "Predicts the transfer velocity with nested control/text CFG over [control..., target]. Each branch is "
641
+ "gated by its guidance interval, and the result is masked so conditioned frames get zero velocity."
642
+ )
643
+
644
+ @property
645
+ def expected_components(self) -> list[ComponentSpec]:
646
+ return [ComponentSpec("transformer", Cosmos3OmniTransformer)]
647
+
648
+ @property
649
+ def inputs(self) -> list[InputParam]:
650
+ return [
651
+ # The three pre-packed CFG sequence variants (cond_full / cond_no_control / uncond_full) flow in as
652
+ # denoiser_input_fields, gathered generically like the other Cosmos3 denoisers.
653
+ InputParam.template("denoiser_input_fields"),
654
+ InputParam(
655
+ name="vision_tokens_full",
656
+ type_hint=list[torch.Tensor],
657
+ required=True,
658
+ description="Token list for the [control..., target] forward passes.",
659
+ ),
660
+ InputParam(
661
+ name="vision_tokens_target",
662
+ type_hint=list[torch.Tensor],
663
+ required=True,
664
+ description="Token list for the target-only (no-control) forward pass.",
665
+ ),
666
+ InputParam(
667
+ name="vision_timesteps",
668
+ type_hint=torch.Tensor,
669
+ required=True,
670
+ description="Timesteps for the noisy target tokens.",
671
+ ),
672
+ InputParam(
673
+ name="velocity_mask",
674
+ type_hint=torch.Tensor,
675
+ required=True,
676
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
677
+ ),
678
+ InputParam(
679
+ name="guidance_scale",
680
+ type_hint=float,
681
+ default=6.0,
682
+ description="Scale for text classifier-free guidance.",
683
+ ),
684
+ InputParam(
685
+ name="control_guidance",
686
+ type_hint=float,
687
+ default=1.0,
688
+ description="Scale for the control (structural) guidance axis.",
689
+ ),
690
+ InputParam(
691
+ name="guidance_interval",
692
+ type_hint=tuple,
693
+ default=None,
694
+ description="Timestep interval [lo, hi] over which text guidance is active (None = always).",
695
+ ),
696
+ InputParam(
697
+ name="control_guidance_interval",
698
+ type_hint=tuple,
699
+ default=None,
700
+ description="Timestep interval [lo, hi] over which control guidance is active (None = always).",
701
+ ),
702
+ ]
703
+
704
+ @property
705
+ def intermediate_outputs(self) -> list[OutputParam]:
706
+ return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")]
707
+
708
+ @staticmethod
709
+ def _forward(components, static, vision_tokens, vision_timesteps):
710
+ preds_vision, _, _ = components.transformer(
711
+ input_ids=static["input_ids"],
712
+ text_indexes=static["text_indexes"],
713
+ position_ids=static["position_ids"],
714
+ und_len=static["und_len"],
715
+ sequence_length=static["sequence_length"],
716
+ vision_tokens=vision_tokens,
717
+ vision_token_shapes=static["vision_token_shapes"],
718
+ vision_sequence_indexes=static["vision_sequence_indexes"],
719
+ vision_mse_loss_indexes=static["vision_mse_loss_indexes"],
720
+ vision_timesteps=vision_timesteps,
721
+ vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"],
722
+ return_dict=False,
723
+ )
724
+ return preds_vision[-1]
725
+
726
+ @torch.no_grad()
727
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
728
+ # active-at: a None interval is always active; otherwise the timestep must fall within [lo, hi].
729
+ guidance_interval = block_state.guidance_interval
730
+ guidance_active = guidance_interval is None or (
731
+ float(guidance_interval[0]) <= float(t.item()) <= float(guidance_interval[1])
732
+ )
733
+ control_interval = block_state.control_guidance_interval
734
+ control_active = control_interval is None or (
735
+ float(control_interval[0]) <= float(t.item()) <= float(control_interval[1])
736
+ )
737
+ step_guidance = block_state.guidance_scale if guidance_active else 1.0
738
+ step_control = block_state.control_guidance if control_active else 1.0
739
+ needs_text_cfg = step_guidance > 1.0
740
+ needs_control_cfg = step_control != 1.0
741
+
742
+ denoiser_input_fields = block_state.denoiser_input_fields
743
+ cond_full_static = denoiser_input_fields["cond_full_static"]
744
+ cond_no_control_static = denoiser_input_fields["cond_no_control_static"]
745
+ uncond_full_static = denoiser_input_fields["uncond_full_static"]
746
+
747
+ cond_full = self._forward(
748
+ components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps
749
+ )
750
+
751
+ cond_no_control = None
752
+ if needs_control_cfg:
753
+ cond_no_control = self._forward(
754
+ components,
755
+ cond_no_control_static,
756
+ block_state.vision_tokens_target,
757
+ block_state.vision_timesteps,
758
+ )
759
+
760
+ uncond_full = None
761
+ if needs_text_cfg:
762
+ uncond_full = self._forward(
763
+ components,
764
+ uncond_full_static,
765
+ block_state.vision_tokens_full,
766
+ block_state.vision_timesteps,
767
+ )
768
+
769
+ if needs_control_cfg and needs_text_cfg:
770
+ control_cond = cond_no_control + step_control * (cond_full - cond_no_control)
771
+ velocity = uncond_full + step_guidance * (control_cond - uncond_full)
772
+ elif needs_control_cfg:
773
+ velocity = cond_no_control + step_control * (cond_full - cond_no_control)
774
+ elif needs_text_cfg:
775
+ velocity = uncond_full + step_guidance * (cond_full - uncond_full)
776
+ else:
777
+ velocity = cond_full
778
+
779
+ block_state.velocity = velocity * block_state.velocity_mask
780
+ return components, block_state
781
+
782
+
783
+ class Cosmos3TransferLoopSchedulerStep(ModularPipelineBlocks):
784
+ model_name = "cosmos3-omni"
785
+
786
+ @property
787
+ def description(self) -> str:
788
+ return "Steps the scheduler and re-pins the conditioned frames exactly for one transfer iteration."
789
+
790
+ @property
791
+ def expected_components(self) -> list[ComponentSpec]:
792
+ return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
793
+
794
+ @property
795
+ def inputs(self) -> list[InputParam]:
796
+ return [
797
+ InputParam(
798
+ name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to update."
799
+ ),
800
+ InputParam(
801
+ name="velocity",
802
+ type_hint=torch.Tensor,
803
+ required=True,
804
+ description="Predicted (masked) transfer velocity.",
805
+ ),
806
+ InputParam(
807
+ name="velocity_mask",
808
+ type_hint=torch.Tensor,
809
+ required=True,
810
+ description="Mask that zeroes the velocity on conditioned (clean) latent frames.",
811
+ ),
812
+ InputParam(
813
+ name="condition_latents",
814
+ type_hint=torch.Tensor,
815
+ required=True,
816
+ description="Clean target latents on the conditioned frames (the autoregressive seed).",
817
+ ),
818
+ ]
819
+
820
+ @property
821
+ def intermediate_outputs(self) -> list[OutputParam]:
822
+ return [OutputParam("latents", type_hint=torch.Tensor, description="Updated target latents for this chunk.")]
823
+
824
+ @torch.no_grad()
825
+ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
826
+ block_state.latents = components.scheduler.step(
827
+ block_state.velocity.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False
828
+ )[0].squeeze(0)
829
+ # Re-pin conditioned frames exactly (the autoregressive seed), guarding multistep drift.
830
+ block_state.latents = (
831
+ block_state.velocity_mask * block_state.latents
832
+ + (1.0 - block_state.velocity_mask) * block_state.condition_latents
833
+ )
834
+ return components, block_state
835
+
836
+
837
+ # auto_docstring
838
+ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
839
+ """
840
+ Runs the per-chunk transfer denoising loop over scheduler timesteps.
841
+
842
+ Components:
843
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
844
+
845
+ Inputs:
846
+ timesteps (`Tensor`):
847
+ Timesteps for the denoising process.
848
+ num_inference_steps (`int`):
849
+ The number of denoising steps.
850
+ num_warmup_steps (`int`):
851
+ Number of scheduler warmup steps.
852
+ control_latents (`list`):
853
+ Clean control latents for this chunk, one per hint in canonical order.
854
+ latents (`Tensor`):
855
+ Noisy target latents to denoise.
856
+ num_noisy_vision_tokens (`int`):
857
+ Number of noisy target vision tokens denoised each step.
858
+ **denoiser_input_fields (`None`, *optional*):
859
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
860
+ velocity_mask (`Tensor`):
861
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
862
+ guidance_scale (`float`, *optional*, defaults to 6.0):
863
+ Scale for text classifier-free guidance.
864
+ control_guidance (`float`, *optional*, defaults to 1.0):
865
+ Scale for the control (structural) guidance axis.
866
+ guidance_interval (`tuple`, *optional*):
867
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
868
+ control_guidance_interval (`tuple`, *optional*):
869
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
870
+ latents (`Tensor`):
871
+ Noisy target latents to update.
872
+ condition_latents (`Tensor`):
873
+ Clean target latents on the conditioned frames (the autoregressive seed).
874
+
875
+ Outputs:
876
+ latents (`Tensor`):
877
+ Updated target latents for this chunk.
878
+ """
879
+
880
+ block_classes = [
881
+ Cosmos3TransferLoopPrepareStep,
882
+ Cosmos3TransferLoopDenoiser,
883
+ Cosmos3TransferLoopSchedulerStep,
884
+ ]
885
+ block_names = ["prepare_transfer", "denoiser", "update_transfer"]
886
+
887
+ @property
888
+ def description(self) -> str:
889
+ return "Runs the per-chunk transfer denoising loop over scheduler timesteps."
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py ADDED
@@ -0,0 +1,1056 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoTokenizer
3
+
4
+ from ...configuration_utils import FrozenDict
5
+ from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
6
+ from ...pipelines.cosmos.pipeline_cosmos3_omni import (
7
+ _ACTION_RESOLUTION_BINS,
8
+ CosmosActionCondition,
9
+ )
10
+ from ...utils import logging
11
+ from ...video_processor import VideoProcessor
12
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
13
+ from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
14
+ from .modular_pipeline import Cosmos3OmniModularPipeline
15
+
16
+
17
+ logger = logging.get_logger(__name__)
18
+
19
+
20
+ # Transfer conditions on control signals (edge/blur/depth/seg/wsm), so it uses its own system prompt instead of the
21
+ # plain image/video ones. Defined here (not on the task pipeline) so the transfer text block is self-contained.
22
+ _SYSTEM_PROMPT_TRANSFER = (
23
+ "You are a helpful assistant that generates images or videos following the user's instructions"
24
+ " and control signals (edge maps, blur, depth, or segmentation)."
25
+ )
26
+
27
+
28
+ class Cosmos3TextEncoderStep(ModularPipelineBlocks):
29
+ model_name = "cosmos3-omni"
30
+
31
+ @property
32
+ def description(self) -> str:
33
+ return "Prepares non-action prompt token IDs for downstream text-segment packing."
34
+
35
+ @staticmethod
36
+ def _check_inputs(block_state) -> None:
37
+ prompt = block_state.prompt
38
+ negative_prompt = block_state.negative_prompt
39
+
40
+ if not isinstance(prompt, str):
41
+ raise ValueError(
42
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
43
+ )
44
+ if negative_prompt is not None and not isinstance(negative_prompt, str):
45
+ raise ValueError(
46
+ "`negative_prompt` must be a str or None; batched prompts are not supported, "
47
+ f"got {type(negative_prompt).__name__}."
48
+ )
49
+
50
+ @property
51
+ def expected_components(self) -> list[ComponentSpec]:
52
+ return [
53
+ ComponentSpec("text_tokenizer", AutoTokenizer),
54
+ ]
55
+
56
+ @property
57
+ def expected_configs(self) -> list[ConfigSpec]:
58
+ return [
59
+ ConfigSpec(name="default_use_system_prompt", default=True),
60
+ ConfigSpec(name="enable_safety_checker", default=True),
61
+ ]
62
+
63
+ @property
64
+ def inputs(self) -> list[InputParam]:
65
+ return [
66
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
67
+ InputParam.template(
68
+ "negative_prompt", description="The negative text prompt used for classifier-free guidance."
69
+ ),
70
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
71
+ InputParam(
72
+ name="height",
73
+ type_hint=int,
74
+ default=None,
75
+ description="Height of the generated video or image in pixels.",
76
+ ),
77
+ InputParam(
78
+ name="width",
79
+ type_hint=int,
80
+ default=None,
81
+ description="Width of the generated video or image in pixels.",
82
+ ),
83
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
84
+ InputParam(
85
+ name="use_system_prompt",
86
+ type_hint=bool | None,
87
+ default=None,
88
+ description="Whether to prepend the Cosmos3 system prompt.",
89
+ ),
90
+ InputParam(
91
+ name="add_resolution_template",
92
+ type_hint=bool,
93
+ default=True,
94
+ description="Whether to add resolution metadata to the prompt.",
95
+ ),
96
+ InputParam(
97
+ name="add_duration_template",
98
+ type_hint=bool,
99
+ default=True,
100
+ description="Whether to add duration metadata to the prompt.",
101
+ ),
102
+ ]
103
+
104
+ @property
105
+ def intermediate_outputs(self) -> list[OutputParam]:
106
+ return [
107
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
108
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
109
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
110
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
111
+ OutputParam(
112
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
113
+ ),
114
+ ]
115
+
116
+ @torch.no_grad()
117
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
118
+ block_state = self.get_block_state(state)
119
+ if block_state.num_frames is None:
120
+ block_state.num_frames = 189
121
+ if block_state.height is None:
122
+ block_state.height = 720
123
+ if block_state.width is None:
124
+ block_state.width = 1280
125
+ if block_state.use_system_prompt is None:
126
+ block_state.use_system_prompt = components.config.default_use_system_prompt
127
+
128
+ self._check_inputs(block_state)
129
+ if components.requires_safety_checker:
130
+ if getattr(components, "safety_checker", None) is None:
131
+ raise ValueError(
132
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
133
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
134
+ )
135
+ device = components._execution_device
136
+ components.safety_checker.to(device)
137
+ try:
138
+ if not components.safety_checker.check_text_safety(block_state.prompt):
139
+ raise ValueError(
140
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
141
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
142
+ )
143
+ finally:
144
+ components.safety_checker.to("cpu")
145
+
146
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
147
+ block_state.prompt,
148
+ block_state.negative_prompt,
149
+ num_frames=block_state.num_frames,
150
+ height=block_state.height,
151
+ width=block_state.width,
152
+ fps=block_state.fps,
153
+ use_system_prompt=block_state.use_system_prompt,
154
+ add_resolution_template=block_state.add_resolution_template,
155
+ add_duration_template=block_state.add_duration_template,
156
+ action_mode=None,
157
+ action_view_point=None,
158
+ )
159
+
160
+ self.set_block_state(state, block_state)
161
+ return components, state
162
+
163
+
164
+ class Cosmos3TransferTextStep(ModularPipelineBlocks):
165
+ model_name = "cosmos3-omni"
166
+
167
+ @property
168
+ def description(self) -> str:
169
+ return (
170
+ "Tokenizes the transfer prompt with the transfer system prompt. Transfer prompts are pre-upsampled JSON "
171
+ "captions passed through verbatim (no resolution/duration templates), so this is self-contained and does "
172
+ "not reuse the standard text step."
173
+ )
174
+
175
+ @staticmethod
176
+ def _check_inputs(block_state) -> None:
177
+ prompt = block_state.prompt
178
+ negative_prompt = block_state.negative_prompt
179
+
180
+ if not isinstance(prompt, (str, list)) or (
181
+ isinstance(prompt, list) and not all(isinstance(p, str) for p in prompt)
182
+ ):
183
+ raise ValueError(f"`prompt` must be a str or list of str, got {type(prompt).__name__}.")
184
+ if negative_prompt is not None and not isinstance(negative_prompt, (str, list)):
185
+ raise ValueError(
186
+ f"`negative_prompt` must be a str, list of str, or None, got {type(negative_prompt).__name__}."
187
+ )
188
+
189
+ @property
190
+ def expected_components(self) -> list[ComponentSpec]:
191
+ return [
192
+ ComponentSpec("text_tokenizer", AutoTokenizer),
193
+ ]
194
+
195
+ @property
196
+ def inputs(self) -> list[InputParam]:
197
+ return [
198
+ InputParam(
199
+ name="prompt",
200
+ type_hint=str,
201
+ required=True,
202
+ description="The text prompt that guides Cosmos3 generation.",
203
+ ),
204
+ InputParam(
205
+ name="negative_prompt",
206
+ type_hint=str,
207
+ default=None,
208
+ description="The negative text prompt used for classifier-free guidance.",
209
+ ),
210
+ InputParam(
211
+ name="use_system_prompt",
212
+ type_hint=bool,
213
+ default=True,
214
+ description="Whether to prepend the Cosmos3 transfer system prompt.",
215
+ ),
216
+ ]
217
+
218
+ @property
219
+ def intermediate_outputs(self) -> list[OutputParam]:
220
+ return [
221
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
222
+ OutputParam(
223
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
224
+ ),
225
+ ]
226
+
227
+ @torch.no_grad()
228
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
229
+ block_state = self.get_block_state(state)
230
+ self._check_inputs(block_state)
231
+
232
+ if isinstance(block_state.prompt, list):
233
+ block_state.prompt = block_state.prompt[0]
234
+ if isinstance(block_state.negative_prompt, list):
235
+ block_state.negative_prompt = block_state.negative_prompt[0]
236
+
237
+ if components.requires_safety_checker:
238
+ if getattr(components, "safety_checker", None) is None:
239
+ raise ValueError(
240
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
241
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
242
+ )
243
+ device = components._execution_device
244
+ components.safety_checker.to(device)
245
+ try:
246
+ if not components.safety_checker.check_text_safety(block_state.prompt):
247
+ raise ValueError(
248
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
249
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
250
+ )
251
+ finally:
252
+ components.safety_checker.to("cpu")
253
+
254
+ # Transfer prompts are pre-upsampled JSON captions: tokenize them verbatim (no resolution/duration templates)
255
+ # under the transfer system prompt. Kept self-contained here rather than adding a flag to the standard step.
256
+ negative_prompt = block_state.negative_prompt if block_state.negative_prompt is not None else ""
257
+ special_tokens = components.llm_special_tokens
258
+
259
+ def _tokenize(text: str) -> list[int]:
260
+ conversations = []
261
+ if block_state.use_system_prompt:
262
+ conversations.append({"role": "system", "content": _SYSTEM_PROMPT_TRANSFER})
263
+ conversations.append({"role": "user", "content": text})
264
+ encoding = components.text_tokenizer.apply_chat_template(
265
+ conversations,
266
+ tokenize=True,
267
+ add_generation_prompt=True,
268
+ add_vision_id=False,
269
+ return_dict=True,
270
+ )
271
+ return list(encoding.input_ids) + [
272
+ special_tokens["eos_token_id"],
273
+ special_tokens["start_of_generation"],
274
+ ]
275
+
276
+ block_state.cond_input_ids = _tokenize(block_state.prompt)
277
+ block_state.uncond_input_ids = _tokenize(negative_prompt)
278
+
279
+ self.set_block_state(state, block_state)
280
+ return components, state
281
+
282
+
283
+ class Cosmos3DistilledTextEncoderStep(ModularPipelineBlocks):
284
+ model_name = "cosmos3-omni"
285
+
286
+ @property
287
+ def description(self) -> str:
288
+ return (
289
+ "Prepares distilled prompt token IDs. Classifier-free guidance is baked into the weights, so "
290
+ "`negative_prompt` is not exposed and the unconditional branch is derived from an empty prompt."
291
+ )
292
+
293
+ @staticmethod
294
+ def _check_inputs(block_state) -> None:
295
+ prompt = block_state.prompt
296
+ if not isinstance(prompt, str):
297
+ raise ValueError(
298
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
299
+ )
300
+
301
+ @property
302
+ def expected_components(self) -> list[ComponentSpec]:
303
+ return [
304
+ ComponentSpec("text_tokenizer", AutoTokenizer),
305
+ ]
306
+
307
+ @property
308
+ def expected_configs(self) -> list[ConfigSpec]:
309
+ return [
310
+ ConfigSpec(name="default_use_system_prompt", default=True),
311
+ ConfigSpec(name="enable_safety_checker", default=True),
312
+ ]
313
+
314
+ @property
315
+ def inputs(self) -> list[InputParam]:
316
+ return [
317
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
318
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
319
+ InputParam(
320
+ name="height",
321
+ type_hint=int,
322
+ default=None,
323
+ description="Height of the generated video or image in pixels.",
324
+ ),
325
+ InputParam(
326
+ name="width",
327
+ type_hint=int,
328
+ default=None,
329
+ description="Width of the generated video or image in pixels.",
330
+ ),
331
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
332
+ InputParam(
333
+ name="use_system_prompt",
334
+ type_hint=bool,
335
+ default=True,
336
+ description="Whether to prepend the Cosmos3 system prompt.",
337
+ ),
338
+ InputParam(
339
+ name="add_resolution_template",
340
+ type_hint=bool,
341
+ default=True,
342
+ description="Whether to add resolution metadata to the prompt.",
343
+ ),
344
+ InputParam(
345
+ name="add_duration_template",
346
+ type_hint=bool,
347
+ default=True,
348
+ description="Whether to add duration metadata to the prompt.",
349
+ ),
350
+ ]
351
+
352
+ @property
353
+ def intermediate_outputs(self) -> list[OutputParam]:
354
+ return [
355
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
356
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
357
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
358
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
359
+ OutputParam(
360
+ "uncond_input_ids",
361
+ type_hint=torch.Tensor,
362
+ description="Token IDs for the unconditional prompt (empty prompt; guidance is baked in).",
363
+ ),
364
+ ]
365
+
366
+ @torch.no_grad()
367
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
368
+ block_state = self.get_block_state(state)
369
+ if block_state.num_frames is None:
370
+ block_state.num_frames = 189
371
+ if block_state.height is None:
372
+ block_state.height = 720
373
+ if block_state.width is None:
374
+ block_state.width = 1280
375
+
376
+ self._check_inputs(block_state)
377
+ if components.requires_safety_checker:
378
+ if getattr(components, "safety_checker", None) is None:
379
+ raise ValueError(
380
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
381
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
382
+ )
383
+ device = components._execution_device
384
+ components.safety_checker.to(device)
385
+ try:
386
+ if not components.safety_checker.check_text_safety(block_state.prompt):
387
+ raise ValueError(
388
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
389
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
390
+ )
391
+ finally:
392
+ components.safety_checker.to("cpu")
393
+
394
+ # Guidance is baked into distilled weights: the unconditional branch is built from an empty prompt
395
+ # (negative_prompt is not a user-facing input) so the downstream text-segment packing contract still holds.
396
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
397
+ block_state.prompt,
398
+ None,
399
+ num_frames=block_state.num_frames,
400
+ height=block_state.height,
401
+ width=block_state.width,
402
+ fps=block_state.fps,
403
+ use_system_prompt=block_state.use_system_prompt,
404
+ add_resolution_template=block_state.add_resolution_template,
405
+ add_duration_template=block_state.add_duration_template,
406
+ action_mode=None,
407
+ action_view_point=None,
408
+ )
409
+
410
+ self.set_block_state(state, block_state)
411
+ return components, state
412
+
413
+
414
+ class Cosmos3ActionTextStep(ModularPipelineBlocks):
415
+ model_name = "cosmos3-omni"
416
+
417
+ @property
418
+ def description(self) -> str:
419
+ return "Prepares action prompt token IDs from prompt + action metadata."
420
+
421
+ @staticmethod
422
+ def _check_inputs(block_state) -> None:
423
+ prompt = block_state.prompt
424
+ negative_prompt = block_state.negative_prompt
425
+ action = block_state.action
426
+ num_frames = block_state.num_frames
427
+ height = block_state.height
428
+ width = block_state.width
429
+ if not isinstance(prompt, str):
430
+ raise ValueError(
431
+ f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}."
432
+ )
433
+ if negative_prompt is not None and not isinstance(negative_prompt, str):
434
+ raise ValueError(
435
+ "`negative_prompt` must be a str or None; batched prompts are not supported, "
436
+ f"got {type(negative_prompt).__name__}."
437
+ )
438
+ if action is None:
439
+ raise ValueError("`action` is required for Cosmos3ActionTextStep.")
440
+ if action.image is None and action.video is None:
441
+ raise ValueError("`action.image` or `action.video` must be provided for action-conditioned generation.")
442
+ if num_frames is not None:
443
+ raise ValueError("`num_frames` has to be None if action is not None.")
444
+ if height is not None or width is not None:
445
+ raise ValueError("`height` and `width` have to be None if action is not None.")
446
+
447
+ @property
448
+ def expected_components(self) -> list[ComponentSpec]:
449
+ return [
450
+ ComponentSpec("text_tokenizer", AutoTokenizer),
451
+ ComponentSpec(
452
+ "video_processor",
453
+ VideoProcessor,
454
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
455
+ default_creation_method="from_config",
456
+ ),
457
+ ]
458
+
459
+ @property
460
+ def expected_configs(self) -> list[ConfigSpec]:
461
+ return [
462
+ ConfigSpec(name="default_use_system_prompt", default=True),
463
+ ConfigSpec(name="enable_safety_checker", default=True),
464
+ ]
465
+
466
+ @property
467
+ def inputs(self) -> list[InputParam]:
468
+ return [
469
+ InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."),
470
+ InputParam.template(
471
+ "negative_prompt", description="The negative text prompt used for classifier-free guidance."
472
+ ),
473
+ InputParam(
474
+ name="action",
475
+ type_hint=CosmosActionCondition,
476
+ required=True,
477
+ description="Action-conditioning metadata and its reference visual input.",
478
+ ),
479
+ InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."),
480
+ InputParam(
481
+ name="height",
482
+ type_hint=int,
483
+ default=None,
484
+ description="Height of the generated video or image in pixels.",
485
+ ),
486
+ InputParam(
487
+ name="width",
488
+ type_hint=int,
489
+ default=None,
490
+ description="Width of the generated video or image in pixels.",
491
+ ),
492
+ InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
493
+ InputParam(
494
+ name="use_system_prompt",
495
+ type_hint=bool | None,
496
+ default=None,
497
+ description="Whether to prepend the Cosmos3 system prompt.",
498
+ ),
499
+ InputParam(
500
+ name="add_resolution_template",
501
+ type_hint=bool,
502
+ default=True,
503
+ description="Whether to add resolution metadata to the prompt.",
504
+ ),
505
+ InputParam(
506
+ name="add_duration_template",
507
+ type_hint=bool,
508
+ default=True,
509
+ description="Whether to add duration metadata to the prompt.",
510
+ ),
511
+ ]
512
+
513
+ @property
514
+ def intermediate_outputs(self) -> list[OutputParam]:
515
+ return [
516
+ OutputParam("action_mode", type_hint=str, description="Requested action-generation mode."),
517
+ OutputParam("num_frames", type_hint=int, description="Number of frames to generate."),
518
+ OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."),
519
+ OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."),
520
+ OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."),
521
+ OutputParam(
522
+ "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt."
523
+ ),
524
+ ]
525
+
526
+ @torch.no_grad()
527
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
528
+ block_state = self.get_block_state(state)
529
+ self._check_inputs(block_state)
530
+ if block_state.use_system_prompt is None:
531
+ block_state.use_system_prompt = components.config.default_use_system_prompt
532
+
533
+ action = block_state.action
534
+ block_state.action_mode = action.mode
535
+ block_state.num_frames = action.chunk_size + 1
536
+ conditioning_clip = [action.image] if action.image is not None else action.video
537
+ probe = components.video_processor.preprocess_video(conditioning_clip)
538
+ source_h, source_w = int(probe.shape[-2]), int(probe.shape[-1])
539
+ resolution_key = str(action.resolution_tier)
540
+ block_state.height, block_state.width = VideoProcessor.classify_height_width_bin(
541
+ source_h, source_w, ratios=_ACTION_RESOLUTION_BINS[resolution_key]
542
+ )
543
+
544
+ if components.requires_safety_checker:
545
+ if getattr(components, "safety_checker", None) is None:
546
+ raise ValueError(
547
+ "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it "
548
+ "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`."
549
+ )
550
+ device = components._execution_device
551
+ components.safety_checker.to(device)
552
+ try:
553
+ if not components.safety_checker.check_text_safety(block_state.prompt):
554
+ raise ValueError(
555
+ f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. "
556
+ "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement."
557
+ )
558
+ finally:
559
+ components.safety_checker.to("cpu")
560
+
561
+ block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt(
562
+ block_state.prompt,
563
+ block_state.negative_prompt,
564
+ num_frames=block_state.num_frames,
565
+ height=block_state.height,
566
+ width=block_state.width,
567
+ fps=block_state.fps,
568
+ use_system_prompt=block_state.use_system_prompt,
569
+ add_resolution_template=block_state.add_resolution_template,
570
+ add_duration_template=block_state.add_duration_template,
571
+ action_mode=block_state.action_mode,
572
+ action_view_point=action.view_point,
573
+ )
574
+
575
+ self.set_block_state(state, block_state)
576
+ return components, state
577
+
578
+
579
+ class Cosmos3ImageVaeEncoderStep(ModularPipelineBlocks):
580
+ model_name = "cosmos3-omni"
581
+
582
+ @property
583
+ def description(self) -> str:
584
+ return "Encodes non-action image-to-video conditioning into Cosmos3 vision latents."
585
+
586
+ @property
587
+ def expected_components(self) -> list[ComponentSpec]:
588
+ return [
589
+ ComponentSpec("vae", AutoencoderKLWan),
590
+ ComponentSpec(
591
+ "video_processor",
592
+ VideoProcessor,
593
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
594
+ default_creation_method="from_config",
595
+ ),
596
+ ]
597
+
598
+ @property
599
+ def inputs(self) -> list[InputParam]:
600
+ return [
601
+ InputParam(name="image", default=None, description="Reference image for image-to-video conditioning."),
602
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
603
+ InputParam(
604
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
605
+ ),
606
+ InputParam(
607
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
608
+ ),
609
+ ]
610
+
611
+ @property
612
+ def intermediate_outputs(self) -> list[OutputParam]:
613
+ return [
614
+ OutputParam(
615
+ "x0_tokens_vision",
616
+ type_hint=torch.Tensor,
617
+ description="Vision latents encoded from the conditioning image or video.",
618
+ ),
619
+ OutputParam(
620
+ "vision_condition_frames",
621
+ type_hint=list[int],
622
+ description="Latent-frame indexes fixed by visual conditioning.",
623
+ ),
624
+ ]
625
+
626
+ @torch.no_grad()
627
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
628
+ block_state = self.get_block_state(state)
629
+
630
+ device = components._execution_device
631
+ dtype = components.vae.dtype
632
+
633
+ if block_state.image is None:
634
+ raise ValueError("`Cosmos3ImageVaeEncoderStep` requires an `image` input.")
635
+ if block_state.num_frames == 1:
636
+ raise ValueError(
637
+ "`image` conditioning requires `num_frames` > 1; image-to-image generation is not supported."
638
+ )
639
+ if block_state.num_frames < 1:
640
+ raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.")
641
+
642
+ sf = int(components.vae.config.scale_factor_spatial)
643
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
644
+ raise ValueError(
645
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
646
+ )
647
+
648
+ conditioning_frame_2d = components.video_processor.preprocess(
649
+ block_state.image, height=block_state.height, width=block_state.width
650
+ ).to(device=device, dtype=dtype)
651
+
652
+ vision_tensor = torch.zeros(
653
+ 1,
654
+ 3,
655
+ block_state.num_frames,
656
+ block_state.height,
657
+ block_state.width,
658
+ dtype=dtype,
659
+ device=device,
660
+ )
661
+ vision_tensor[:, :, 0] = conditioning_frame_2d
662
+ vision_tensor[:, :, 1:] = conditioning_frame_2d.unsqueeze(2).expand(-1, -1, block_state.num_frames - 1, -1, -1)
663
+
664
+ block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
665
+ block_state.vision_condition_frames = [0]
666
+
667
+ self.set_block_state(state, block_state)
668
+ return components, state
669
+
670
+
671
+ class Cosmos3VideoVaeEncoderStep(ModularPipelineBlocks):
672
+ model_name = "cosmos3-omni"
673
+
674
+ @property
675
+ def description(self) -> str:
676
+ return "Encodes non-action video conditioning into Cosmos3 vision latents."
677
+
678
+ @property
679
+ def expected_components(self) -> list[ComponentSpec]:
680
+ return [
681
+ ComponentSpec("vae", AutoencoderKLWan),
682
+ ComponentSpec(
683
+ "video_processor",
684
+ VideoProcessor,
685
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
686
+ default_creation_method="from_config",
687
+ ),
688
+ ]
689
+
690
+ @property
691
+ def inputs(self) -> list[InputParam]:
692
+ return [
693
+ InputParam(name="video", default=None, description="Reference video for video-to-video conditioning."),
694
+ InputParam(
695
+ name="condition_frame_indexes_vision",
696
+ type_hint=tuple[int, ...] | list[int],
697
+ default=(0, 1),
698
+ description="Latent-frame indexes to preserve from the conditioning video.",
699
+ ),
700
+ InputParam(
701
+ name="condition_video_keep",
702
+ type_hint=str,
703
+ default="first",
704
+ description="Which end of a longer conditioning video to use: `first` or `last`.",
705
+ ),
706
+ InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."),
707
+ InputParam(
708
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
709
+ ),
710
+ InputParam(
711
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
712
+ ),
713
+ ]
714
+
715
+ @property
716
+ def intermediate_outputs(self) -> list[OutputParam]:
717
+ return [
718
+ OutputParam(
719
+ "x0_tokens_vision",
720
+ type_hint=torch.Tensor,
721
+ description="Vision latents encoded from the conditioning image or video.",
722
+ ),
723
+ OutputParam(
724
+ "vision_condition_frames",
725
+ type_hint=list[int],
726
+ description="Latent-frame indexes fixed by visual conditioning.",
727
+ ),
728
+ ]
729
+
730
+ @torch.no_grad()
731
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
732
+ block_state = self.get_block_state(state)
733
+
734
+ device = components._execution_device
735
+ dtype = components.vae.dtype
736
+
737
+ if block_state.video is None:
738
+ raise ValueError("`Cosmos3VideoVaeEncoderStep` requires a `video` input.")
739
+ if block_state.num_frames == 1:
740
+ raise ValueError("`video` conditioning requires `num_frames` > 1.")
741
+ if block_state.num_frames < 1:
742
+ raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.")
743
+
744
+ sf = int(components.vae.config.scale_factor_spatial)
745
+ if block_state.height % sf != 0 or block_state.width % sf != 0:
746
+ raise ValueError(
747
+ f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})."
748
+ )
749
+
750
+ if not isinstance(block_state.condition_frame_indexes_vision, (list, tuple)) or isinstance(
751
+ block_state.condition_frame_indexes_vision, (str, bytes)
752
+ ):
753
+ raise ValueError(
754
+ "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got "
755
+ f"{block_state.condition_frame_indexes_vision!r}."
756
+ )
757
+ if not all(isinstance(index, int) and index >= 0 for index in block_state.condition_frame_indexes_vision):
758
+ raise ValueError(
759
+ "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got "
760
+ f"{block_state.condition_frame_indexes_vision!r}."
761
+ )
762
+ if block_state.condition_video_keep not in {"first", "last"}:
763
+ raise ValueError("`condition_video_keep` must be either 'first' or 'last'.")
764
+
765
+ indexes = tuple(block_state.condition_frame_indexes_vision)
766
+ if not indexes:
767
+ raise ValueError("`condition_frame_indexes_vision` must contain at least one index.")
768
+ latent_t = (block_state.num_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1
769
+ if max(indexes) >= latent_t:
770
+ raise ValueError(
771
+ f"`condition_frame_indexes_vision` {indexes} contains an index outside the latent timeline "
772
+ f"(latent_frames={latent_t} for num_frames={block_state.num_frames})."
773
+ )
774
+
775
+ condition_indexes_vision = indexes
776
+ conditioning_frames_3d = components.video_processor.preprocess_video(
777
+ block_state.video, height=block_state.height, width=block_state.width
778
+ ).to(device=device, dtype=dtype)
779
+ temporal_compression = int(components.vae.config.scale_factor_temporal)
780
+ max_cond_frames = max(condition_indexes_vision) * temporal_compression + 1
781
+ if block_state.condition_video_keep == "first":
782
+ conditioning_frames_3d = conditioning_frames_3d[:, :, :max_cond_frames]
783
+ else:
784
+ conditioning_frames_3d = conditioning_frames_3d[:, :, -max_cond_frames:]
785
+
786
+ vision_tensor = torch.zeros(
787
+ 1,
788
+ 3,
789
+ block_state.num_frames,
790
+ block_state.height,
791
+ block_state.width,
792
+ dtype=dtype,
793
+ device=device,
794
+ )
795
+ t_fill = min(conditioning_frames_3d.shape[2], block_state.num_frames)
796
+ vision_tensor[:, :, :t_fill] = conditioning_frames_3d[:, :, :t_fill]
797
+ if t_fill < block_state.num_frames:
798
+ vision_tensor[:, :, t_fill:] = vision_tensor[:, :, t_fill - 1 : t_fill].expand(
799
+ -1, -1, block_state.num_frames - t_fill, -1, -1
800
+ )
801
+ vision_condition_frames = list(condition_indexes_vision)
802
+
803
+ block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
804
+ block_state.vision_condition_frames = vision_condition_frames
805
+
806
+ self.set_block_state(state, block_state)
807
+ return components, state
808
+
809
+
810
+ class Cosmos3TransferChunkVaeEncoderStep(ModularPipelineBlocks):
811
+ model_name = "cosmos3-omni"
812
+
813
+ @property
814
+ def description(self) -> str:
815
+ return (
816
+ "Per-chunk transfer VAE encode: slices + pads this chunk's control maps, seeds the target's conditioning "
817
+ "frames (first chunk from the input video, later chunks from the previous chunk's tail), and encodes both "
818
+ "the controls and the seeded target into clean Cosmos3 vision latents. Runs inside the autoregressive "
819
+ "chunk loop."
820
+ )
821
+
822
+ @property
823
+ def expected_components(self) -> list[ComponentSpec]:
824
+ return [
825
+ ComponentSpec("vae", AutoencoderKLWan),
826
+ ComponentSpec(
827
+ "video_processor",
828
+ VideoProcessor,
829
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
830
+ default_creation_method="from_config",
831
+ ),
832
+ ]
833
+
834
+ @property
835
+ def inputs(self) -> list[InputParam]:
836
+ return [
837
+ InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."),
838
+ InputParam(
839
+ name="previous_output",
840
+ default=None,
841
+ description="Decoded pixels of the previous chunk, used to seed later chunks.",
842
+ ),
843
+ InputParam(
844
+ name="control_frames",
845
+ type_hint=dict,
846
+ required=True,
847
+ description="Preprocessed, time-padded control maps in canonical hint order.",
848
+ ),
849
+ InputParam(name="chunk_frames", type_hint=int, required=True, description="Pixel frames per chunk."),
850
+ InputParam(
851
+ name="total_frames", type_hint=int, required=True, description="Total number of output frames."
852
+ ),
853
+ InputParam(name="stride", type_hint=int, required=True, description="Frame stride between chunks."),
854
+ InputParam(
855
+ name="height", type_hint=int, required=True, description="Height of the generated video in pixels."
856
+ ),
857
+ InputParam(
858
+ name="width", type_hint=int, required=True, description="Width of the generated video in pixels."
859
+ ),
860
+ InputParam(
861
+ name="video",
862
+ default=None,
863
+ description="Optional input video that seeds the first chunk's conditioning.",
864
+ ),
865
+ InputParam(
866
+ name="num_first_chunk_conditional_frames",
867
+ type_hint=int,
868
+ default=0,
869
+ description="Number of frames the first chunk reuses from the input video.",
870
+ ),
871
+ InputParam(
872
+ name="num_conditional_frames",
873
+ type_hint=int,
874
+ default=1,
875
+ description="Number of frames each later chunk reuses from the previous chunk's tail.",
876
+ ),
877
+ ]
878
+
879
+ @property
880
+ def intermediate_outputs(self) -> list[OutputParam]:
881
+ return [
882
+ OutputParam(
883
+ "control_latents",
884
+ type_hint=list[torch.Tensor],
885
+ description="Clean control latents for this chunk, one per hint in canonical order.",
886
+ ),
887
+ OutputParam(
888
+ "x0_tokens_vision",
889
+ type_hint=torch.Tensor,
890
+ description="Clean target vision latents encoded from the seeded target frames.",
891
+ ),
892
+ OutputParam(
893
+ "current_conditional_frames",
894
+ type_hint=int,
895
+ description="Number of pixel frames actually used to seed this chunk's target.",
896
+ ),
897
+ ]
898
+
899
+ @torch.no_grad()
900
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
901
+ block_state = self.get_block_state(state)
902
+ device = components._execution_device
903
+ dtype = components.vae.dtype
904
+
905
+ chunk_id = block_state.chunk_id
906
+ chunk_frames = block_state.chunk_frames
907
+ height = block_state.height
908
+ width = block_state.width
909
+
910
+ # Slice this chunk's window out of the (padded) control maps and reflect-pad it up to a full chunk (repeat the
911
+ # last frame once too short to keep reflecting). control_frames is already in canonical hint order.
912
+ start_frame = chunk_id * block_state.stride
913
+ end_frame = min(start_frame + chunk_frames, block_state.total_frames)
914
+ chunk_controls = []
915
+ for frames in block_state.control_frames.values():
916
+ frames = frames[:, :, start_frame:end_frame]
917
+ while frames.shape[2] < chunk_frames:
918
+ pad_len = min(frames.shape[2] - 1, chunk_frames - frames.shape[2])
919
+ if pad_len <= 0:
920
+ pad_frame = frames[:, :, -1:].repeat(1, 1, chunk_frames - frames.shape[2], 1, 1)
921
+ frames = torch.cat([frames, pad_frame], dim=2)
922
+ break
923
+ frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2)
924
+ chunk_controls.append(frames)
925
+
926
+ # Seed the target with conditioning frames (first chunk from the input video, later chunks from the
927
+ # previous chunk's tail), repeat-padding the remaining frames so the whole clip is well-defined.
928
+ target = torch.zeros(1, 3, chunk_frames, height, width, device=device, dtype=dtype)
929
+ current_conditional_frames = 0
930
+ if chunk_id == 0 and block_state.num_first_chunk_conditional_frames > 0 and block_state.video is not None:
931
+ input_frames = components.video_processor.preprocess_video(
932
+ block_state.video, height=height, width=width
933
+ ).to(device=device, dtype=dtype)
934
+ current_conditional_frames = min(
935
+ block_state.num_first_chunk_conditional_frames, input_frames.shape[2], chunk_frames
936
+ )
937
+ if current_conditional_frames > 0:
938
+ target[:, :, :current_conditional_frames] = input_frames[:, :, :current_conditional_frames]
939
+ elif chunk_id > 0 and block_state.previous_output is not None:
940
+ current_conditional_frames = min(
941
+ block_state.num_conditional_frames, block_state.previous_output.shape[2], chunk_frames
942
+ )
943
+ if current_conditional_frames > 0:
944
+ target[:, :, :current_conditional_frames] = block_state.previous_output[
945
+ :, :, -current_conditional_frames:
946
+ ].to(device=device, dtype=dtype)
947
+ if 0 < current_conditional_frames < chunk_frames:
948
+ fill = target[:, :, current_conditional_frames - 1 : current_conditional_frames]
949
+ target[:, :, current_conditional_frames:] = fill.expand(
950
+ -1, -1, chunk_frames - current_conditional_frames, -1, -1
951
+ )
952
+
953
+ block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls]
954
+ block_state.x0_tokens_vision = components._encode_video(target).contiguous().float()
955
+ block_state.current_conditional_frames = current_conditional_frames
956
+
957
+ self.set_block_state(state, block_state)
958
+ return components, state
959
+
960
+
961
+ class Cosmos3ActionVisionVaeEncoderStep(ModularPipelineBlocks):
962
+ model_name = "cosmos3-omni"
963
+
964
+ @property
965
+ def description(self) -> str:
966
+ return (
967
+ "Prepares action-conditioned vision latents and action frame metadata. "
968
+ "Only the action visual reference (image/video) is VAE-encoded; action vectors are handled separately."
969
+ )
970
+
971
+ @property
972
+ def expected_components(self) -> list[ComponentSpec]:
973
+ return [
974
+ ComponentSpec("vae", AutoencoderKLWan),
975
+ ComponentSpec(
976
+ "video_processor",
977
+ VideoProcessor,
978
+ config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}),
979
+ default_creation_method="from_config",
980
+ ),
981
+ ]
982
+
983
+ @property
984
+ def inputs(self) -> list[InputParam]:
985
+ return [
986
+ InputParam(
987
+ name="action",
988
+ type_hint=CosmosActionCondition,
989
+ required=True,
990
+ description="Action-conditioning metadata and its reference visual input.",
991
+ ),
992
+ ]
993
+
994
+ @property
995
+ def intermediate_outputs(self) -> list[OutputParam]:
996
+ return [
997
+ OutputParam(
998
+ "x0_tokens_vision",
999
+ type_hint=torch.Tensor,
1000
+ description="Vision latents encoded from the conditioning image or video.",
1001
+ ),
1002
+ OutputParam(
1003
+ "vision_condition_frames",
1004
+ type_hint=list[int],
1005
+ description="Latent-frame indexes fixed by visual conditioning.",
1006
+ ),
1007
+ OutputParam(
1008
+ "action_condition_frame_indexes",
1009
+ type_hint=list[int],
1010
+ description="Action-frame indexes fixed by action conditioning.",
1011
+ ),
1012
+ ]
1013
+
1014
+ @torch.no_grad()
1015
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
1016
+ block_state = self.get_block_state(state)
1017
+
1018
+ device = components._execution_device
1019
+ dtype = components.vae.dtype
1020
+
1021
+ action = block_state.action
1022
+ target_frames = action.chunk_size + 1
1023
+ conditioning_clip = [action.image] if action.image is not None else action.video
1024
+ vision_tensor, action_image_size, _, _ = components._prepare_action_video_conditioning(
1025
+ conditioning_clip,
1026
+ action.resolution_tier,
1027
+ target_frames,
1028
+ device=device,
1029
+ dtype=dtype,
1030
+ )
1031
+
1032
+ if action.mode == "forward_dynamics":
1033
+ vision_condition_frames = [0]
1034
+ action_condition_frame_indexes = list(range(action.chunk_size))
1035
+ elif action.mode == "policy":
1036
+ vision_condition_frames = [0]
1037
+ action_condition_frame_indexes = []
1038
+ elif action.mode == "inverse_dynamics":
1039
+ latent_frames = (target_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1
1040
+ vision_condition_frames = list(range(latent_frames))
1041
+ action_condition_frame_indexes = []
1042
+ else:
1043
+ raise ValueError(
1044
+ f"Unsupported action_mode={action.mode!r}; expected one of ['forward_dynamics', 'inverse_dynamics', 'policy']."
1045
+ )
1046
+
1047
+ x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float()
1048
+ if action_image_size is not None:
1049
+ x0_tokens_vision = components._remove_action_video_padding_from_latent(x0_tokens_vision, action_image_size)
1050
+
1051
+ block_state.x0_tokens_vision = x0_tokens_vision
1052
+ block_state.vision_condition_frames = vision_condition_frames
1053
+ block_state.action_condition_frame_indexes = action_condition_frame_indexes
1054
+
1055
+ self.set_block_state(state, block_state)
1056
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py ADDED
@@ -0,0 +1,1301 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ..modular_pipeline import (
4
+ AutoPipelineBlocks,
5
+ ConditionalPipelineBlocks,
6
+ PipelineState,
7
+ SequentialPipelineBlocks,
8
+ )
9
+ from ..modular_pipeline_utils import InputParam, OutputParam
10
+ from .after_decode import Cosmos3ActionOutputStep
11
+ from .before_denoise import (
12
+ Cosmos3ActionDenoiseInputStep,
13
+ Cosmos3ActionPackSequenceStep,
14
+ Cosmos3ActionPrepareLatentsStep,
15
+ Cosmos3PrepareTextSegmentsStep,
16
+ Cosmos3SetTimestepsStep,
17
+ Cosmos3SoundDenoiseInputStep,
18
+ Cosmos3SoundPackSequenceStep,
19
+ Cosmos3SoundPrepareLatentsStep,
20
+ Cosmos3TransferPackSequenceStep,
21
+ Cosmos3TransferPrepareLatentsStep,
22
+ Cosmos3TransferSetTimestepsStep,
23
+ Cosmos3VisionDenoiseInputStep,
24
+ Cosmos3VisionPackSequenceStep,
25
+ Cosmos3VisionPrepareLatentsStep,
26
+ )
27
+ from .before_encoder import Cosmos3TransferSetupStep
28
+ from .decoders import (
29
+ Cosmos3SoundDecodeStep,
30
+ Cosmos3TransferDecodeChunkStep,
31
+ Cosmos3TransferStitchStep,
32
+ Cosmos3VideoDecodeStep,
33
+ )
34
+ from .denoise import (
35
+ Cosmos3TransferDenoiseStep,
36
+ Cosmos3VisionActionDenoiseStep,
37
+ Cosmos3VisionDenoiseStep,
38
+ Cosmos3VisionSoundActionDenoiseStep,
39
+ Cosmos3VisionSoundDenoiseStep,
40
+ )
41
+ from .encoders import (
42
+ Cosmos3ActionTextStep,
43
+ Cosmos3ActionVisionVaeEncoderStep,
44
+ Cosmos3ImageVaeEncoderStep,
45
+ Cosmos3TextEncoderStep,
46
+ Cosmos3TransferChunkVaeEncoderStep,
47
+ Cosmos3TransferTextStep,
48
+ Cosmos3VideoVaeEncoderStep,
49
+ )
50
+ from .modular_pipeline import Cosmos3OmniModularPipeline
51
+
52
+
53
+ # auto_docstring
54
+ class Cosmos3TransferTextBlocks(SequentialPipelineBlocks):
55
+ """
56
+ Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in
57
+ transfer mode using the per-chunk frame count.
58
+
59
+ Components:
60
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
61
+
62
+ Inputs:
63
+ control_videos (`dict`):
64
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
65
+ height (`int`, *optional*):
66
+ Height of the generated video in pixels.
67
+ width (`int`, *optional*):
68
+ Width of the generated video in pixels.
69
+ num_frames (`int`, *optional*):
70
+ Optional cap on the number of output frames (defaults to the control video length).
71
+ num_video_frames_per_chunk (`int`, *optional*):
72
+ Number of pixel frames generated per autoregressive chunk.
73
+ num_conditional_frames (`int`, *optional*, defaults to 1):
74
+ Number of frames each chunk reuses from the previous chunk's tail.
75
+ prompt (`str`):
76
+ The text prompt that guides Cosmos3 generation.
77
+ negative_prompt (`str`, *optional*):
78
+ The negative text prompt used for classifier-free guidance.
79
+ use_system_prompt (`bool`, *optional*, defaults to True):
80
+ Whether to prepend the Cosmos3 transfer system prompt.
81
+
82
+ Outputs:
83
+ height (`int`):
84
+ Resolved output height in pixels.
85
+ width (`int`):
86
+ Resolved output width in pixels.
87
+ control_frames (`dict`):
88
+ Preprocessed, time-padded control maps in canonical hint order.
89
+ total_frames (`int`):
90
+ Total number of output frames to generate.
91
+ chunk_frames (`int`):
92
+ Number of pixel frames per autoregressive chunk.
93
+ num_chunks (`int`):
94
+ Number of autoregressive chunks.
95
+ stride (`int`):
96
+ Frame stride between consecutive chunks.
97
+ cond_input_ids (`Tensor`):
98
+ Token IDs for the conditional prompt.
99
+ uncond_input_ids (`Tensor`):
100
+ Token IDs for the unconditional prompt.
101
+ """
102
+
103
+ model_name = "cosmos3-omni"
104
+ block_classes = [Cosmos3TransferSetupStep, Cosmos3TransferTextStep]
105
+ block_names = ["setup", "transfer_text"]
106
+
107
+ @property
108
+ def description(self):
109
+ return (
110
+ "Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) "
111
+ "prompt in transfer mode using the per-chunk frame count."
112
+ )
113
+
114
+
115
+ # auto_docstring
116
+ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks):
117
+ """
118
+ Auto text encoder block for Cosmos3.
119
+ - Cosmos3TransferTextBlocks runs when control_videos are provided.
120
+ - Cosmos3ActionTextStep runs when action is provided.
121
+ - Cosmos3TextEncoderStep runs otherwise.
122
+
123
+ Components:
124
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`)
125
+
126
+ Configs:
127
+ default_use_system_prompt (default: True) enable_safety_checker (default: True)
128
+
129
+ Inputs:
130
+ control_videos (`dict`, *optional*):
131
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
132
+ height (`int`, *optional*):
133
+ Height of the generated video in pixels.
134
+ width (`int`, *optional*):
135
+ Width of the generated video in pixels.
136
+ num_frames (`int`, *optional*):
137
+ Optional cap on the number of output frames (defaults to the control video length).
138
+ num_video_frames_per_chunk (`int`, *optional*):
139
+ Number of pixel frames generated per autoregressive chunk.
140
+ num_conditional_frames (`int`, *optional*, defaults to 1):
141
+ Number of frames each chunk reuses from the previous chunk's tail.
142
+ prompt (`str`):
143
+ The text prompt that guides Cosmos3 generation.
144
+ negative_prompt (`str`, *optional*):
145
+ The negative text prompt used for classifier-free guidance.
146
+ use_system_prompt (`bool`, *optional*, defaults to True):
147
+ Whether to prepend the Cosmos3 transfer system prompt.
148
+ action (`CosmosActionCondition`, *optional*):
149
+ Action-conditioning metadata and its reference visual input.
150
+ fps (`float`, *optional*, defaults to 24.0):
151
+ Frame rate of the generated video.
152
+ add_resolution_template (`bool`, *optional*, defaults to True):
153
+ Whether to add resolution metadata to the prompt.
154
+ add_duration_template (`bool`, *optional*, defaults to True):
155
+ Whether to add duration metadata to the prompt.
156
+
157
+ Outputs:
158
+ height (`int`):
159
+ Resolved output height in pixels.
160
+ width (`int`):
161
+ Resolved output width in pixels.
162
+ control_frames (`dict`):
163
+ Preprocessed, time-padded control maps in canonical hint order.
164
+ total_frames (`int`):
165
+ Total number of output frames to generate.
166
+ chunk_frames (`int`):
167
+ Number of pixel frames per autoregressive chunk.
168
+ num_chunks (`int`):
169
+ Number of autoregressive chunks.
170
+ stride (`int`):
171
+ Frame stride between consecutive chunks.
172
+ cond_input_ids (`Tensor`):
173
+ Token IDs for the conditional prompt.
174
+ uncond_input_ids (`Tensor`):
175
+ Token IDs for the unconditional prompt.
176
+ action_mode (`str`):
177
+ Requested action-generation mode.
178
+ num_frames (`int`):
179
+ Number of frames to generate.
180
+ """
181
+
182
+ model_name = "cosmos3-omni"
183
+ block_classes = [Cosmos3TransferTextBlocks, Cosmos3ActionTextStep, Cosmos3TextEncoderStep]
184
+ block_names = ["transfer_text", "action_text", "text"]
185
+ block_trigger_inputs = ["control_videos", "action", None]
186
+
187
+ @property
188
+ def description(self):
189
+ return (
190
+ "Auto text encoder block for Cosmos3.\n"
191
+ + " - Cosmos3TransferTextBlocks runs when control_videos are provided.\n"
192
+ + " - Cosmos3ActionTextStep runs when action is provided.\n"
193
+ + " - Cosmos3TextEncoderStep runs otherwise."
194
+ )
195
+
196
+
197
+ # auto_docstring
198
+ class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks):
199
+ """
200
+ Auto VAE conditioning block for Cosmos3.
201
+ - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.
202
+ - Cosmos3VideoVaeEncoderStep runs for the non-action video path.
203
+ - Cosmos3ImageVaeEncoderStep runs for the non-action image path.
204
+ - when no action, image, or video conditioning is provided, this block is skipped.
205
+
206
+ Components:
207
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
208
+
209
+ Inputs:
210
+ action (`CosmosActionCondition`, *optional*):
211
+ Action-conditioning metadata and its reference visual input.
212
+ video (`None`, *optional*):
213
+ Reference video for video-to-video conditioning.
214
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
215
+ Latent-frame indexes to preserve from the conditioning video.
216
+ condition_video_keep (`str`, *optional*, defaults to first):
217
+ Which end of a longer conditioning video to use: `first` or `last`.
218
+ num_frames (`int`, *optional*):
219
+ Number of frames to generate.
220
+ height (`int`, *optional*):
221
+ Height of the generated video in pixels.
222
+ width (`int`, *optional*):
223
+ Width of the generated video in pixels.
224
+ image (`None`, *optional*):
225
+ Reference image for image-to-video conditioning.
226
+
227
+ Outputs:
228
+ x0_tokens_vision (`Tensor`):
229
+ Vision latents encoded from the conditioning image or video.
230
+ vision_condition_frames (`list`):
231
+ Latent-frame indexes fixed by visual conditioning.
232
+ action_condition_frame_indexes (`list`):
233
+ Action-frame indexes fixed by action conditioning.
234
+ """
235
+
236
+ model_name = "cosmos3-omni"
237
+ block_classes = [Cosmos3ActionVisionVaeEncoderStep, Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep]
238
+ block_names = ["action_conditioning", "video_conditioning", "image_conditioning"]
239
+ block_trigger_inputs = ["action", "video", "image", "control_videos"]
240
+ default_block_name = None
241
+
242
+ def select_block(self, **kwargs) -> str | None:
243
+ action = kwargs.get("action")
244
+ image = kwargs.get("image")
245
+ video = kwargs.get("video")
246
+ # Transfer preprocesses/encodes its control maps inside the denoise chunk loop, so the standard VAE
247
+ # conditioning stage is skipped when control_videos drive the workflow.
248
+ if kwargs.get("control_videos") is not None:
249
+ return None
250
+ if action is not None:
251
+ if image is not None or video is not None:
252
+ raise ValueError(
253
+ "Pass action conditioning via `action.image` / `action.video`, not top-level image/video."
254
+ )
255
+ return "action_conditioning"
256
+ if image is not None and video is not None:
257
+ raise ValueError("Pass either image or video, not both.")
258
+ if video is not None:
259
+ return "video_conditioning"
260
+ if image is not None:
261
+ return "image_conditioning"
262
+ return None
263
+
264
+ @property
265
+ def description(self):
266
+ return (
267
+ "Auto VAE conditioning block for Cosmos3.\n"
268
+ + " - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.\n"
269
+ + " - Cosmos3VideoVaeEncoderStep runs for the non-action video path.\n"
270
+ + " - Cosmos3ImageVaeEncoderStep runs for the non-action image path.\n"
271
+ + " - when no action, image, or video conditioning is provided, this block is skipped."
272
+ )
273
+
274
+
275
+ # auto_docstring
276
+ class Cosmos3AutoSoundDecodeStep(AutoPipelineBlocks):
277
+ """
278
+ Auto sound decoder block for Cosmos3.
279
+ - Cosmos3SoundDecodeStep runs when sound_latents are present.
280
+ - if sound_latents are not provided, this block is skipped.
281
+
282
+ Components:
283
+ sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
284
+
285
+ Inputs:
286
+ sound_latents (`Tensor`, *optional*):
287
+ Denoised sound latents to decode.
288
+
289
+ Outputs:
290
+ sound (`Tensor`):
291
+ Generated waveform.
292
+ sampling_rate (`int`):
293
+ Sample rate of the generated waveform in Hz.
294
+ """
295
+
296
+ model_name = "cosmos3-omni"
297
+ block_classes = [Cosmos3SoundDecodeStep]
298
+ block_names = ["decode"]
299
+ block_trigger_inputs = ["sound_latents"]
300
+
301
+ @property
302
+ def description(self):
303
+ return (
304
+ "Auto sound decoder block for Cosmos3.\n"
305
+ + " - Cosmos3SoundDecodeStep runs when sound_latents are present.\n"
306
+ + " - if sound_latents are not provided, this block is skipped."
307
+ )
308
+
309
+
310
+ # auto_docstring
311
+ class Cosmos3DecodeStep(SequentialPipelineBlocks):
312
+ """
313
+ Decodes denoised latents into modality outputs.
314
+
315
+ Components:
316
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`)
317
+
318
+ Inputs:
319
+ latents (`Tensor`):
320
+ Denoised vision latents to decode.
321
+ output_type (`str`, *optional*, defaults to pil):
322
+ Output format: 'pil', 'np', 'pt'.
323
+ sound_latents (`Tensor`, *optional*):
324
+ Denoised sound latents to decode.
325
+
326
+ Outputs:
327
+ videos (`list`):
328
+ The generated videos.
329
+ sound (`Tensor`):
330
+ Generated waveform.
331
+ sampling_rate (`int`):
332
+ Sample rate of the generated waveform in Hz.
333
+ """
334
+
335
+ model_name = "cosmos3-omni"
336
+ block_classes = [Cosmos3VideoDecodeStep, Cosmos3AutoSoundDecodeStep]
337
+ block_names = ["video", "sound"]
338
+
339
+ @property
340
+ def description(self) -> str:
341
+ return "Decodes denoised latents into modality outputs."
342
+
343
+
344
+ class Cosmos3AutoDecodeStep(ConditionalPipelineBlocks):
345
+ model_name = "cosmos3-omni"
346
+ block_classes = [Cosmos3TransferStitchStep, Cosmos3DecodeStep]
347
+ block_names = ["transfer", "standard"]
348
+ block_trigger_inputs = ["control_videos"]
349
+ default_block_name = "standard"
350
+
351
+ def select_block(self, **kwargs) -> str | None:
352
+ if kwargs.get("control_videos") is not None:
353
+ return "transfer"
354
+ return "standard"
355
+
356
+ @property
357
+ def description(self) -> str:
358
+ return (
359
+ "Selects the Cosmos3 decode workflow.\n"
360
+ + " - Cosmos3TransferStitchStep stitches the decoded transfer chunks when control_videos are provided.\n"
361
+ + " - Cosmos3DecodeStep decodes the denoised latents otherwise."
362
+ )
363
+
364
+
365
+ # auto_docstring
366
+ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
367
+ """
368
+ Runs the text-and-vision Cosmos3 denoising workflow.
369
+
370
+ Components:
371
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
372
+
373
+ Configs:
374
+ use_native_flow_schedule (default: False)
375
+
376
+ Inputs:
377
+ cond_input_ids (`None`):
378
+ Token IDs for the conditional prompt.
379
+ uncond_input_ids (`None`):
380
+ Token IDs for the unconditional prompt.
381
+ x0_tokens_vision (`Tensor`, *optional*):
382
+ Vision latents encoded from the conditioning image or video.
383
+ vision_condition_frames (`list`, *optional*):
384
+ Latent-frame indexes fixed by visual conditioning.
385
+ num_frames (`int`):
386
+ Number of frames to generate.
387
+ height (`int`):
388
+ Height of the generated video in pixels.
389
+ width (`int`):
390
+ Width of the generated video in pixels.
391
+ fps (`float`, *optional*, defaults to 24.0):
392
+ Frame rate of the generated video.
393
+ latents (`Tensor`, *optional*):
394
+ Pre-generated noisy vision latents.
395
+ generator (`Generator`, *optional*):
396
+ Torch generator for deterministic generation.
397
+ num_inference_steps (`int`):
398
+ The number of denoising steps.
399
+ **denoiser_input_fields (`None`, *optional*):
400
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
401
+ guidance_scale (`float`, *optional*, defaults to 6.0):
402
+ Scale for classifier-free guidance.
403
+
404
+ Outputs:
405
+ latents (`Tensor`):
406
+ Denoised latents.
407
+ """
408
+
409
+ model_name = "cosmos3-omni"
410
+ block_classes = [
411
+ Cosmos3PrepareTextSegmentsStep,
412
+ Cosmos3VisionPrepareLatentsStep,
413
+ Cosmos3VisionPackSequenceStep,
414
+ Cosmos3VisionDenoiseInputStep,
415
+ Cosmos3SetTimestepsStep,
416
+ Cosmos3VisionDenoiseStep,
417
+ ]
418
+ block_names = [
419
+ "prepare_text_segments",
420
+ "prepare_vision_latents",
421
+ "pack_vision_sequence",
422
+ "prepare_vision_denoiser_inputs",
423
+ "set_timesteps",
424
+ "denoise",
425
+ ]
426
+
427
+ @property
428
+ def description(self):
429
+ return "Runs the text-and-vision Cosmos3 denoising workflow."
430
+
431
+ @property
432
+ def outputs(self):
433
+ return [OutputParam.template("latents")]
434
+
435
+
436
+ # auto_docstring
437
+ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
438
+ """
439
+ Runs the text, vision, and sound Cosmos3 denoising workflow.
440
+
441
+ Components:
442
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
443
+
444
+ Configs:
445
+ use_native_flow_schedule (default: False)
446
+
447
+ Inputs:
448
+ cond_input_ids (`None`):
449
+ Token IDs for the conditional prompt.
450
+ uncond_input_ids (`None`):
451
+ Token IDs for the unconditional prompt.
452
+ x0_tokens_vision (`Tensor`, *optional*):
453
+ Vision latents encoded from the conditioning image or video.
454
+ vision_condition_frames (`list`, *optional*):
455
+ Latent-frame indexes fixed by visual conditioning.
456
+ num_frames (`int`):
457
+ Number of frames to generate.
458
+ height (`int`):
459
+ Height of the generated video in pixels.
460
+ width (`int`):
461
+ Width of the generated video in pixels.
462
+ fps (`float`, *optional*, defaults to 24.0):
463
+ Frame rate of the generated video.
464
+ latents (`Tensor`, *optional*):
465
+ Pre-generated noisy vision latents.
466
+ generator (`Generator`, *optional*):
467
+ Torch generator for deterministic generation.
468
+ num_inference_steps (`int`):
469
+ The number of denoising steps.
470
+ sound_latents (`Tensor`, *optional*):
471
+ Pre-generated noisy sound latents.
472
+ **denoiser_input_fields (`None`, *optional*):
473
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
474
+ guidance_scale (`float`, *optional*, defaults to 6.0):
475
+ Scale for classifier-free guidance.
476
+
477
+ Outputs:
478
+ latents (`Tensor`):
479
+ Denoised latents.
480
+ sound_latents (`Tensor`):
481
+ Denoised sound latents.
482
+ """
483
+
484
+ model_name = "cosmos3-omni"
485
+ block_classes = [
486
+ Cosmos3PrepareTextSegmentsStep,
487
+ Cosmos3VisionPrepareLatentsStep,
488
+ Cosmos3VisionPackSequenceStep,
489
+ Cosmos3VisionDenoiseInputStep,
490
+ Cosmos3SetTimestepsStep,
491
+ Cosmos3SoundPrepareLatentsStep,
492
+ Cosmos3SoundPackSequenceStep,
493
+ Cosmos3SoundDenoiseInputStep,
494
+ Cosmos3VisionSoundDenoiseStep,
495
+ ]
496
+ block_names = [
497
+ "prepare_text_segments",
498
+ "prepare_vision_latents",
499
+ "pack_vision_sequence",
500
+ "prepare_vision_denoiser_inputs",
501
+ "set_timesteps",
502
+ "prepare_sound_latents",
503
+ "pack_sound_sequence",
504
+ "prepare_sound_denoiser_inputs",
505
+ "denoise",
506
+ ]
507
+
508
+ @property
509
+ def description(self):
510
+ return "Runs the text, vision, and sound Cosmos3 denoising workflow."
511
+
512
+ @property
513
+ def outputs(self):
514
+ return [
515
+ OutputParam.template("latents"),
516
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."),
517
+ ]
518
+
519
+
520
+ # auto_docstring
521
+ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
522
+ """
523
+ Runs the text, vision, and action Cosmos3 denoising workflow.
524
+
525
+ Components:
526
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
527
+
528
+ Configs:
529
+ use_native_flow_schedule (default: False)
530
+
531
+ Inputs:
532
+ cond_input_ids (`None`):
533
+ Token IDs for the conditional prompt.
534
+ uncond_input_ids (`None`):
535
+ Token IDs for the unconditional prompt.
536
+ x0_tokens_vision (`Tensor`, *optional*):
537
+ Vision latents encoded from the conditioning image or video.
538
+ vision_condition_frames (`list`, *optional*):
539
+ Latent-frame indexes fixed by visual conditioning.
540
+ num_frames (`int`):
541
+ Number of frames to generate.
542
+ height (`int`):
543
+ Height of the generated video in pixels.
544
+ width (`int`):
545
+ Width of the generated video in pixels.
546
+ fps (`float`, *optional*, defaults to 24.0):
547
+ Frame rate of the generated video.
548
+ latents (`Tensor`, *optional*):
549
+ Pre-generated noisy vision latents.
550
+ generator (`Generator`, *optional*):
551
+ Torch generator for deterministic generation.
552
+ num_inference_steps (`int`):
553
+ The number of denoising steps.
554
+ action (`CosmosActionCondition`):
555
+ Action-conditioning metadata.
556
+ action_condition_frame_indexes (`list`, *optional*):
557
+ Action-frame indexes fixed by action conditioning.
558
+ action_latents (`Tensor`, *optional*):
559
+ Pre-generated noisy action latents.
560
+ **denoiser_input_fields (`None`, *optional*):
561
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
562
+ guidance_scale (`float`, *optional*, defaults to 6.0):
563
+ Scale for classifier-free guidance.
564
+
565
+ Outputs:
566
+ latents (`Tensor`):
567
+ Denoised latents.
568
+ action_latents (`Tensor`):
569
+ Denoised action latents.
570
+ """
571
+
572
+ model_name = "cosmos3-omni"
573
+ block_classes = [
574
+ Cosmos3PrepareTextSegmentsStep,
575
+ Cosmos3VisionPrepareLatentsStep,
576
+ Cosmos3VisionPackSequenceStep,
577
+ Cosmos3VisionDenoiseInputStep,
578
+ Cosmos3SetTimestepsStep,
579
+ Cosmos3ActionPrepareLatentsStep,
580
+ Cosmos3ActionPackSequenceStep,
581
+ Cosmos3ActionDenoiseInputStep,
582
+ Cosmos3VisionActionDenoiseStep,
583
+ ]
584
+ block_names = [
585
+ "prepare_text_segments",
586
+ "prepare_vision_latents",
587
+ "pack_vision_sequence",
588
+ "prepare_vision_denoiser_inputs",
589
+ "set_timesteps",
590
+ "prepare_action_latents",
591
+ "pack_action_sequence",
592
+ "prepare_action_denoiser_inputs",
593
+ "denoise",
594
+ ]
595
+
596
+ @property
597
+ def description(self):
598
+ return "Runs the text, vision, and action Cosmos3 denoising workflow."
599
+
600
+ @property
601
+ def outputs(self):
602
+ return [
603
+ OutputParam.template("latents"),
604
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."),
605
+ ]
606
+
607
+
608
+ # auto_docstring
609
+ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
610
+ """
611
+ Runs the text, vision, sound, and action Cosmos3 denoising workflow.
612
+
613
+ Components:
614
+ transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
615
+
616
+ Configs:
617
+ use_native_flow_schedule (default: False)
618
+
619
+ Inputs:
620
+ cond_input_ids (`None`):
621
+ Token IDs for the conditional prompt.
622
+ uncond_input_ids (`None`):
623
+ Token IDs for the unconditional prompt.
624
+ x0_tokens_vision (`Tensor`, *optional*):
625
+ Vision latents encoded from the conditioning image or video.
626
+ vision_condition_frames (`list`, *optional*):
627
+ Latent-frame indexes fixed by visual conditioning.
628
+ num_frames (`int`):
629
+ Number of frames to generate.
630
+ height (`int`):
631
+ Height of the generated video in pixels.
632
+ width (`int`):
633
+ Width of the generated video in pixels.
634
+ fps (`float`, *optional*, defaults to 24.0):
635
+ Frame rate of the generated video.
636
+ latents (`Tensor`, *optional*):
637
+ Pre-generated noisy vision latents.
638
+ generator (`Generator`, *optional*):
639
+ Torch generator for deterministic generation.
640
+ num_inference_steps (`int`):
641
+ The number of denoising steps.
642
+ sound_latents (`Tensor`, *optional*):
643
+ Pre-generated noisy sound latents.
644
+ action (`CosmosActionCondition`):
645
+ Action-conditioning metadata.
646
+ action_condition_frame_indexes (`list`, *optional*):
647
+ Action-frame indexes fixed by action conditioning.
648
+ action_latents (`Tensor`, *optional*):
649
+ Pre-generated noisy action latents.
650
+ **denoiser_input_fields (`None`, *optional*):
651
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
652
+ guidance_scale (`float`, *optional*, defaults to 6.0):
653
+ Scale for classifier-free guidance.
654
+
655
+ Outputs:
656
+ latents (`Tensor`):
657
+ Denoised latents.
658
+ sound_latents (`Tensor`):
659
+ Denoised sound latents.
660
+ action_latents (`Tensor`):
661
+ Denoised action latents.
662
+ """
663
+
664
+ model_name = "cosmos3-omni"
665
+ block_classes = [
666
+ Cosmos3PrepareTextSegmentsStep,
667
+ Cosmos3VisionPrepareLatentsStep,
668
+ Cosmos3VisionPackSequenceStep,
669
+ Cosmos3VisionDenoiseInputStep,
670
+ Cosmos3SetTimestepsStep,
671
+ Cosmos3SoundPrepareLatentsStep,
672
+ Cosmos3SoundPackSequenceStep,
673
+ Cosmos3SoundDenoiseInputStep,
674
+ Cosmos3ActionPrepareLatentsStep,
675
+ Cosmos3ActionPackSequenceStep,
676
+ Cosmos3ActionDenoiseInputStep,
677
+ Cosmos3VisionSoundActionDenoiseStep,
678
+ ]
679
+ block_names = [
680
+ "prepare_text_segments",
681
+ "prepare_vision_latents",
682
+ "pack_vision_sequence",
683
+ "prepare_vision_denoiser_inputs",
684
+ "set_timesteps",
685
+ "prepare_sound_latents",
686
+ "pack_sound_sequence",
687
+ "prepare_sound_denoiser_inputs",
688
+ "prepare_action_latents",
689
+ "pack_action_sequence",
690
+ "prepare_action_denoiser_inputs",
691
+ "denoise",
692
+ ]
693
+
694
+ @property
695
+ def description(self):
696
+ return "Runs the text, vision, sound, and action Cosmos3 denoising workflow."
697
+
698
+ @property
699
+ def outputs(self):
700
+ return [
701
+ OutputParam.template("latents"),
702
+ OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."),
703
+ OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."),
704
+ ]
705
+
706
+
707
+ # auto_docstring
708
+ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
709
+ """
710
+ Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf
711
+ LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk
712
+ cross-carry (previous_output, output_chunks) lives on PipelineState.
713
+
714
+ Components:
715
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`Cosmos3OmniTransformer`) scheduler
716
+ (`UniPCMultistepScheduler`)
717
+
718
+ Inputs:
719
+ chunk_id (`int`, *optional*, defaults to 0):
720
+ Index of the current chunk.
721
+ previous_output (`None`, *optional*):
722
+ Decoded pixels of the previous chunk, used to seed later chunks.
723
+ control_frames (`dict`):
724
+ Preprocessed, time-padded control maps in canonical hint order.
725
+ chunk_frames (`int`):
726
+ Pixel frames per chunk.
727
+ total_frames (`int`):
728
+ Total number of output frames.
729
+ stride (`int`):
730
+ Frame stride between chunks.
731
+ height (`int`):
732
+ Height of the generated video in pixels.
733
+ width (`int`):
734
+ Width of the generated video in pixels.
735
+ video (`None`, *optional*):
736
+ Optional input video that seeds the first chunk's conditioning.
737
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
738
+ Number of frames the first chunk reuses from the input video.
739
+ num_conditional_frames (`int`, *optional*, defaults to 1):
740
+ Number of frames each later chunk reuses from the previous chunk's tail.
741
+ generator (`Generator`, *optional*):
742
+ Torch generator for deterministic generation.
743
+ cond_text_segment (`dict`):
744
+ Conditional text segment.
745
+ uncond_text_segment (`dict`):
746
+ Unconditional text segment.
747
+ fps (`float`, *optional*, defaults to 24.0):
748
+ Frame rate of the generated video.
749
+ num_inference_steps (`int`):
750
+ The number of denoising steps.
751
+ **denoiser_input_fields (`None`, *optional*):
752
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
753
+ guidance_scale (`float`, *optional*, defaults to 6.0):
754
+ Scale for text classifier-free guidance.
755
+ control_guidance (`float`, *optional*, defaults to 1.0):
756
+ Scale for the control (structural) guidance axis.
757
+ guidance_interval (`tuple`, *optional*):
758
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
759
+ control_guidance_interval (`tuple`, *optional*):
760
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
761
+ output_chunks (`list`):
762
+ Decoded pixel chunks accumulated so far.
763
+ num_chunks (`int`):
764
+ Number of autoregressive chunks.
765
+
766
+ Outputs:
767
+ control_latents (`list`):
768
+ Clean control latents for this chunk, one per hint in canonical order.
769
+ x0_tokens_vision (`Tensor`):
770
+ Clean target vision latents encoded from the seeded target frames.
771
+ current_conditional_frames (`int`):
772
+ Number of pixel frames actually used to seed this chunk's target.
773
+ latents (`Tensor`):
774
+ Noisy target latents for this chunk.
775
+ velocity_mask (`Tensor`):
776
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
777
+ condition_latents (`Tensor`):
778
+ Clean target latents on the conditioned frames (the autoregressive seed).
779
+ target_condition_indexes (`list`):
780
+ Latent-frame indexes fixed by the chunk's conditioning.
781
+ cond_full_static (`dict`):
782
+ Conditional [control..., target] transfer sequence carrying every control item.
783
+ cond_no_control_static (`dict`):
784
+ Conditional [target] transfer sequence with the control items dropped.
785
+ uncond_full_static (`dict`):
786
+ Unconditional [control..., target] transfer sequence for text CFG.
787
+ num_noisy_vision_tokens (`int`):
788
+ Number of noisy target vision tokens denoised each step.
789
+ timesteps (`Tensor`):
790
+ Scheduler timesteps for this chunk.
791
+ num_warmup_steps (`int`):
792
+ Number of scheduler warmup steps for this chunk.
793
+ vision_tokens_full (`list`):
794
+ Token list for the [control..., target] forward passes.
795
+ vision_tokens_target (`list`):
796
+ Token list for the target-only (no-control) forward pass.
797
+ vision_timesteps (`Tensor`):
798
+ Timesteps for the noisy target tokens.
799
+ velocity (`Tensor`):
800
+ Predicted (masked) transfer velocity.
801
+ previous_output (`Tensor`):
802
+ Decoded pixels of this chunk, used to seed the next chunk.
803
+ output_chunks (`list`):
804
+ Decoded pixel chunks accumulated so far (with this chunk appended).
805
+ """
806
+
807
+ model_name = "cosmos3-omni"
808
+ block_classes = [
809
+ Cosmos3TransferChunkVaeEncoderStep,
810
+ Cosmos3TransferPrepareLatentsStep,
811
+ Cosmos3TransferPackSequenceStep,
812
+ Cosmos3TransferSetTimestepsStep,
813
+ Cosmos3TransferDenoiseStep,
814
+ Cosmos3TransferDecodeChunkStep,
815
+ ]
816
+ block_names = [
817
+ "encode_transfer_chunk",
818
+ "prepare_transfer_latents",
819
+ "pack_transfer_sequence",
820
+ "set_timesteps",
821
+ "denoise",
822
+ "decode_chunk",
823
+ ]
824
+
825
+ @property
826
+ def description(self) -> str:
827
+ return (
828
+ "Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a "
829
+ "non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). "
830
+ "Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState."
831
+ )
832
+
833
+ @property
834
+ def inputs(self) -> list[InputParam]:
835
+ return super().inputs + [
836
+ InputParam(name="num_chunks", type_hint=int, required=True, description="Number of autoregressive chunks.")
837
+ ]
838
+
839
+ @torch.no_grad()
840
+ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
841
+ num_chunks = state.get("num_chunks")
842
+ state.set("output_chunks", [])
843
+ state.set("previous_output", None)
844
+ for chunk_id in range(num_chunks):
845
+ state.set("chunk_id", chunk_id)
846
+ for _, block in self.sub_blocks.items():
847
+ components, state = block(components, state)
848
+ return components, state
849
+
850
+
851
+ # auto_docstring
852
+ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
853
+ """
854
+ Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop.
855
+
856
+ Components:
857
+ transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
858
+ (`UniPCMultistepScheduler`)
859
+
860
+ Inputs:
861
+ cond_input_ids (`None`):
862
+ Token IDs for the conditional prompt.
863
+ uncond_input_ids (`None`):
864
+ Token IDs for the unconditional prompt.
865
+ chunk_id (`int`, *optional*, defaults to 0):
866
+ Index of the current chunk.
867
+ previous_output (`None`, *optional*):
868
+ Decoded pixels of the previous chunk, used to seed later chunks.
869
+ control_frames (`dict`):
870
+ Preprocessed, time-padded control maps in canonical hint order.
871
+ chunk_frames (`int`):
872
+ Pixel frames per chunk.
873
+ total_frames (`int`):
874
+ Total number of output frames.
875
+ stride (`int`):
876
+ Frame stride between chunks.
877
+ height (`int`):
878
+ Height of the generated video in pixels.
879
+ width (`int`):
880
+ Width of the generated video in pixels.
881
+ video (`None`, *optional*):
882
+ Optional input video that seeds the first chunk's conditioning.
883
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
884
+ Number of frames the first chunk reuses from the input video.
885
+ num_conditional_frames (`int`, *optional*, defaults to 1):
886
+ Number of frames each later chunk reuses from the previous chunk's tail.
887
+ generator (`Generator`, *optional*):
888
+ Torch generator for deterministic generation.
889
+ fps (`float`, *optional*, defaults to 24.0):
890
+ Frame rate of the generated video.
891
+ num_inference_steps (`int`):
892
+ The number of denoising steps.
893
+ **denoiser_input_fields (`None`, *optional*):
894
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
895
+ guidance_scale (`float`, *optional*, defaults to 6.0):
896
+ Scale for text classifier-free guidance.
897
+ control_guidance (`float`, *optional*, defaults to 1.0):
898
+ Scale for the control (structural) guidance axis.
899
+ guidance_interval (`tuple`, *optional*):
900
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
901
+ control_guidance_interval (`tuple`, *optional*):
902
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
903
+ output_chunks (`list`):
904
+ Decoded pixel chunks accumulated so far.
905
+ num_chunks (`int`):
906
+ Number of autoregressive chunks.
907
+
908
+ Outputs:
909
+ cond_text_segment (`dict`):
910
+ Conditional text segment for the denoiser.
911
+ uncond_text_segment (`dict`):
912
+ Unconditional text segment for the denoiser.
913
+ control_latents (`list`):
914
+ Clean control latents for this chunk, one per hint in canonical order.
915
+ x0_tokens_vision (`Tensor`):
916
+ Clean target vision latents encoded from the seeded target frames.
917
+ current_conditional_frames (`int`):
918
+ Number of pixel frames actually used to seed this chunk's target.
919
+ latents (`Tensor`):
920
+ Noisy target latents for this chunk.
921
+ velocity_mask (`Tensor`):
922
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
923
+ condition_latents (`Tensor`):
924
+ Clean target latents on the conditioned frames (the autoregressive seed).
925
+ target_condition_indexes (`list`):
926
+ Latent-frame indexes fixed by the chunk's conditioning.
927
+ cond_full_static (`dict`):
928
+ Conditional [control..., target] transfer sequence carrying every control item.
929
+ cond_no_control_static (`dict`):
930
+ Conditional [target] transfer sequence with the control items dropped.
931
+ uncond_full_static (`dict`):
932
+ Unconditional [control..., target] transfer sequence for text CFG.
933
+ num_noisy_vision_tokens (`int`):
934
+ Number of noisy target vision tokens denoised each step.
935
+ timesteps (`Tensor`):
936
+ Scheduler timesteps for this chunk.
937
+ num_warmup_steps (`int`):
938
+ Number of scheduler warmup steps for this chunk.
939
+ vision_tokens_full (`list`):
940
+ Token list for the [control..., target] forward passes.
941
+ vision_tokens_target (`list`):
942
+ Token list for the target-only (no-control) forward pass.
943
+ vision_timesteps (`Tensor`):
944
+ Timesteps for the noisy target tokens.
945
+ velocity (`Tensor`):
946
+ Predicted (masked) transfer velocity.
947
+ previous_output (`Tensor`):
948
+ Decoded pixels of this chunk, used to seed the next chunk.
949
+ output_chunks (`list`):
950
+ Decoded pixel chunks accumulated so far (with this chunk appended).
951
+ """
952
+
953
+ model_name = "cosmos3-omni"
954
+ block_classes = [
955
+ Cosmos3PrepareTextSegmentsStep,
956
+ Cosmos3TransferChunkDenoiseStep,
957
+ ]
958
+ block_names = ["prepare_text_segments", "chunk_denoise"]
959
+
960
+ @property
961
+ def description(self) -> str:
962
+ return "Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop."
963
+
964
+
965
+ # auto_docstring
966
+ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
967
+ """
968
+ Selects the Cosmos3 core denoising workflow.
969
+ - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.
970
+ - vision_sound_action runs when action and enable_sound are provided.
971
+ - vision_action runs when action is provided.
972
+ - vision_sound runs when enable_sound is true.
973
+ - vision runs otherwise.
974
+
975
+ Components:
976
+ transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler
977
+ (`UniPCMultistepScheduler`)
978
+
979
+ Configs:
980
+ use_native_flow_schedule (default: False)
981
+
982
+ Inputs:
983
+ cond_input_ids (`None`):
984
+ Token IDs for the conditional prompt.
985
+ uncond_input_ids (`None`):
986
+ Token IDs for the unconditional prompt.
987
+ chunk_id (`int`, *optional*, defaults to 0):
988
+ Index of the current chunk.
989
+ previous_output (`None`, *optional*):
990
+ Decoded pixels of the previous chunk, used to seed later chunks.
991
+ control_frames (`dict`, *optional*):
992
+ Preprocessed, time-padded control maps in canonical hint order.
993
+ chunk_frames (`int`, *optional*):
994
+ Pixel frames per chunk.
995
+ total_frames (`int`, *optional*):
996
+ Total number of output frames.
997
+ stride (`int`, *optional*):
998
+ Frame stride between chunks.
999
+ height (`int`):
1000
+ Height of the generated video in pixels.
1001
+ width (`int`):
1002
+ Width of the generated video in pixels.
1003
+ video (`None`, *optional*):
1004
+ Optional input video that seeds the first chunk's conditioning.
1005
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
1006
+ Number of frames the first chunk reuses from the input video.
1007
+ num_conditional_frames (`int`, *optional*, defaults to 1):
1008
+ Number of frames each later chunk reuses from the previous chunk's tail.
1009
+ generator (`Generator`, *optional*):
1010
+ Torch generator for deterministic generation.
1011
+ fps (`float`, *optional*, defaults to 24.0):
1012
+ Frame rate of the generated video.
1013
+ num_inference_steps (`int`):
1014
+ The number of denoising steps.
1015
+ **denoiser_input_fields (`None`, *optional*):
1016
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
1017
+ guidance_scale (`float`, *optional*, defaults to 6.0):
1018
+ Scale for text classifier-free guidance.
1019
+ control_guidance (`float`, *optional*, defaults to 1.0):
1020
+ Scale for the control (structural) guidance axis.
1021
+ guidance_interval (`tuple`, *optional*):
1022
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
1023
+ control_guidance_interval (`tuple`, *optional*):
1024
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
1025
+ output_chunks (`list`, *optional*):
1026
+ Decoded pixel chunks accumulated so far.
1027
+ num_chunks (`int`, *optional*):
1028
+ Number of autoregressive chunks.
1029
+ x0_tokens_vision (`Tensor`, *optional*):
1030
+ Vision latents encoded from the conditioning image or video.
1031
+ vision_condition_frames (`list`, *optional*):
1032
+ Latent-frame indexes fixed by visual conditioning.
1033
+ num_frames (`int`, *optional*):
1034
+ Number of frames to generate.
1035
+ latents (`Tensor`):
1036
+ Pre-generated noisy vision latents.
1037
+ sound_latents (`Tensor`, *optional*):
1038
+ Pre-generated noisy sound latents.
1039
+ action (`CosmosActionCondition`, *optional*):
1040
+ Action-conditioning metadata.
1041
+ action_condition_frame_indexes (`list`, *optional*):
1042
+ Action-frame indexes fixed by action conditioning.
1043
+ action_latents (`Tensor`, *optional*):
1044
+ Pre-generated noisy action latents.
1045
+ enable_sound (`bool`, *optional*, defaults to False):
1046
+ Whether to generate a synchronized sound track.
1047
+
1048
+ Outputs:
1049
+ cond_text_segment (`dict`):
1050
+ Conditional text segment for the denoiser.
1051
+ uncond_text_segment (`dict`):
1052
+ Unconditional text segment for the denoiser.
1053
+ control_latents (`list`):
1054
+ Clean control latents for this chunk, one per hint in canonical order.
1055
+ x0_tokens_vision (`Tensor`):
1056
+ Clean target vision latents encoded from the seeded target frames.
1057
+ current_conditional_frames (`int`):
1058
+ Number of pixel frames actually used to seed this chunk's target.
1059
+ latents (`Tensor`):
1060
+ Noisy target latents for this chunk.
1061
+ velocity_mask (`Tensor`):
1062
+ Mask that zeroes the velocity on conditioned (clean) latent frames.
1063
+ condition_latents (`Tensor`):
1064
+ Clean target latents on the conditioned frames (the autoregressive seed).
1065
+ target_condition_indexes (`list`):
1066
+ Latent-frame indexes fixed by the chunk's conditioning.
1067
+ cond_full_static (`dict`):
1068
+ Conditional [control..., target] transfer sequence carrying every control item.
1069
+ cond_no_control_static (`dict`):
1070
+ Conditional [target] transfer sequence with the control items dropped.
1071
+ uncond_full_static (`dict`):
1072
+ Unconditional [control..., target] transfer sequence for text CFG.
1073
+ num_noisy_vision_tokens (`int`):
1074
+ Number of noisy target vision tokens denoised each step.
1075
+ timesteps (`Tensor`):
1076
+ Scheduler timesteps for this chunk.
1077
+ num_warmup_steps (`int`):
1078
+ Number of scheduler warmup steps for this chunk.
1079
+ vision_tokens_full (`list`):
1080
+ Token list for the [control..., target] forward passes.
1081
+ vision_tokens_target (`list`):
1082
+ Token list for the target-only (no-control) forward pass.
1083
+ vision_timesteps (`Tensor`):
1084
+ Timesteps for the noisy target tokens.
1085
+ velocity (`Tensor`):
1086
+ Predicted (masked) transfer velocity.
1087
+ previous_output (`Tensor`):
1088
+ Decoded pixels of this chunk, used to seed the next chunk.
1089
+ output_chunks (`list`):
1090
+ Decoded pixel chunks accumulated so far (with this chunk appended).
1091
+ sound_latents (`Tensor`):
1092
+ Denoised sound latents.
1093
+ action_latents (`Tensor`):
1094
+ Denoised action latents.
1095
+ """
1096
+
1097
+ model_name = "cosmos3-omni"
1098
+ block_classes = [
1099
+ Cosmos3TransferCoreDenoiseStep,
1100
+ Cosmos3VisionSoundActionCoreDenoiseStep,
1101
+ Cosmos3VisionActionCoreDenoiseStep,
1102
+ Cosmos3VisionSoundCoreDenoiseStep,
1103
+ Cosmos3VisionCoreDenoiseStep,
1104
+ ]
1105
+ block_names = ["transfer", "vision_sound_action", "vision_action", "vision_sound", "vision"]
1106
+ block_trigger_inputs = ["action", "enable_sound", "control_videos"]
1107
+ default_block_name = "vision"
1108
+
1109
+ @property
1110
+ def inputs(self):
1111
+ inputs = super().inputs
1112
+ inputs.append(
1113
+ InputParam(
1114
+ name="enable_sound",
1115
+ type_hint=bool,
1116
+ default=False,
1117
+ description="Whether to generate a synchronized sound track.",
1118
+ )
1119
+ )
1120
+ return inputs
1121
+
1122
+ def select_block(self, **kwargs) -> str | None:
1123
+ action = kwargs.get("action")
1124
+ enable_sound = kwargs.get("enable_sound")
1125
+ if kwargs.get("control_videos") is not None:
1126
+ return "transfer"
1127
+ if action is not None and enable_sound:
1128
+ return "vision_sound_action"
1129
+ if action is not None:
1130
+ return "vision_action"
1131
+ if enable_sound:
1132
+ return "vision_sound"
1133
+ return "vision"
1134
+
1135
+ @property
1136
+ def description(self):
1137
+ return (
1138
+ "Selects the Cosmos3 core denoising workflow.\n"
1139
+ + " - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.\n"
1140
+ + " - vision_sound_action runs when action and enable_sound are provided.\n"
1141
+ + " - vision_action runs when action is provided.\n"
1142
+ + " - vision_sound runs when enable_sound is true.\n"
1143
+ + " - vision runs otherwise."
1144
+ )
1145
+
1146
+
1147
+ # auto_docstring
1148
+ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
1149
+ """
1150
+ Modular pipeline blocks for Cosmos3 generation modes.
1151
+
1152
+ Supported workflows:
1153
+ - `text2image`: requires `prompt`, `num_frames`
1154
+ - `text2video`: requires `prompt`
1155
+ - `image2video`: requires `prompt`, `image`
1156
+ - `video2video`: requires `prompt`, `video`
1157
+ - `text2video_with_sound`: requires `prompt`, `enable_sound`
1158
+ - `image2video_with_sound`: requires `prompt`, `image`, `enable_sound`
1159
+ - `video2video_with_sound`: requires `prompt`, `video`, `enable_sound`
1160
+ - `action_policy`: requires `prompt`, `action`
1161
+ - `action_forward_dynamics`: requires `prompt`, `action`
1162
+ - `action_inverse_dynamics`: requires `prompt`, `action`
1163
+
1164
+ Components:
1165
+ video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer
1166
+ (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer
1167
+ (`Cosmos3AVAEAudioTokenizer`)
1168
+
1169
+ Configs:
1170
+ default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule
1171
+ (default: False)
1172
+
1173
+ Inputs:
1174
+ control_videos (`dict`, *optional*):
1175
+ Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.
1176
+ height (`int`, *optional*):
1177
+ Height of the generated video in pixels.
1178
+ width (`int`, *optional*):
1179
+ Width of the generated video in pixels.
1180
+ num_frames (`int`, *optional*):
1181
+ Optional cap on the number of output frames (defaults to the control video length).
1182
+ num_video_frames_per_chunk (`int`, *optional*):
1183
+ Number of pixel frames generated per autoregressive chunk.
1184
+ num_conditional_frames (`int`, *optional*, defaults to 1):
1185
+ Number of frames each chunk reuses from the previous chunk's tail.
1186
+ prompt (`str`):
1187
+ The text prompt that guides Cosmos3 generation.
1188
+ negative_prompt (`str`, *optional*):
1189
+ The negative text prompt used for classifier-free guidance.
1190
+ use_system_prompt (`bool`, *optional*, defaults to True):
1191
+ Whether to prepend the Cosmos3 transfer system prompt.
1192
+ action (`CosmosActionCondition`, *optional*):
1193
+ Action-conditioning metadata and its reference visual input.
1194
+ fps (`float`, *optional*, defaults to 24.0):
1195
+ Frame rate of the generated video.
1196
+ add_resolution_template (`bool`, *optional*, defaults to True):
1197
+ Whether to add resolution metadata to the prompt.
1198
+ add_duration_template (`bool`, *optional*, defaults to True):
1199
+ Whether to add duration metadata to the prompt.
1200
+ video (`None`, *optional*):
1201
+ Reference video for video-to-video conditioning.
1202
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
1203
+ Latent-frame indexes to preserve from the conditioning video.
1204
+ condition_video_keep (`str`, *optional*, defaults to first):
1205
+ Which end of a longer conditioning video to use: `first` or `last`.
1206
+ image (`None`, *optional*):
1207
+ Reference image for image-to-video conditioning.
1208
+ chunk_id (`int`, *optional*, defaults to 0):
1209
+ Index of the current chunk.
1210
+ previous_output (`None`, *optional*):
1211
+ Decoded pixels of the previous chunk, used to seed later chunks.
1212
+ num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0):
1213
+ Number of frames the first chunk reuses from the input video.
1214
+ generator (`Generator`, *optional*):
1215
+ Torch generator for deterministic generation.
1216
+ num_inference_steps (`int`):
1217
+ The number of denoising steps.
1218
+ **denoiser_input_fields (`None`, *optional*):
1219
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
1220
+ guidance_scale (`float`, *optional*, defaults to 6.0):
1221
+ Scale for text classifier-free guidance.
1222
+ control_guidance (`float`, *optional*, defaults to 1.0):
1223
+ Scale for the control (structural) guidance axis.
1224
+ guidance_interval (`tuple`, *optional*):
1225
+ Timestep interval [lo, hi] over which text guidance is active (None = always).
1226
+ control_guidance_interval (`tuple`, *optional*):
1227
+ Timestep interval [lo, hi] over which control guidance is active (None = always).
1228
+ output_chunks (`list`, *optional*):
1229
+ Decoded pixel chunks accumulated so far.
1230
+ x0_tokens_vision (`Tensor`, *optional*):
1231
+ Vision latents encoded from the conditioning image or video.
1232
+ vision_condition_frames (`list`, *optional*):
1233
+ Latent-frame indexes fixed by visual conditioning.
1234
+ latents (`Tensor`):
1235
+ Pre-generated noisy vision latents.
1236
+ sound_latents (`Tensor`, *optional*):
1237
+ Pre-generated noisy sound latents.
1238
+ action_condition_frame_indexes (`list`, *optional*):
1239
+ Action-frame indexes fixed by action conditioning.
1240
+ action_latents (`Tensor`, *optional*):
1241
+ Pre-generated noisy action latents.
1242
+ enable_sound (`bool`, *optional*, defaults to False):
1243
+ Whether to generate a synchronized sound track.
1244
+ output_type (`str`, *optional*, defaults to pil):
1245
+ Output format: 'pil', 'np', 'pt'.
1246
+
1247
+ Outputs:
1248
+ videos (`list`):
1249
+ The generated videos.
1250
+ sound (`Tensor`):
1251
+ Generated waveform.
1252
+ sampling_rate (`int`):
1253
+ Sample rate of the generated waveform in Hz.
1254
+ action (`list`):
1255
+ Generated action vectors.
1256
+ """
1257
+
1258
+ model_name = "cosmos3-omni"
1259
+ block_classes = [
1260
+ Cosmos3AutoTextEncoderStep,
1261
+ Cosmos3AutoVaeEncoderStep,
1262
+ Cosmos3AutoCoreDenoiseStep,
1263
+ Cosmos3AutoDecodeStep,
1264
+ Cosmos3ActionOutputStep,
1265
+ ]
1266
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"]
1267
+ _workflow_map = {
1268
+ "text2image": {"prompt": True, "num_frames": 1},
1269
+ "text2video": {"prompt": True},
1270
+ "image2video": {"prompt": True, "image": True},
1271
+ "video2video": {"prompt": True, "video": True},
1272
+ "text2video_with_sound": {"prompt": True, "enable_sound": True},
1273
+ "image2video_with_sound": {"prompt": True, "image": True, "enable_sound": True},
1274
+ "video2video_with_sound": {"prompt": True, "video": True, "enable_sound": True},
1275
+ "action_policy": {"prompt": True, "action": True},
1276
+ "action_forward_dynamics": {"prompt": True, "action": True},
1277
+ "action_inverse_dynamics": {"prompt": True, "action": True},
1278
+ }
1279
+
1280
+ @property
1281
+ def description(self):
1282
+ return "Modular pipeline blocks for Cosmos3 generation modes."
1283
+
1284
+ def get_workflow(self, workflow_name: str):
1285
+ if workflow_name == "transfer":
1286
+ raise NotImplementedError(
1287
+ 'The standalone "transfer" workflow is temporarily unavailable because its nested autoregressive '
1288
+ "chunk and denoising loops cannot be preserved by the current workflow extraction logic. Transfer "
1289
+ "remains available through the full Cosmos3OmniBlocks pipeline. The standalone workflow will be "
1290
+ "enabled after migration to the upcoming composable nested-loop abstraction."
1291
+ )
1292
+ return super().get_workflow(workflow_name)
1293
+
1294
+ @property
1295
+ def outputs(self):
1296
+ return [
1297
+ OutputParam.template("videos"),
1298
+ OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."),
1299
+ OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."),
1300
+ OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors."),
1301
+ ]
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
2
+ from ..modular_pipeline_utils import OutputParam
3
+ from .before_denoise import (
4
+ Cosmos3DistilledSetTimestepsStep,
5
+ Cosmos3PrepareTextSegmentsStep,
6
+ Cosmos3VisionDenoiseInputStep,
7
+ Cosmos3VisionPackSequenceStep,
8
+ Cosmos3VisionPrepareLatentsStep,
9
+ )
10
+ from .decoders import Cosmos3VideoDecodeStep
11
+ from .denoise import Cosmos3DistilledVisionDenoiseStep
12
+ from .encoders import (
13
+ Cosmos3DistilledTextEncoderStep,
14
+ Cosmos3ImageVaeEncoderStep,
15
+ Cosmos3VideoVaeEncoderStep,
16
+ )
17
+
18
+
19
+ # auto_docstring
20
+ class Cosmos3DistilledAutoVaeEncoderStep(ConditionalPipelineBlocks):
21
+ """
22
+ Auto VAE conditioning block for distilled Cosmos3.
23
+ - Cosmos3VideoVaeEncoderStep runs for the video path.
24
+ - Cosmos3ImageVaeEncoderStep runs for the image path.
25
+ - when no image or video conditioning is provided, this block is skipped.
26
+
27
+ Components:
28
+ vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`)
29
+
30
+ Inputs:
31
+ video (`None`, *optional*):
32
+ Reference video for video-to-video conditioning.
33
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
34
+ Latent-frame indexes to preserve from the conditioning video.
35
+ condition_video_keep (`str`, *optional*, defaults to first):
36
+ Which end of a longer conditioning video to use: `first` or `last`.
37
+ num_frames (`int`, *optional*):
38
+ Number of frames to generate.
39
+ height (`int`, *optional*):
40
+ Height of the generated video in pixels.
41
+ width (`int`, *optional*):
42
+ Width of the generated video in pixels.
43
+ image (`None`, *optional*):
44
+ Reference image for image-to-video conditioning.
45
+
46
+ Outputs:
47
+ x0_tokens_vision (`Tensor`):
48
+ Vision latents encoded from the conditioning image or video.
49
+ vision_condition_frames (`list`):
50
+ Latent-frame indexes fixed by visual conditioning.
51
+ """
52
+
53
+ model_name = "cosmos3-omni"
54
+ block_classes = [Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep]
55
+ block_names = ["video_conditioning", "image_conditioning"]
56
+ block_trigger_inputs = ["video", "image"]
57
+ default_block_name = None
58
+
59
+ def select_block(self, **kwargs) -> str | None:
60
+ image = kwargs.get("image")
61
+ video = kwargs.get("video")
62
+ if image is not None and video is not None:
63
+ raise ValueError("Pass either image or video, not both.")
64
+ if video is not None:
65
+ return "video_conditioning"
66
+ if image is not None:
67
+ return "image_conditioning"
68
+ return None
69
+
70
+ @property
71
+ def description(self):
72
+ return (
73
+ "Auto VAE conditioning block for distilled Cosmos3.\n"
74
+ + " - Cosmos3VideoVaeEncoderStep runs for the video path.\n"
75
+ + " - Cosmos3ImageVaeEncoderStep runs for the image path.\n"
76
+ + " - when no image or video conditioning is provided, this block is skipped."
77
+ )
78
+
79
+
80
+ # auto_docstring
81
+ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
82
+ """
83
+ Runs the text-and-vision distilled Cosmos3 denoising workflow.
84
+
85
+ Components:
86
+ transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
87
+
88
+ Configs:
89
+ is_distilled (default: True) distilled_sigmas (default: None)
90
+
91
+ Inputs:
92
+ cond_input_ids (`None`):
93
+ Token IDs for the conditional prompt.
94
+ uncond_input_ids (`None`):
95
+ Token IDs for the unconditional prompt.
96
+ x0_tokens_vision (`Tensor`, *optional*):
97
+ Vision latents encoded from the conditioning image or video.
98
+ vision_condition_frames (`list`, *optional*):
99
+ Latent-frame indexes fixed by visual conditioning.
100
+ num_frames (`int`):
101
+ Number of frames to generate.
102
+ height (`int`):
103
+ Height of the generated video in pixels.
104
+ width (`int`):
105
+ Width of the generated video in pixels.
106
+ fps (`float`, *optional*, defaults to 24.0):
107
+ Frame rate of the generated video.
108
+ latents (`Tensor`, *optional*):
109
+ Pre-generated noisy vision latents.
110
+ generator (`Generator`, *optional*):
111
+ Torch generator for deterministic generation.
112
+ num_inference_steps (`int`, *optional*):
113
+ The number of denoising steps.
114
+ guidance_scale (`float`, *optional*):
115
+ Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
116
+ forced to 1.0. Passing a value other than 1.0 raises an error.
117
+ **denoiser_input_fields (`None`, *optional*):
118
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
119
+
120
+ Outputs:
121
+ latents (`Tensor`):
122
+ Denoised latents.
123
+ """
124
+
125
+ model_name = "cosmos3-omni"
126
+ block_classes = [
127
+ Cosmos3PrepareTextSegmentsStep,
128
+ Cosmos3VisionPrepareLatentsStep,
129
+ Cosmos3VisionPackSequenceStep,
130
+ Cosmos3VisionDenoiseInputStep,
131
+ Cosmos3DistilledSetTimestepsStep,
132
+ Cosmos3DistilledVisionDenoiseStep,
133
+ ]
134
+ block_names = [
135
+ "prepare_text_segments",
136
+ "prepare_vision_latents",
137
+ "pack_vision_sequence",
138
+ "prepare_vision_denoiser_inputs",
139
+ "set_timesteps",
140
+ "denoise",
141
+ ]
142
+
143
+ @property
144
+ def description(self):
145
+ return "Runs the text-and-vision distilled Cosmos3 denoising workflow."
146
+
147
+ @property
148
+ def outputs(self):
149
+ return [OutputParam.template("latents")]
150
+
151
+
152
+ # auto_docstring
153
+ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
154
+ """
155
+ Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes.
156
+
157
+ Supported workflows:
158
+ - `text2image`: requires `prompt`, `num_frames`
159
+ - `text2video`: requires `prompt`
160
+ - `image2video`: requires `prompt`, `image`
161
+ - `video2video`: requires `prompt`, `video`
162
+
163
+ Components:
164
+ text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer
165
+ (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`)
166
+
167
+ Configs:
168
+ default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True)
169
+ distilled_sigmas (default: None)
170
+
171
+ Inputs:
172
+ prompt (`str`):
173
+ The text prompt that guides Cosmos3 generation.
174
+ num_frames (`int`, *optional*):
175
+ Number of frames to generate.
176
+ height (`int`, *optional*):
177
+ Height of the generated video or image in pixels.
178
+ width (`int`, *optional*):
179
+ Width of the generated video or image in pixels.
180
+ fps (`float`, *optional*, defaults to 24.0):
181
+ Frame rate of the generated video.
182
+ use_system_prompt (`bool`, *optional*, defaults to True):
183
+ Whether to prepend the Cosmos3 system prompt.
184
+ add_resolution_template (`bool`, *optional*, defaults to True):
185
+ Whether to add resolution metadata to the prompt.
186
+ add_duration_template (`bool`, *optional*, defaults to True):
187
+ Whether to add duration metadata to the prompt.
188
+ video (`None`, *optional*):
189
+ Reference video for video-to-video conditioning.
190
+ condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)):
191
+ Latent-frame indexes to preserve from the conditioning video.
192
+ condition_video_keep (`str`, *optional*, defaults to first):
193
+ Which end of a longer conditioning video to use: `first` or `last`.
194
+ image (`None`, *optional*):
195
+ Reference image for image-to-video conditioning.
196
+ x0_tokens_vision (`Tensor`, *optional*):
197
+ Vision latents encoded from the conditioning image or video.
198
+ vision_condition_frames (`list`, *optional*):
199
+ Latent-frame indexes fixed by visual conditioning.
200
+ latents (`Tensor`, *optional*):
201
+ Pre-generated noisy vision latents.
202
+ generator (`Generator`, *optional*):
203
+ Torch generator for deterministic generation.
204
+ num_inference_steps (`int`, *optional*):
205
+ The number of denoising steps.
206
+ guidance_scale (`float`, *optional*):
207
+ Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
208
+ forced to 1.0. Passing a value other than 1.0 raises an error.
209
+ **denoiser_input_fields (`None`, *optional*):
210
+ conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
211
+ output_type (`str`, *optional*, defaults to pil):
212
+ Output format: 'pil', 'np', 'pt'.
213
+
214
+ Outputs:
215
+ videos (`list`):
216
+ The generated videos.
217
+ """
218
+
219
+ model_name = "cosmos3-omni"
220
+ block_classes = [
221
+ Cosmos3DistilledTextEncoderStep,
222
+ Cosmos3DistilledAutoVaeEncoderStep,
223
+ Cosmos3DistilledVisionCoreDenoiseStep,
224
+ Cosmos3VideoDecodeStep,
225
+ ]
226
+ block_names = ["text_encoder", "vae_encoder", "denoise", "decode"]
227
+ _workflow_map = {
228
+ "text2image": {"prompt": True, "num_frames": 1},
229
+ "text2video": {"prompt": True},
230
+ "image2video": {"prompt": True, "image": True},
231
+ "video2video": {"prompt": True, "video": True},
232
+ }
233
+
234
+ @property
235
+ def description(self):
236
+ return "Modular pipeline blocks for distilled (few-step) Cosmos3 generation modes."
237
+
238
+ @property
239
+ def outputs(self):
240
+ return [OutputParam.template("videos")]
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from ...pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline, CosmosSafetyChecker
4
+ from ..modular_pipeline import ModularPipeline
5
+
6
+
7
+ class Cosmos3OmniModularPipeline(ModularPipeline):
8
+ """
9
+ A ModularPipeline for Cosmos 3 omni generation.
10
+ """
11
+
12
+ default_blocks_name = "Cosmos3OmniBlocks"
13
+
14
+ duration_template = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
15
+ image_resolution_template = "This image is of {height}x{width} resolution."
16
+ video_resolution_template = "This video is of {height}x{width} resolution."
17
+ inverse_duration_template = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS."
18
+ inverse_image_resolution_template = "This image is not of {height}x{width} resolution."
19
+ inverse_video_resolution_template = "This video is not of {height}x{width} resolution."
20
+
21
+ @property
22
+ def vae_scale_factor_spatial(self):
23
+ if getattr(self, "vae", None) is not None:
24
+ return int(self.vae.config.scale_factor_spatial)
25
+ return 16
26
+
27
+ @property
28
+ def vae_scale_factor_temporal(self):
29
+ if getattr(self, "vae", None) is not None:
30
+ return int(self.vae.config.scale_factor_temporal)
31
+ return 4
32
+
33
+ @property
34
+ def num_channels_latents(self):
35
+ if getattr(self, "transformer", None) is not None:
36
+ return int(self.transformer.config.latent_channel)
37
+ return 48
38
+
39
+ @property
40
+ def sound_sampling_rate(self):
41
+ if getattr(self, "sound_tokenizer", None) is not None:
42
+ return int(self.sound_tokenizer.config.sampling_rate)
43
+ return 48000
44
+
45
+ @property
46
+ def sound_hop_size(self):
47
+ if getattr(self, "sound_tokenizer", None) is not None:
48
+ return int(self.sound_tokenizer._hop_size)
49
+ return 1920
50
+
51
+ @property
52
+ def _vae_latents_mean(self):
53
+ return torch.tensor(self.vae.config.latents_mean, dtype=self.vae.dtype)
54
+
55
+ @property
56
+ def _vae_latents_inv_std(self):
57
+ return 1.0 / torch.tensor(self.vae.config.latents_std, dtype=self.vae.dtype)
58
+
59
+ @property
60
+ def llm_special_tokens(self):
61
+ if getattr(self, "text_tokenizer", None) is None:
62
+ return None
63
+ return {
64
+ "start_of_generation": self.text_tokenizer.convert_tokens_to_ids("<|vision_start|>"),
65
+ "eos_token_id": self.text_tokenizer.eos_token_id,
66
+ }
67
+
68
+ def enable_safety_checker(self, safety_checker=None):
69
+ if safety_checker is not None:
70
+ self.safety_checker = safety_checker
71
+ elif getattr(self, "safety_checker", None) is None:
72
+ self.safety_checker = CosmosSafetyChecker()
73
+ self._is_safety_checker_enabled = True
74
+
75
+ def disable_safety_checker(self):
76
+ self._is_safety_checker_enabled = False
77
+
78
+ @property
79
+ def requires_safety_checker(self):
80
+ return getattr(self, "_is_safety_checker_enabled", self.config.enable_safety_checker)
81
+
82
+ def _encode_video(self, x):
83
+ return Cosmos3OmniPipeline._encode_video(self, x)
84
+
85
+ def decode_sound(self, latent):
86
+ return Cosmos3OmniPipeline.decode_sound(self, latent)
87
+
88
+ def _prepare_text_segment(self, input_ids, device):
89
+ return Cosmos3OmniPipeline._prepare_text_segment(self, input_ids, device)
90
+
91
+ def _prepare_vision_segment(self, *args, **kwargs):
92
+ return Cosmos3OmniPipeline._prepare_vision_segment(self, *args, **kwargs)
93
+
94
+ def _prepare_sound_segment(self, *args, **kwargs):
95
+ return Cosmos3OmniPipeline._prepare_sound_segment(self, *args, **kwargs)
96
+
97
+ def _prepare_action_segment(self, *args, **kwargs):
98
+ return Cosmos3OmniPipeline._prepare_action_segment(self, *args, **kwargs)
99
+
100
+ def _prepare_action_video_conditioning(self, *args, **kwargs):
101
+ return Cosmos3OmniPipeline._prepare_action_video_conditioning(self, *args, **kwargs)
102
+
103
+ def _remove_action_video_padding_from_latent(self, *args, **kwargs):
104
+ return Cosmos3OmniPipeline._remove_action_video_padding_from_latent(self, *args, **kwargs)
105
+
106
+ @staticmethod
107
+ def _build_action_json_prompt(*args, **kwargs):
108
+ return Cosmos3OmniPipeline._build_action_json_prompt(*args, **kwargs)
109
+
110
+ def tokenize_prompt(self, *args, **kwargs):
111
+ return Cosmos3OmniPipeline.tokenize_prompt(self, *args, **kwargs)
112
+
113
+ @staticmethod
114
+ def _mask_velocity_predictions(*args, **kwargs):
115
+ return Cosmos3OmniPipeline._mask_velocity_predictions(*args, **kwargs)
116
+
117
+ def _apply_video_safety_check(self, *args, **kwargs):
118
+ return Cosmos3OmniPipeline._apply_video_safety_check(self, *args, **kwargs)
119
+
120
+
121
+ class Cosmos3DistilledModularPipeline(Cosmos3OmniModularPipeline):
122
+ """
123
+ A ModularPipeline for distilled (few-step) Cosmos 3 omni generation.
124
+
125
+ Distilled checkpoints bake classifier-free guidance into the weights and sample on a fixed schedule read from the
126
+ pipeline's `distilled_sigmas` config (populated from `modular_model_index.json`), so `guidance_scale` and
127
+ `num_inference_steps` are fixed and `negative_prompt` is not supported.
128
+ """
129
+
130
+ default_blocks_name = "Cosmos3DistilledBlocks"
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ...utils import (
4
+ DIFFUSERS_SLOW_IMPORT,
5
+ OptionalDependencyNotAvailable,
6
+ _LazyModule,
7
+ get_objects_from_module,
8
+ is_torch_available,
9
+ is_transformers_available,
10
+ )
11
+
12
+
13
+ _dummy_objects = {}
14
+ _import_structure = {}
15
+
16
+ try:
17
+ if not (is_transformers_available() and is_torch_available()):
18
+ raise OptionalDependencyNotAvailable()
19
+ except OptionalDependencyNotAvailable:
20
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
21
+
22
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
23
+ else:
24
+ _import_structure["modular_blocks_ernie_image"] = ["ErnieImageAutoBlocks"]
25
+ _import_structure["modular_pipeline"] = ["ErnieImageModularPipeline"]
26
+
27
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
28
+ try:
29
+ if not (is_transformers_available() and is_torch_available()):
30
+ raise OptionalDependencyNotAvailable()
31
+ except OptionalDependencyNotAvailable:
32
+ from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
33
+ else:
34
+ from .modular_blocks_ernie_image import ErnieImageAutoBlocks
35
+ from .modular_pipeline import ErnieImageModularPipeline
36
+ else:
37
+ import sys
38
+
39
+ sys.modules[__name__] = _LazyModule(
40
+ __name__,
41
+ globals()["__file__"],
42
+ _import_structure,
43
+ module_spec=__spec__,
44
+ )
45
+
46
+ for name, value in _dummy_objects.items():
47
+ setattr(sys.modules[__name__], name, value)
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py ADDED
@@ -0,0 +1,270 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import torch
16
+
17
+ from ...models import ErnieImageTransformer2DModel
18
+ from ...schedulers import FlowMatchEulerDiscreteScheduler
19
+ from ...utils import logging
20
+ from ...utils.torch_utils import randn_tensor
21
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
22
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
23
+ from .modular_pipeline import ErnieImageModularPipeline
24
+
25
+
26
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
27
+
28
+
29
+ def _pad_text(
30
+ text_hiddens: list[torch.Tensor], device: torch.device, dtype: torch.dtype, text_in_dim: int
31
+ ) -> tuple[torch.Tensor, torch.Tensor]:
32
+ """Pad a list of variable-length text hidden states to a common length and return (padded, lengths)."""
33
+ batch_size = len(text_hiddens)
34
+ if batch_size == 0:
35
+ return (
36
+ torch.zeros((0, 0, text_in_dim), device=device, dtype=dtype),
37
+ torch.zeros((0,), device=device, dtype=torch.long),
38
+ )
39
+ normalized = [t.squeeze(1).to(device).to(dtype) if t.dim() == 3 else t.to(device).to(dtype) for t in text_hiddens]
40
+ lengths = torch.tensor([t.shape[0] for t in normalized], device=device, dtype=torch.long)
41
+ max_length = int(lengths.max().item())
42
+ padded = torch.zeros((batch_size, max_length, text_in_dim), device=device, dtype=dtype)
43
+ for i, t in enumerate(normalized):
44
+ padded[i, : t.shape[0], :] = t
45
+ return padded, lengths
46
+
47
+
48
+ class ErnieImageTextInputStep(ModularPipelineBlocks):
49
+ model_name = "ernie-image"
50
+
51
+ @property
52
+ def description(self) -> str:
53
+ return (
54
+ "Input processing step that pads the variable-length text hidden states to a common length and "
55
+ "produces `text_bth` / `text_lens` tensors consumed by the denoiser."
56
+ )
57
+
58
+ @property
59
+ def expected_components(self) -> list[ComponentSpec]:
60
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
61
+
62
+ @property
63
+ def inputs(self) -> list[InputParam]:
64
+ return [
65
+ InputParam(
66
+ "prompt_embeds",
67
+ required=True,
68
+ type_hint=list,
69
+ description="List of per-prompt text embeddings from the text encoder step.",
70
+ ),
71
+ InputParam(
72
+ "negative_prompt_embeds",
73
+ type_hint=list,
74
+ description="List of per-prompt negative text embeddings from the text encoder step.",
75
+ ),
76
+ InputParam(
77
+ "num_images_per_prompt",
78
+ type_hint=int,
79
+ default=1,
80
+ description="Number of images to generate per prompt.",
81
+ ),
82
+ ]
83
+
84
+ @property
85
+ def intermediate_outputs(self) -> list[OutputParam]:
86
+ return [
87
+ OutputParam("batch_size", type_hint=int, description="The number of prompts in the batch."),
88
+ OutputParam(
89
+ "text_bth",
90
+ type_hint=torch.Tensor,
91
+ kwargs_type="denoiser_input_fields",
92
+ description="Padded text hidden states of shape (B, T_max, H) fed into the transformer.",
93
+ ),
94
+ OutputParam(
95
+ "text_lens",
96
+ type_hint=torch.Tensor,
97
+ kwargs_type="denoiser_input_fields",
98
+ description="Actual per-prompt text lengths used to build the transformer attention mask.",
99
+ ),
100
+ OutputParam(
101
+ "negative_text_bth",
102
+ type_hint=torch.Tensor,
103
+ kwargs_type="denoiser_input_fields",
104
+ description="Padded negative text hidden states, when classifier-free guidance is enabled.",
105
+ ),
106
+ OutputParam(
107
+ "negative_text_lens",
108
+ type_hint=torch.Tensor,
109
+ kwargs_type="denoiser_input_fields",
110
+ description="Actual per-prompt negative text lengths, when classifier-free guidance is enabled.",
111
+ ),
112
+ ]
113
+
114
+ @staticmethod
115
+ def _expand(hiddens: list[torch.Tensor], num_images_per_prompt: int) -> list[torch.Tensor]:
116
+ if num_images_per_prompt == 1:
117
+ return list(hiddens)
118
+ return [h for h in hiddens for _ in range(num_images_per_prompt)]
119
+
120
+ @torch.no_grad()
121
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
122
+ block_state = self.get_block_state(state)
123
+ device = components._execution_device
124
+ dtype = components.transformer.dtype
125
+ text_in_dim = components.text_in_dim
126
+ num_images_per_prompt = block_state.num_images_per_prompt
127
+
128
+ prompt_embeds = block_state.prompt_embeds
129
+ block_state.batch_size = len(prompt_embeds)
130
+
131
+ prompt_embeds = self._expand(prompt_embeds, num_images_per_prompt)
132
+ text_bth, text_lens = _pad_text(prompt_embeds, device, dtype, text_in_dim)
133
+ block_state.text_bth = text_bth
134
+ block_state.text_lens = text_lens
135
+
136
+ negative_prompt_embeds = block_state.negative_prompt_embeds
137
+ if negative_prompt_embeds is not None:
138
+ negative_prompt_embeds = self._expand(negative_prompt_embeds, num_images_per_prompt)
139
+ negative_text_bth, negative_text_lens = _pad_text(negative_prompt_embeds, device, dtype, text_in_dim)
140
+ block_state.negative_text_bth = negative_text_bth
141
+ block_state.negative_text_lens = negative_text_lens
142
+ else:
143
+ block_state.negative_text_bth = None
144
+ block_state.negative_text_lens = None
145
+
146
+ self.set_block_state(state, block_state)
147
+ return components, state
148
+
149
+
150
+ class ErnieImageSetTimestepsStep(ModularPipelineBlocks):
151
+ model_name = "ernie-image"
152
+
153
+ @property
154
+ def description(self) -> str:
155
+ return "Step that sets the scheduler's timesteps for inference using a linear sigma schedule."
156
+
157
+ @property
158
+ def expected_components(self) -> list[ComponentSpec]:
159
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
160
+
161
+ @property
162
+ def inputs(self) -> list[InputParam]:
163
+ return [
164
+ InputParam(
165
+ "num_inference_steps",
166
+ type_hint=int,
167
+ default=50,
168
+ description="Number of denoising steps.",
169
+ ),
170
+ ]
171
+
172
+ @property
173
+ def intermediate_outputs(self) -> list[OutputParam]:
174
+ return [
175
+ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference."),
176
+ OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps."),
177
+ ]
178
+
179
+ @torch.no_grad()
180
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
181
+ block_state = self.get_block_state(state)
182
+ device = components._execution_device
183
+ num_inference_steps = block_state.num_inference_steps
184
+
185
+ sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1]
186
+ components.scheduler.set_timesteps(sigmas=sigmas, device=device)
187
+
188
+ block_state.timesteps = components.scheduler.timesteps
189
+ block_state.num_inference_steps = num_inference_steps
190
+
191
+ self.set_block_state(state, block_state)
192
+ return components, state
193
+
194
+
195
+ class ErnieImagePrepareLatentsStep(ModularPipelineBlocks):
196
+ model_name = "ernie-image"
197
+
198
+ @property
199
+ def description(self) -> str:
200
+ return "Prepare random noise latents for the ErnieImage text-to-image denoising process."
201
+
202
+ @property
203
+ def expected_components(self) -> list[ComponentSpec]:
204
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
205
+
206
+ @property
207
+ def inputs(self) -> list[InputParam]:
208
+ return [
209
+ InputParam("height", type_hint=int, description="The height in pixels of the generated image."),
210
+ InputParam("width", type_hint=int, description="The width in pixels of the generated image."),
211
+ InputParam(
212
+ "latents",
213
+ type_hint=torch.Tensor,
214
+ description="Pre-generated noisy latents. If provided, skips noise sampling.",
215
+ ),
216
+ InputParam(
217
+ "generator",
218
+ type_hint=torch.Generator,
219
+ description="Torch generator for deterministic noise sampling.",
220
+ ),
221
+ InputParam(
222
+ "text_bth",
223
+ required=True,
224
+ type_hint=torch.Tensor,
225
+ description="Padded text hidden states; used to derive the total batch size for the latents.",
226
+ ),
227
+ ]
228
+
229
+ @property
230
+ def intermediate_outputs(self) -> list[OutputParam]:
231
+ return [
232
+ OutputParam("latents", type_hint=torch.Tensor, description="The initial noise latents to denoise."),
233
+ OutputParam("height", type_hint=int, description="The resolved image height in pixels."),
234
+ OutputParam("width", type_hint=int, description="The resolved image width in pixels."),
235
+ ]
236
+
237
+ @staticmethod
238
+ def _check_inputs(components: ErnieImageModularPipeline, height: int, width: int) -> None:
239
+ vae_scale_factor = components.vae_scale_factor
240
+ if height % vae_scale_factor != 0 or width % vae_scale_factor != 0:
241
+ raise ValueError(
242
+ f"`height` and `width` must be divisible by {vae_scale_factor}, got {height} and {width}."
243
+ )
244
+
245
+ @torch.no_grad()
246
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
247
+ block_state = self.get_block_state(state)
248
+ device = components._execution_device
249
+ dtype = components.transformer.dtype
250
+
251
+ height = block_state.height or components.default_height
252
+ width = block_state.width or components.default_width
253
+ self._check_inputs(components, height, width)
254
+
255
+ total_batch_size = block_state.text_bth.shape[0]
256
+ latent_h = height // components.vae_scale_factor
257
+ latent_w = width // components.vae_scale_factor
258
+ num_channels_latents = components.num_channels_latents
259
+
260
+ shape = (total_batch_size, num_channels_latents, latent_h, latent_w)
261
+ if block_state.latents is None:
262
+ block_state.latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=dtype)
263
+ else:
264
+ block_state.latents = block_state.latents.to(device=device, dtype=dtype)
265
+
266
+ block_state.height = height
267
+ block_state.width = width
268
+
269
+ self.set_block_state(state, block_state)
270
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import torch
16
+
17
+ from ...configuration_utils import FrozenDict
18
+ from ...image_processor import VaeImageProcessor
19
+ from ...models import AutoencoderKLFlux2
20
+ from ...utils import logging
21
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
22
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
23
+ from .modular_pipeline import ErnieImageModularPipeline, ErnieImagePachifier
24
+
25
+
26
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
27
+
28
+
29
+ class ErnieImageVaeDecoderStep(ModularPipelineBlocks):
30
+ model_name = "ernie-image"
31
+
32
+ @property
33
+ def description(self) -> str:
34
+ return "Step that decodes the denoised latents into images (unpachify, BN denormalization, VAE decode)."
35
+
36
+ @property
37
+ def expected_components(self) -> list[ComponentSpec]:
38
+ return [
39
+ ComponentSpec("vae", AutoencoderKLFlux2),
40
+ ComponentSpec(
41
+ "pachifier",
42
+ ErnieImagePachifier,
43
+ config=FrozenDict({"patch_size": 2}),
44
+ default_creation_method="from_config",
45
+ ),
46
+ ComponentSpec(
47
+ "image_processor",
48
+ VaeImageProcessor,
49
+ config=FrozenDict({"vae_scale_factor": 16}),
50
+ default_creation_method="from_config",
51
+ ),
52
+ ]
53
+
54
+ @property
55
+ def inputs(self) -> list[InputParam]:
56
+ return [
57
+ InputParam(
58
+ "latents",
59
+ required=True,
60
+ type_hint=torch.Tensor,
61
+ description="The latents to decode into images.",
62
+ ),
63
+ InputParam(
64
+ "output_type",
65
+ type_hint=str,
66
+ default="pil",
67
+ description="Output format: 'pil', 'np', or 'pt'.",
68
+ ),
69
+ ]
70
+
71
+ @property
72
+ def intermediate_outputs(self) -> list[OutputParam]:
73
+ return [OutputParam("images", type_hint=list, description="The generated images.")]
74
+
75
+ @torch.no_grad()
76
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
77
+ block_state = self.get_block_state(state)
78
+ vae = components.vae
79
+ device = block_state.latents.device
80
+
81
+ latents = block_state.latents
82
+ bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device=device, dtype=latents.dtype)
83
+ bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + 1e-5).to(device=device, dtype=latents.dtype)
84
+ latents = latents * bn_std + bn_mean
85
+
86
+ latents = components.pachifier.unpack_latents(latents)
87
+
88
+ images = vae.decode(latents.to(vae.dtype), return_dict=False)[0]
89
+ block_state.images = components.image_processor.postprocess(images, output_type=block_state.output_type)
90
+
91
+ self.set_block_state(state, block_state)
92
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py ADDED
@@ -0,0 +1,236 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import torch
16
+
17
+ from ...configuration_utils import FrozenDict
18
+ from ...guiders import ClassifierFreeGuidance
19
+ from ...models import ErnieImageTransformer2DModel
20
+ from ...schedulers import FlowMatchEulerDiscreteScheduler
21
+ from ...utils import logging
22
+ from ..modular_pipeline import (
23
+ BlockState,
24
+ LoopSequentialPipelineBlocks,
25
+ ModularPipelineBlocks,
26
+ PipelineState,
27
+ )
28
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
29
+ from .modular_pipeline import ErnieImageModularPipeline
30
+
31
+
32
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
33
+
34
+
35
+ class ErnieImageLoopBeforeDenoiser(ModularPipelineBlocks):
36
+ model_name = "ernie-image"
37
+
38
+ @property
39
+ def description(self) -> str:
40
+ return (
41
+ "Step within the denoising loop that prepares the latent model input and timestep tensor. "
42
+ "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
43
+ "object (e.g. `ErnieImageDenoiseLoopWrapper`)."
44
+ )
45
+
46
+ @property
47
+ def expected_components(self) -> list[ComponentSpec]:
48
+ return [ComponentSpec("transformer", ErnieImageTransformer2DModel)]
49
+
50
+ @property
51
+ def inputs(self) -> list[InputParam]:
52
+ return [
53
+ InputParam(
54
+ "latents",
55
+ required=True,
56
+ type_hint=torch.Tensor,
57
+ description="The latents to denoise.",
58
+ ),
59
+ ]
60
+
61
+ @torch.no_grad()
62
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
63
+ latents = block_state.latents
64
+ block_state.latent_model_input = latents.to(components.transformer.dtype)
65
+ block_state.timestep = t.expand(latents.shape[0]).to(components.transformer.dtype)
66
+ return components, block_state
67
+
68
+
69
+ class ErnieImageLoopDenoiser(ModularPipelineBlocks):
70
+ model_name = "ernie-image"
71
+
72
+ @property
73
+ def expected_components(self) -> list[ComponentSpec]:
74
+ return [
75
+ ComponentSpec("transformer", ErnieImageTransformer2DModel),
76
+ ComponentSpec(
77
+ "guider",
78
+ ClassifierFreeGuidance,
79
+ config=FrozenDict({"guidance_scale": 4.0}),
80
+ default_creation_method="from_config",
81
+ ),
82
+ ]
83
+
84
+ @property
85
+ def description(self) -> str:
86
+ return (
87
+ "Step within the denoising loop that runs the ErnieImage transformer with classifier-free guidance via "
88
+ "the configured guider."
89
+ )
90
+
91
+ @property
92
+ def inputs(self) -> list[InputParam]:
93
+ return [
94
+ InputParam(
95
+ "text_bth",
96
+ required=True,
97
+ type_hint=torch.Tensor,
98
+ description="Padded text hidden states fed into the transformer.",
99
+ ),
100
+ InputParam(
101
+ "text_lens",
102
+ required=True,
103
+ type_hint=torch.Tensor,
104
+ description="Per-prompt text lengths used by the transformer attention mask.",
105
+ ),
106
+ InputParam(
107
+ "negative_text_bth",
108
+ type_hint=torch.Tensor,
109
+ description="Padded negative text hidden states for classifier-free guidance.",
110
+ ),
111
+ InputParam(
112
+ "negative_text_lens",
113
+ type_hint=torch.Tensor,
114
+ description="Per-prompt negative text lengths for classifier-free guidance.",
115
+ ),
116
+ InputParam(
117
+ "num_inference_steps",
118
+ required=True,
119
+ type_hint=int,
120
+ description="Total number of denoising steps. Used by the guider for step-aware scheduling.",
121
+ ),
122
+ ]
123
+
124
+ @torch.no_grad()
125
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
126
+ guider_inputs = {
127
+ "text_bth": (block_state.text_bth, block_state.negative_text_bth),
128
+ "text_lens": (block_state.text_lens, block_state.negative_text_lens),
129
+ }
130
+
131
+ components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
132
+ guider_state = components.guider.prepare_inputs(guider_inputs)
133
+
134
+ for guider_state_batch in guider_state:
135
+ components.guider.prepare_models(components.transformer)
136
+ cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs.keys()}
137
+ noise_pred = components.transformer(
138
+ hidden_states=block_state.latent_model_input,
139
+ timestep=block_state.timestep,
140
+ return_dict=False,
141
+ **cond_kwargs,
142
+ )[0]
143
+ guider_state_batch.noise_pred = noise_pred
144
+ components.guider.cleanup_models(components.transformer)
145
+
146
+ block_state.noise_pred = components.guider(guider_state)[0]
147
+ return components, block_state
148
+
149
+
150
+ class ErnieImageLoopAfterDenoiser(ModularPipelineBlocks):
151
+ model_name = "ernie-image"
152
+
153
+ @property
154
+ def expected_components(self) -> list[ComponentSpec]:
155
+ return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
156
+
157
+ @property
158
+ def description(self) -> str:
159
+ return "Step within the denoising loop that updates the latents using the scheduler step."
160
+
161
+ @torch.no_grad()
162
+ def __call__(self, components: ErnieImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
163
+ latents_dtype = block_state.latents.dtype
164
+ block_state.latents = components.scheduler.step(
165
+ block_state.noise_pred, t, block_state.latents, return_dict=False
166
+ )[0]
167
+ if block_state.latents.dtype != latents_dtype and torch.backends.mps.is_available():
168
+ block_state.latents = block_state.latents.to(latents_dtype)
169
+ return components, block_state
170
+
171
+
172
+ class ErnieImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
173
+ model_name = "ernie-image"
174
+
175
+ @property
176
+ def description(self) -> str:
177
+ return (
178
+ "Pipeline block that iteratively denoises the latents over `timesteps`. "
179
+ "The specific steps within each iteration can be customized with `sub_blocks` attribute."
180
+ )
181
+
182
+ @property
183
+ def loop_expected_components(self) -> list[ComponentSpec]:
184
+ return [
185
+ ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
186
+ ComponentSpec("transformer", ErnieImageTransformer2DModel),
187
+ ]
188
+
189
+ @property
190
+ def loop_inputs(self) -> list[InputParam]:
191
+ return [
192
+ InputParam(
193
+ "timesteps",
194
+ required=True,
195
+ type_hint=torch.Tensor,
196
+ description="The timesteps to use for inference.",
197
+ ),
198
+ InputParam(
199
+ "num_inference_steps",
200
+ required=True,
201
+ type_hint=int,
202
+ description="The number of denoising steps.",
203
+ ),
204
+ ]
205
+
206
+ @property
207
+ def intermediate_outputs(self) -> list[OutputParam]:
208
+ return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents.")]
209
+
210
+ @torch.no_grad()
211
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
212
+ block_state = self.get_block_state(state)
213
+ with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
214
+ for i, t in enumerate(block_state.timesteps):
215
+ components, block_state = self.loop_step(components, block_state, i=i, t=t)
216
+ progress_bar.update()
217
+ self.set_block_state(state, block_state)
218
+ return components, state
219
+
220
+
221
+ class ErnieImageDenoiseStep(ErnieImageDenoiseLoopWrapper):
222
+ block_classes = [
223
+ ErnieImageLoopBeforeDenoiser,
224
+ ErnieImageLoopDenoiser,
225
+ ErnieImageLoopAfterDenoiser,
226
+ ]
227
+ block_names = ["before_denoiser", "denoiser", "after_denoiser"]
228
+
229
+ @property
230
+ def description(self) -> str:
231
+ return (
232
+ "Denoise step that iteratively denoises the latents. At each iteration it runs:\n"
233
+ " - `ErnieImageLoopBeforeDenoiser`\n"
234
+ " - `ErnieImageLoopDenoiser`\n"
235
+ " - `ErnieImageLoopAfterDenoiser`"
236
+ )
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py ADDED
@@ -0,0 +1,264 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import json
16
+
17
+ import torch
18
+ from transformers import AutoTokenizer, Mistral3Model
19
+
20
+ from ...configuration_utils import FrozenDict
21
+ from ...guiders import ClassifierFreeGuidance
22
+ from ...utils import logging
23
+ from ...utils.import_utils import is_transformers_version
24
+ from ..modular_pipeline import ModularPipelineBlocks, PipelineState
25
+ from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
26
+ from .modular_pipeline import ErnieImageModularPipeline
27
+
28
+
29
+ if is_transformers_version("<", "5.0.0"):
30
+ raise ImportError("`ErnieImageModularPipeline` requires `transformers>=5.0.0` for `Ministral3ForCausalLM`.")
31
+
32
+ from transformers import Ministral3ForCausalLM # noqa: E402
33
+
34
+
35
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
36
+
37
+
38
+ class ErnieImagePromptEnhancerStep(ModularPipelineBlocks):
39
+ model_name = "ernie-image"
40
+
41
+ @property
42
+ def description(self) -> str:
43
+ return "Prompt enhancer step that rewrites the input prompt using a causal language model (PE)."
44
+
45
+ @property
46
+ def expected_components(self) -> list[ComponentSpec]:
47
+ return [
48
+ ComponentSpec("pe", Ministral3ForCausalLM),
49
+ ComponentSpec("pe_tokenizer", AutoTokenizer),
50
+ ]
51
+
52
+ @property
53
+ def inputs(self) -> list[InputParam]:
54
+ return [
55
+ InputParam(
56
+ "prompt",
57
+ required=True,
58
+ type_hint=str,
59
+ description="The prompt or prompts to guide image generation.",
60
+ ),
61
+ InputParam("height", type_hint=int, description="The height in pixels of the generated image."),
62
+ InputParam("width", type_hint=int, description="The width in pixels of the generated image."),
63
+ InputParam(
64
+ "pe_system_prompt",
65
+ type_hint=str,
66
+ default=None,
67
+ description="Optional system prompt passed to the prompt enhancer.",
68
+ ),
69
+ InputParam(
70
+ "pe_temperature",
71
+ type_hint=float,
72
+ default=0.6,
73
+ description="Sampling temperature used when generating with the prompt enhancer.",
74
+ ),
75
+ InputParam(
76
+ "pe_top_p",
77
+ type_hint=float,
78
+ default=0.95,
79
+ description="Nucleus sampling `top_p` used when generating with the prompt enhancer.",
80
+ ),
81
+ ]
82
+
83
+ @property
84
+ def intermediate_outputs(self) -> list[OutputParam]:
85
+ return [
86
+ OutputParam("prompt", type_hint=list, description="The prompt list after prompt-enhancer rewriting."),
87
+ OutputParam("height", type_hint=int, description="The resolved image height in pixels."),
88
+ OutputParam("width", type_hint=int, description="The resolved image width in pixels."),
89
+ ]
90
+
91
+ @staticmethod
92
+ def _enhance_prompt(
93
+ pe: Ministral3ForCausalLM,
94
+ pe_tokenizer: AutoTokenizer,
95
+ prompt: str,
96
+ device: torch.device,
97
+ width: int,
98
+ height: int,
99
+ system_prompt: str | None,
100
+ temperature: float,
101
+ top_p: float,
102
+ ) -> str:
103
+ user_content = json.dumps({"prompt": prompt, "width": width, "height": height}, ensure_ascii=False)
104
+ messages = []
105
+ if system_prompt is not None:
106
+ messages.append({"role": "system", "content": system_prompt})
107
+ messages.append({"role": "user", "content": user_content})
108
+
109
+ input_text = pe_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
110
+ inputs = pe_tokenizer(input_text, return_tensors="pt").to(device)
111
+ output_ids = pe.generate(
112
+ **inputs,
113
+ max_new_tokens=pe_tokenizer.model_max_length,
114
+ do_sample=temperature != 1.0 or top_p != 1.0,
115
+ temperature=temperature,
116
+ top_p=top_p,
117
+ pad_token_id=pe_tokenizer.pad_token_id,
118
+ eos_token_id=pe_tokenizer.eos_token_id,
119
+ )
120
+ generated_ids = output_ids[0][inputs["input_ids"].shape[1] :]
121
+ return pe_tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
122
+
123
+ @torch.no_grad()
124
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
125
+ block_state = self.get_block_state(state)
126
+ device = components._execution_device
127
+
128
+ prompt = block_state.prompt
129
+ if isinstance(prompt, str):
130
+ prompt = [prompt]
131
+
132
+ height = block_state.height or components.default_height
133
+ width = block_state.width or components.default_width
134
+
135
+ revised = [
136
+ self._enhance_prompt(
137
+ pe=components.pe,
138
+ pe_tokenizer=components.pe_tokenizer,
139
+ prompt=p,
140
+ device=device,
141
+ width=width,
142
+ height=height,
143
+ system_prompt=block_state.pe_system_prompt,
144
+ temperature=block_state.pe_temperature,
145
+ top_p=block_state.pe_top_p,
146
+ )
147
+ for p in prompt
148
+ ]
149
+
150
+ block_state.prompt = revised
151
+ block_state.height = height
152
+ block_state.width = width
153
+
154
+ self.set_block_state(state, block_state)
155
+ return components, state
156
+
157
+
158
+ class ErnieImageTextEncoderStep(ModularPipelineBlocks):
159
+ model_name = "ernie-image"
160
+
161
+ @property
162
+ def description(self) -> str:
163
+ return (
164
+ "Text encoder step that encodes prompts into variable-length hidden states for the ErnieImage transformer."
165
+ )
166
+
167
+ @property
168
+ def expected_components(self) -> list[ComponentSpec]:
169
+ return [
170
+ ComponentSpec("text_encoder", Mistral3Model),
171
+ ComponentSpec("tokenizer", AutoTokenizer),
172
+ ComponentSpec(
173
+ "guider",
174
+ ClassifierFreeGuidance,
175
+ config=FrozenDict({"guidance_scale": 4.0}),
176
+ default_creation_method="from_config",
177
+ ),
178
+ ]
179
+
180
+ @property
181
+ def inputs(self) -> list[InputParam]:
182
+ return [
183
+ InputParam("prompt", type_hint=str, description="The prompt or prompts to guide image generation."),
184
+ InputParam(
185
+ "negative_prompt",
186
+ type_hint=str,
187
+ description="The prompt or prompts to avoid during image generation.",
188
+ ),
189
+ ]
190
+
191
+ @property
192
+ def intermediate_outputs(self) -> list[OutputParam]:
193
+ return [
194
+ OutputParam(
195
+ "prompt_embeds",
196
+ type_hint=list,
197
+ kwargs_type="denoiser_input_fields",
198
+ description="List of per-prompt text embeddings of shape (T, H).",
199
+ ),
200
+ OutputParam(
201
+ "negative_prompt_embeds",
202
+ type_hint=list,
203
+ kwargs_type="denoiser_input_fields",
204
+ description="List of per-prompt negative text embeddings for classifier-free guidance.",
205
+ ),
206
+ ]
207
+
208
+ @staticmethod
209
+ def _encode(
210
+ text_encoder: Mistral3Model,
211
+ tokenizer: AutoTokenizer,
212
+ prompt: list[str],
213
+ device: torch.device,
214
+ ) -> list[torch.Tensor]:
215
+ text_hiddens = []
216
+ for p in prompt:
217
+ ids = tokenizer(p, add_special_tokens=True, truncation=True, padding=False)["input_ids"]
218
+ if len(ids) == 0:
219
+ ids = [tokenizer.bos_token_id if tokenizer.bos_token_id is not None else 0]
220
+ input_ids = torch.tensor([ids], device=device)
221
+ outputs = text_encoder(input_ids=input_ids, output_hidden_states=True)
222
+ text_hiddens.append(outputs.hidden_states[-2][0])
223
+ return text_hiddens
224
+
225
+ @torch.no_grad()
226
+ def __call__(self, components: ErnieImageModularPipeline, state: PipelineState) -> PipelineState:
227
+ block_state = self.get_block_state(state)
228
+ device = components._execution_device
229
+
230
+ prompt = block_state.prompt
231
+ if prompt is None:
232
+ prompt = [""]
233
+ if isinstance(prompt, str):
234
+ prompt = [prompt]
235
+
236
+ block_state.prompt_embeds = self._encode(
237
+ text_encoder=components.text_encoder,
238
+ tokenizer=components.tokenizer,
239
+ prompt=prompt,
240
+ device=device,
241
+ )
242
+
243
+ if components.requires_unconditional_embeds:
244
+ negative_prompt = block_state.negative_prompt
245
+ if negative_prompt is None:
246
+ negative_prompt = ""
247
+ if isinstance(negative_prompt, str):
248
+ negative_prompt = [negative_prompt] * len(prompt)
249
+ if len(negative_prompt) != len(prompt):
250
+ raise ValueError(
251
+ f"`negative_prompt` must have the same length as `prompt` ({len(prompt)}), "
252
+ f"got {len(negative_prompt)}."
253
+ )
254
+ block_state.negative_prompt_embeds = self._encode(
255
+ text_encoder=components.text_encoder,
256
+ tokenizer=components.tokenizer,
257
+ prompt=negative_prompt,
258
+ device=device,
259
+ )
260
+ else:
261
+ block_state.negative_prompt_embeds = None
262
+
263
+ self.set_block_state(state, block_state)
264
+ return components, state
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py ADDED
@@ -0,0 +1,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from ...utils import logging
16
+ from ..modular_pipeline import ConditionalPipelineBlocks, SequentialPipelineBlocks
17
+ from ..modular_pipeline_utils import OutputParam
18
+ from .before_denoise import (
19
+ ErnieImagePrepareLatentsStep,
20
+ ErnieImageSetTimestepsStep,
21
+ ErnieImageTextInputStep,
22
+ )
23
+ from .decoders import ErnieImageVaeDecoderStep
24
+ from .denoise import ErnieImageDenoiseStep
25
+ from .encoders import ErnieImagePromptEnhancerStep, ErnieImageTextEncoderStep
26
+
27
+
28
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
29
+
30
+
31
+ # auto_docstring
32
+ class ErnieImageAutoPromptEnhancerStep(ConditionalPipelineBlocks):
33
+ """
34
+ Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.
35
+ - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.
36
+ - If `use_pe` is `None` or `False`, the step is skipped.
37
+
38
+ Components:
39
+ pe (`Ministral3ForCausalLM`) pe_tokenizer (`AutoTokenizer`)
40
+
41
+ Inputs:
42
+ prompt (`str`, *optional*):
43
+ The prompt or prompts to guide image generation.
44
+ height (`int`, *optional*):
45
+ The height in pixels of the generated image.
46
+ width (`int`, *optional*):
47
+ The width in pixels of the generated image.
48
+ pe_system_prompt (`str`, *optional*):
49
+ Optional system prompt passed to the prompt enhancer.
50
+ pe_temperature (`float`, *optional*, defaults to 0.6):
51
+ Sampling temperature used when generating with the prompt enhancer.
52
+ pe_top_p (`float`, *optional*, defaults to 0.95):
53
+ Nucleus sampling `top_p` used when generating with the prompt enhancer.
54
+
55
+ Outputs:
56
+ prompt (`list`):
57
+ The prompt list after prompt-enhancer rewriting.
58
+ height (`int`):
59
+ The resolved image height in pixels.
60
+ width (`int`):
61
+ The resolved image width in pixels.
62
+ """
63
+
64
+ model_name = "ernie-image"
65
+ block_classes = [ErnieImagePromptEnhancerStep]
66
+ block_names = ["prompt_enhancer"]
67
+ block_trigger_inputs = ["use_pe"]
68
+
69
+ def select_block(self, use_pe=None) -> str | None:
70
+ if use_pe:
71
+ return "prompt_enhancer"
72
+ return None
73
+
74
+ @property
75
+ def description(self):
76
+ return (
77
+ "Conditional block that runs the optional prompt enhancer when `use_pe` is truthy.\n"
78
+ " - `ErnieImagePromptEnhancerStep` is used when `use_pe=True`.\n"
79
+ " - If `use_pe` is `None` or `False`, the step is skipped."
80
+ )
81
+
82
+
83
+ # auto_docstring
84
+ class ErnieImageCoreDenoiseStep(SequentialPipelineBlocks):
85
+ """
86
+ Denoise block that takes encoded conditions and runs the denoising process for ErnieImage.
87
+
88
+ Components:
89
+ transformer (`ErnieImageTransformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider
90
+ (`ClassifierFreeGuidance`)
91
+
92
+ Inputs:
93
+ prompt_embeds (`list`):
94
+ List of per-prompt text embeddings from the text encoder step.
95
+ negative_prompt_embeds (`list`, *optional*):
96
+ List of per-prompt negative text embeddings from the text encoder step.
97
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
98
+ Number of images to generate per prompt.
99
+ num_inference_steps (`int`, *optional*, defaults to 50):
100
+ Number of denoising steps.
101
+ height (`int`, *optional*):
102
+ The height in pixels of the generated image.
103
+ width (`int`, *optional*):
104
+ The width in pixels of the generated image.
105
+ latents (`Tensor`, *optional*):
106
+ Pre-generated noisy latents. If provided, skips noise sampling.
107
+ generator (`Generator`, *optional*):
108
+ Torch generator for deterministic noise sampling.
109
+
110
+ Outputs:
111
+ latents (`Tensor`):
112
+ Denoised latents.
113
+ """
114
+
115
+ model_name = "ernie-image"
116
+ block_classes = [
117
+ ErnieImageTextInputStep,
118
+ ErnieImageSetTimestepsStep,
119
+ ErnieImagePrepareLatentsStep,
120
+ ErnieImageDenoiseStep,
121
+ ]
122
+ block_names = ["input", "set_timesteps", "prepare_latents", "denoise"]
123
+
124
+ @property
125
+ def description(self):
126
+ return "Denoise block that takes encoded conditions and runs the denoising process for ErnieImage."
127
+
128
+ @property
129
+ def outputs(self):
130
+ return [OutputParam.template("latents")]
131
+
132
+
133
+ # auto_docstring
134
+ class ErnieImageAutoBlocks(SequentialPipelineBlocks):
135
+ """
136
+ Auto modular pipeline for ErnieImage text-to-image generation. Supports an optional prompt enhancer when the `pe`
137
+ components are loaded and `use_pe=True`.
138
+
139
+ Supported workflows:
140
+ - `text2image`: requires `prompt`
141
+
142
+ Components:
143
+ pe (`Ministral3ForCausalLM`) pe_tokenizer (`AutoTokenizer`) text_encoder (`Mistral3Model`) tokenizer
144
+ (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer (`ErnieImageTransformer2DModel`) scheduler
145
+ (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLFlux2`) pachifier (`ErnieImagePachifier`)
146
+ image_processor (`VaeImageProcessor`)
147
+
148
+ Inputs:
149
+ prompt (`str`, *optional*):
150
+ The prompt or prompts to guide image generation.
151
+ height (`int`, *optional*):
152
+ The height in pixels of the generated image.
153
+ width (`int`, *optional*):
154
+ The width in pixels of the generated image.
155
+ pe_system_prompt (`str`, *optional*):
156
+ Optional system prompt passed to the prompt enhancer.
157
+ pe_temperature (`float`, *optional*, defaults to 0.6):
158
+ Sampling temperature used when generating with the prompt enhancer.
159
+ pe_top_p (`float`, *optional*, defaults to 0.95):
160
+ Nucleus sampling `top_p` used when generating with the prompt enhancer.
161
+ negative_prompt (`str`, *optional*):
162
+ The prompt or prompts to avoid during image generation.
163
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
164
+ Number of images to generate per prompt.
165
+ num_inference_steps (`int`, *optional*, defaults to 50):
166
+ Number of denoising steps.
167
+ latents (`Tensor`, *optional*):
168
+ Pre-generated noisy latents. If provided, skips noise sampling.
169
+ generator (`Generator`, *optional*):
170
+ Torch generator for deterministic noise sampling.
171
+ output_type (`str`, *optional*, defaults to pil):
172
+ Output format: 'pil', 'np', or 'pt'.
173
+
174
+ Outputs:
175
+ images (`list`):
176
+ Generated images.
177
+ """
178
+
179
+ model_name = "ernie-image"
180
+ block_classes = [
181
+ ErnieImageAutoPromptEnhancerStep,
182
+ ErnieImageTextEncoderStep,
183
+ ErnieImageCoreDenoiseStep,
184
+ ErnieImageVaeDecoderStep,
185
+ ]
186
+ block_names = ["prompt_enhancer", "text_encoder", "denoise", "decode"]
187
+ _workflow_map = {
188
+ "text2image": {"prompt": True},
189
+ }
190
+
191
+ @property
192
+ def description(self):
193
+ return (
194
+ "Auto modular pipeline for ErnieImage text-to-image generation. Supports an optional prompt enhancer "
195
+ "when the `pe` components are loaded and `use_pe=True`."
196
+ )
197
+
198
+ @property
199
+ def outputs(self):
200
+ return [OutputParam.template("images")]
.slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_pipeline.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Baidu ERNIE-Image Team and The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import torch
16
+
17
+ from ...configuration_utils import ConfigMixin, register_to_config
18
+ from ...loaders import ErnieImageLoraLoaderMixin
19
+ from ...utils import logging
20
+ from ..modular_pipeline import ModularPipeline
21
+
22
+
23
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
24
+
25
+
26
+ class ErnieImagePachifier(ConfigMixin):
27
+ """
28
+ A class to pack and unpack latents for ErnieImage.
29
+ """
30
+
31
+ config_name = "config.json"
32
+
33
+ @register_to_config
34
+ def __init__(self, patch_size: int = 2):
35
+ super().__init__()
36
+
37
+ def pack_latents(self, latents: torch.Tensor) -> torch.Tensor:
38
+ batch_size, num_channels, height, width = latents.shape
39
+ patch_size = self.config.patch_size
40
+
41
+ if height % patch_size != 0 or width % patch_size != 0:
42
+ raise ValueError(
43
+ f"Latent height and width must be divisible by {patch_size}, but got {height} and {width}"
44
+ )
45
+
46
+ latents = latents.view(
47
+ batch_size, num_channels, height // patch_size, patch_size, width // patch_size, patch_size
48
+ )
49
+ latents = latents.permute(0, 1, 3, 5, 2, 4)
50
+ return latents.reshape(
51
+ batch_size, num_channels * patch_size * patch_size, height // patch_size, width // patch_size
52
+ )
53
+
54
+ def unpack_latents(self, latents: torch.Tensor) -> torch.Tensor:
55
+ batch_size, num_channels, height, width = latents.shape
56
+ patch_size = self.config.patch_size
57
+
58
+ latents = latents.reshape(
59
+ batch_size, num_channels // (patch_size * patch_size), patch_size, patch_size, height, width
60
+ )
61
+ latents = latents.permute(0, 1, 4, 2, 5, 3)
62
+ return latents.reshape(
63
+ batch_size, num_channels // (patch_size * patch_size), height * patch_size, width * patch_size
64
+ )
65
+
66
+
67
+ class ErnieImageModularPipeline(ModularPipeline, ErnieImageLoraLoaderMixin):
68
+ """
69
+ A ModularPipeline for ErnieImage.
70
+
71
+ > [!WARNING] > This is an experimental feature and is likely to change in the future.
72
+ """
73
+
74
+ default_blocks_name = "ErnieImageAutoBlocks"
75
+
76
+ @property
77
+ def default_height(self):
78
+ return 1024
79
+
80
+ @property
81
+ def default_width(self):
82
+ return 1024
83
+
84
+ @property
85
+ def vae_scale_factor(self):
86
+ vae_scale_factor = 16
87
+ if hasattr(self, "vae") and self.vae is not None:
88
+ vae_scale_factor = 2 ** len(self.vae.config.block_out_channels)
89
+ return vae_scale_factor
90
+
91
+ @property
92
+ def num_channels_latents(self):
93
+ num_channels_latents = 128
94
+ if hasattr(self, "transformer") and self.transformer is not None:
95
+ num_channels_latents = self.transformer.config.in_channels
96
+ return num_channels_latents
97
+
98
+ @property
99
+ def text_in_dim(self):
100
+ text_in_dim = 3584
101
+ if hasattr(self, "transformer") and self.transformer is not None:
102
+ text_in_dim = self.transformer.config.text_in_dim
103
+ return text_in_dim
104
+
105
+ @property
106
+ def requires_unconditional_embeds(self):
107
+ requires_unconditional_embeds = False
108
+ if hasattr(self, "guider") and self.guider is not None:
109
+ requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
110
+ return requires_unconditional_embeds