Spaces:
Running on Zero
Running on Zero
feat: add turbo LoRA option (lightx2v 4-step distillation)
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitignore +3 -0
- .ignore +9 -0
- .slim/clonedeps.json +23 -0
- .slim/clonedeps/repos/huggingface__diffusers/.gitignore +188 -0
- .slim/clonedeps/repos/huggingface__diffusers/AGENTS.md +1 -0
- .slim/clonedeps/repos/huggingface__diffusers/CITATION.cff +52 -0
- .slim/clonedeps/repos/huggingface__diffusers/CLAUDE.md +1 -0
- .slim/clonedeps/repos/huggingface__diffusers/CODE_OF_CONDUCT.md +130 -0
- .slim/clonedeps/repos/huggingface__diffusers/CONTRIBUTING.md +1 -0
- .slim/clonedeps/repos/huggingface__diffusers/LICENSE +201 -0
- .slim/clonedeps/repos/huggingface__diffusers/MANIFEST.in +2 -0
- .slim/clonedeps/repos/huggingface__diffusers/Makefile +120 -0
- .slim/clonedeps/repos/huggingface__diffusers/PHILOSOPHY.md +1 -0
- .slim/clonedeps/repos/huggingface__diffusers/README.md +232 -0
- .slim/clonedeps/repos/huggingface__diffusers/SECURITY.md +106 -0
- .slim/clonedeps/repos/huggingface__diffusers/_typos.toml +13 -0
- .slim/clonedeps/repos/huggingface__diffusers/pyproject.toml +32 -0
- .slim/clonedeps/repos/huggingface__diffusers/setup.py +305 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/__init__.py +1746 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/callbacks.py +244 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/configuration_utils.py +752 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_check.py +34 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/dependency_versions_table.py +57 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/image_processor.py +1468 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/__init__.py +230 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/__init__.py +47 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/before_denoise.py +714 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/decoders.py +120 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/denoise.py +211 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/encoders.py +404 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py +381 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/anima/modular_pipeline.py +52 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/components_manager.py +1109 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/__init__.py +49 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/after_decode.py +113 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_denoise.py +1331 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/before_encoder.py +166 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/decoders.py +259 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/denoise.py +889 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/encoders.py +1056 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +1301 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +240 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +130 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/__init__.py +47 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/before_denoise.py +270 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/decoders.py +92 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/denoise.py +236 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/encoders.py +264 -0
- .slim/clonedeps/repos/huggingface__diffusers/src/diffusers/modular_pipelines/ernie_image/modular_blocks_ernie_image.py +200 -0
- .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
|