USO Windows test OK
This commit is contained in:
unknown
+54
@@ -0,0 +1,54 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# PyInstaller
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Virtual environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Temporary files
|
||||
temp/
|
||||
tmp/
|
||||
*.tmp
|
||||
*.temp
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,139 @@
|
||||
# ComfyUI USO Node
|
||||
|
||||
A custom node for ComfyUI that integrates USO (Unified Style and Subject-Driven Generation) for high-quality image generation with style and subject control.
|
||||
|
||||
## ✨ Features
|
||||
|
||||
- 🎨 **Unified Style & Subject Generation**: Powered by USO model based on FLUX architecture
|
||||
- 🎯 **Style-Driven Generation**: Generate images with specific artistic styles
|
||||
- 👤 **Subject-Driven Generation**: Maintain subject consistency across generations
|
||||
- 🔄 **Multi-Style Support**: Combine multiple styles in a single generation
|
||||
- ⚙️ **Memory Optimization**: FP8 precision support for consumer-grade GPUs (~16GB VRAM)
|
||||
- 🚀 **Flexible Control**: Advanced parameter control for fine-tuning results
|
||||
|
||||
## 🔧 Node List
|
||||
|
||||
### Core Nodes
|
||||
- **RH_USO_Loader**: Load and initialize USO models with optimization options
|
||||
- **RH_USO_Generator**: Generate images with style and subject control
|
||||
|
||||
## 🚀 Quick Installation
|
||||
|
||||
### Step 1: Install the Node
|
||||
```bash
|
||||
# Navigate to ComfyUI custom_nodes directory
|
||||
cd ComfyUI/custom_nodes
|
||||
|
||||
# Clone the repository
|
||||
git clone https://github.com/HM-RunningHub/ComfyUI_RH_USO
|
||||
|
||||
# Install dependencies
|
||||
cd ComfyUI_RH_USO
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Step 2: Download Required Models
|
||||
```bash
|
||||
# Download FLUX.1-dev model (Required base model)
|
||||
huggingface-cli download black-forest-labs/FLUX.1-dev flux1-dev.safetensors --local-dir models/diffusers/FLUX.1-dev
|
||||
huggingface-cli download black-forest-labs/FLUX.1-dev ae.safetensors --local-dir models/diffusers/FLUX.1-dev
|
||||
|
||||
# Download USO model
|
||||
huggingface-cli download bytedance-research/USO --local-dir models/uso
|
||||
|
||||
# Download SigLIP model
|
||||
huggingface-cli download google/siglip-so400m-patch14-384 --local-dir models/clip/siglip-so400m-patch14-384
|
||||
|
||||
# Final model structure should look like:
|
||||
models/
|
||||
├── diffusers/
|
||||
│ └── FLUX.1-dev/
|
||||
│ ├── flux1-dev.safetensors
|
||||
│ └── ae.safetensors
|
||||
├── uso/
|
||||
│ ├── assets/
|
||||
│ │ └── uso.webp
|
||||
│ ├── config.json
|
||||
│ ├── download_repo_enhanced.py
|
||||
│ ├── README.md
|
||||
│ └── uso_flux_v1.0/
|
||||
│ ├── dit_lora.safetensors
|
||||
│ └── projector.safetensors
|
||||
└── clip/
|
||||
└── siglip-so400m-patch14-384/
|
||||
|
||||
# Restart ComfyUI
|
||||
```
|
||||
|
||||
## 📖 Usage
|
||||
|
||||
### Basic Workflow
|
||||
```
|
||||
[RH_USO_Loader] → [RH_USO_Generator] → [Save Image]
|
||||
```
|
||||
|
||||
### Generation Types
|
||||
|
||||
#### Style-Driven Generation
|
||||
- Load style reference images
|
||||
- Input text prompt describing the content
|
||||
- Generate images in the specified style
|
||||
|
||||
#### Subject-Driven Generation
|
||||
- Load subject reference image
|
||||
- Input text prompt with scene description
|
||||
- Generate images maintaining subject identity
|
||||
|
||||
#### Style + Subject Generation
|
||||
- Load both style and subject reference images
|
||||
- Combine style transfer with subject consistency
|
||||
- Generate images with unified style and preserved subjects
|
||||
|
||||
## 🛠️ Technical Requirements
|
||||
|
||||
- **GPU**: 16GB+ VRAM (with FP8 optimization)
|
||||
- **RAM**: 32GB+ recommended
|
||||
- **Storage**: ~35GB for all models
|
||||
- FLUX.1-dev: ~24GB (flux1-dev.safetensors + ae.safetensors)
|
||||
- USO models: ~6GB
|
||||
- SigLIP: ~1.5GB
|
||||
- **CUDA**: Required for optimal performance
|
||||
|
||||
## ⚠️ Important Notes
|
||||
|
||||
- **Model Paths**: Models must be placed in specific directories:
|
||||
- FLUX.1-dev → `models/diffusers/FLUX.1-dev/`
|
||||
- USO models → `models/uso/`
|
||||
- SigLIP → `models/clip/siglip-so400m-patch14-384/`
|
||||
- FP8 mode recommended for consumer GPUs (reduces VRAM usage)
|
||||
- All model files must be downloaded before first use
|
||||
|
||||
## 📄 License
|
||||
|
||||
This project is licensed under Apache 2.0 License.
|
||||
|
||||
## 🔗 References
|
||||
|
||||
- [USO Project Page](https://bytedance.github.io/USO/)
|
||||
- [USO Paper](https://arxiv.org/abs/2508.18966)
|
||||
- [USO HuggingFace](https://huggingface.co/bytedance-research/USO)
|
||||
- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! Please feel free to submit issues and pull requests.
|
||||
|
||||
## ⭐ Citation
|
||||
|
||||
If you find this project useful, please consider citing the original USO paper:
|
||||
|
||||
```bibtex
|
||||
@article{wu2025uso,
|
||||
title={USO: Unified Style and Subject-Driven Generation via Disentangled and Reward Learning},
|
||||
author={Shaojin Wu and Mengqi Huang and Yufeng Cheng and Wenxu Wu and Jiahe Tian and Yiming Luo and Fei Ding and Qian He},
|
||||
year={2025},
|
||||
eprint={2508.18966},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV},
|
||||
}
|
||||
```
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# ComfyUI USO 节点
|
||||
|
||||
一个用于ComfyUI的自定义节点，集成USO（统一风格和主题驱动生成）模型，实现高质量的风格和主题控制图像生成。
|
||||
|
||||
## ✨ 特性
|
||||
|
||||
- 🎨 **统一风格与主题生成**: 基于FLUX架构的USO模型
|
||||
- 🎯 **风格驱动生成**: 根据特定艺术风格生成图像
|
||||
- 👤 **主题驱动生成**: 在生成过程中保持主题一致性
|
||||
- 🔄 **多风格支持**: 在单次生成中结合多种风格
|
||||
- ⚙️ **内存优化**: 支持FP8精度，适用于消费级GPU（约16GB显存）
|
||||
- 🚀 **灵活控制**: 高级参数控制，精细调节生成结果
|
||||
|
||||
## 🔧 节点列表
|
||||
|
||||
### 核心节点
|
||||
- **RH_USO_Loader**: 加载和初始化USO模型，包含优化选项
|
||||
- **RH_USO_Generator**: 具有风格和主题控制的图像生成器
|
||||
|
||||
## 🚀 快速安装
|
||||
|
||||
### 步骤1: 安装节点
|
||||
```bash
|
||||
# 进入ComfyUI自定义节点目录
|
||||
cd ComfyUI/custom_nodes
|
||||
|
||||
# 克隆仓库
|
||||
git clone https://github.com/HM-RunningHub/ComfyUI_RH_USO
|
||||
|
||||
# 安装依赖
|
||||
cd ComfyUI_RH_USO
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 步骤2: 下载所需模型
|
||||
```bash
|
||||
# 下载FLUX.1-dev模型（必需的基础模型）
|
||||
huggingface-cli download black-forest-labs/FLUX.1-dev flux1-dev.safetensors --local-dir models/diffusers/FLUX.1-dev
|
||||
huggingface-cli download black-forest-labs/FLUX.1-dev ae.safetensors --local-dir models/diffusers/FLUX.1-dev
|
||||
|
||||
# 下载USO模型
|
||||
huggingface-cli download bytedance-research/USO --local-dir models/uso
|
||||
|
||||
# 下载SigLIP模型
|
||||
huggingface-cli download google/siglip-so400m-patch14-384 --local-dir models/clip/siglip-so400m-patch14-384
|
||||
|
||||
# 最终模型结构应该如下：
|
||||
models/
|
||||
├── diffusers/
|
||||
│ └── FLUX.1-dev/
|
||||
│ ├── flux1-dev.safetensors
|
||||
│ └── ae.safetensors
|
||||
├── uso/
|
||||
│ ├── assets/
|
||||
│ │ └── uso.webp
|
||||
│ ├── config.json
|
||||
│ ├── download_repo_enhanced.py
|
||||
│ ├── README.md
|
||||
│ └── uso_flux_v1.0/
|
||||
│ ├── dit_lora.safetensors
|
||||
│ └── projector.safetensors
|
||||
└── clip/
|
||||
└── siglip-so400m-patch14-384/
|
||||
|
||||
# 重启ComfyUI
|
||||
```
|
||||
|
||||
## 📖 使用方法
|
||||
|
||||
### 基础工作流
|
||||
```
|
||||
[RH_USO_Loader] → [RH_USO_Generator] → [Save Image]
|
||||
```
|
||||
|
||||
### 生成类型
|
||||
|
||||
#### 风格驱动生成
|
||||
- 加载风格参考图像
|
||||
- 输入描述内容的文本提示
|
||||
- 生成指定风格的图像
|
||||
|
||||
#### 主题驱动生成
|
||||
- 加载主题参考图像
|
||||
- 输入包含场景描述的文本提示
|
||||
- 生成保持主题身份的图像
|
||||
|
||||
#### 风格+主题生成
|
||||
- 同时加载风格和主题参考图像
|
||||
- 结合风格转换与主题一致性
|
||||
- 生成具有统一风格且保持主题的图像
|
||||
|
||||
## 🛠️ 技术要求
|
||||
|
||||
- **GPU**: 16GB+显存（使用FP8优化）
|
||||
- **内存**: 推荐32GB+
|
||||
- **存储**: 约35GB用于所有模型
|
||||
- FLUX.1-dev: ~24GB (flux1-dev.safetensors + ae.safetensors)
|
||||
- USO模型: ~6GB
|
||||
- SigLIP: ~1.5GB
|
||||
- **CUDA**: 优化性能需要CUDA支持
|
||||
|
||||
## ⚠️ 重要提示
|
||||
|
||||
- **模型路径**: 模型必须放置在特定目录：
|
||||
- FLUX.1-dev → `models/diffusers/FLUX.1-dev/`
|
||||
- USO模型 → `models/uso/`
|
||||
- SigLIP → `models/clip/siglip-so400m-patch14-384/`
|
||||
- 推荐消费级GPU使用FP8模式（减少显存占用）
|
||||
- 所有模型文件必须在首次使用前下载完成
|
||||
|
||||
## 📄 许可证
|
||||
|
||||
本项目采用Apache 2.0许可证。
|
||||
|
||||
## 🔗 参考链接
|
||||
|
||||
- [USO项目页面](https://bytedance.github.io/USO/)
|
||||
- [USO论文](https://arxiv.org/abs/2508.18966)
|
||||
- [USO HuggingFace](https://huggingface.co/bytedance-research/USO)
|
||||
- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
||||
|
||||
## 🤝 贡献
|
||||
|
||||
欢迎贡献！请随时提交问题和拉取请求。
|
||||
|
||||
## ⭐ 引用
|
||||
|
||||
如果您觉得这个项目有用，请考虑引用原始USO论文：
|
||||
|
||||
```bibtex
|
||||
@article{wu2025uso,
|
||||
title={USO: Unified Style and Subject-Driven Generation via Disentangled and Reward Learning},
|
||||
author={Shaojin Wu and Mengqi Huang and Yufeng Cheng and Wenxu Wu and Jiahe Tian and Yiming Luo and Fei Ding and Qian He},
|
||||
year={2025},
|
||||
eprint={2508.18966},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV},
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,2 @@
|
||||
from .rh_uso_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,241 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import gradio as gr
|
||||
import torch
|
||||
|
||||
from uso.flux.pipeline import USOPipeline
|
||||
from transformers import SiglipVisionModel, SiglipImageProcessor
|
||||
|
||||
|
||||
with open("assets/uso_text.svg", "r", encoding="utf-8") as svg_file:
|
||||
text_content = svg_file.read()
|
||||
|
||||
with open("assets/uso_logo.svg", "r", encoding="utf-8") as svg_file:
|
||||
logo_content = svg_file.read()
|
||||
|
||||
title = f"""
|
||||
<div style="display: flex; align-items: center; justify-content: center;">
|
||||
<span style="transform: scale(0.7);margin-right: -5px;">{text_content}</span>
|
||||
<span style="font-size: 1.8em;margin-left: -10px;font-weight: bold; font-family: Gill Sans;">by UXO Team</span>
|
||||
<span style="margin-left: 0px; transform: scale(0.85); display: inline-block;">{logo_content}</span>
|
||||
</div>
|
||||
""".strip()
|
||||
|
||||
badges_text = r"""
|
||||
<div style="text-align: center; display: flex; justify-content: center; gap: 5px;">
|
||||
<a href="https://github.com/bytedance/USO"><img src="https://img.shields.io/static/v1?label=GitHub&message=Code&color=green&logo=github"></a>
|
||||
<a href="https://bytedance.github.io/USO/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-USO-yellow"></a>
|
||||
<a href="https://arxiv.org/abs/2504.02160"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-USO-b31b1b.svg"></a>
|
||||
<a href="https://huggingface.co/bytedance-research/USO"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=Model&color=orange"></a>
|
||||
</div>
|
||||
""".strip()
|
||||
|
||||
tips = """
|
||||
**What is USO?** 🎨
|
||||
USO is a unified style-subject optimized customization model and the latest addition to the UXO family (<a href='https://github.com/bytedance/USO' target='_blank'> USO</a> and <a href='https://github.com/bytedance/UNO' target='_blank'> UNO</a>).
|
||||
It can freely combine any subjects with any styles in any scenarios.
|
||||
|
||||
**How to use?** 💡
|
||||
We provide step-by-step instructions in our <a href='https://github.com/bytedance/USO' target='_blank'> Github Repo</a>.
|
||||
Additionally, try the examples provided below the demo to quickly get familiar with USO and spark your creativity!
|
||||
|
||||
<details>
|
||||
<summary style="cursor: pointer; color: #d34c0e; font-weight: 500;">The model is trained on 1024x1024 resolution and supports 3 types of usage. 📌 Tips:</summary>
|
||||
|
||||
* **Only content img**: support following types:
|
||||
* Subject/Identity-driven (supports natural prompt, e.g., *A clock on the table.* *The woman near the sea.*, excels in producing **photorealistic portraits**)
|
||||
* Style edit (layout-preserved): *Transform the image into Ghibli style/Pixel style/Retro comic style/Watercolor painting style...*.
|
||||
* Style edit (layout-shift): *Ghibli style, the man on the beach.*.
|
||||
* **Only style img**: Reference input style and generate anything following prompt. Excelling in this and further support multiple style references (in beta).
|
||||
* **Content img + style img**: Place the content into the desired style.
|
||||
* Layout-preserved: set prompt to **empty**.
|
||||
* Layout-shift: using natural prompt.</details>"""
|
||||
|
||||
star = r"""
|
||||
If USO is helpful, please help to ⭐ our <a href='https://github.com/bytedance/USO' target='_blank'> Github Repo</a>. Thanks a lot!"""
|
||||
|
||||
def get_examples(examples_dir: str = "assets/examples") -> list:
|
||||
examples = Path(examples_dir)
|
||||
ans = []
|
||||
for example in examples.iterdir():
|
||||
if not example.is_dir() or len(os.listdir(example)) == 0:
|
||||
continue
|
||||
with open(example / "config.json") as f:
|
||||
example_dict = json.load(f)
|
||||
|
||||
|
||||
example_list = []
|
||||
example_list.append(example_dict["prompt"]) # prompt
|
||||
|
||||
for key in ["image_ref1", "image_ref2", "image_ref3"]:
|
||||
if key in example_dict:
|
||||
example_list.append(str(example / example_dict[key]))
|
||||
else:
|
||||
example_list.append(None)
|
||||
|
||||
example_list.append(example_dict["seed"])
|
||||
ans.append(example_list)
|
||||
return ans
|
||||
|
||||
|
||||
def create_demo(
|
||||
model_type: str,
|
||||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||||
offload: bool = False,
|
||||
):
|
||||
pipeline = USOPipeline(
|
||||
model_type, device, offload, only_lora=True, lora_rank=128, hf_download=True
|
||||
)
|
||||
print("USOPipeline loaded successfully")
|
||||
|
||||
siglip_processor = SiglipImageProcessor.from_pretrained(
|
||||
"google/siglip-so400m-patch14-384"
|
||||
)
|
||||
siglip_model = SiglipVisionModel.from_pretrained(
|
||||
"google/siglip-so400m-patch14-384"
|
||||
)
|
||||
siglip_model.eval()
|
||||
siglip_model.to(device)
|
||||
pipeline.model.vision_encoder = siglip_model
|
||||
pipeline.model.vision_encoder_processor = siglip_processor
|
||||
print("SigLIP model loaded successfully")
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.Markdown(title)
|
||||
gr.Markdown(badges_text)
|
||||
gr.Markdown(tips)
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
prompt = gr.Textbox(label="Prompt", value="A beautiful woman.")
|
||||
with gr.Row():
|
||||
image_prompt1 = gr.Image(
|
||||
label="Content Reference Img", visible=True, interactive=True, type="pil"
|
||||
)
|
||||
image_prompt2 = gr.Image(
|
||||
label="Style Reference Img", visible=True, interactive=True, type="pil"
|
||||
)
|
||||
image_prompt3 = gr.Image(
|
||||
label="Extra Style Reference Img (Beta)", visible=True, interactive=True, type="pil"
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Row():
|
||||
width = gr.Slider(
|
||||
512, 1536, 1024, step=16, label="Generation Width"
|
||||
)
|
||||
height = gr.Slider(
|
||||
512, 1536, 1024, step=16, label="Generation Height"
|
||||
)
|
||||
with gr.Row():
|
||||
with gr.Row():
|
||||
keep_size = gr.Checkbox(
|
||||
label="Keep input size",
|
||||
value=False,
|
||||
interactive=True
|
||||
)
|
||||
with gr.Column():
|
||||
gr.Markdown("Set it to True if you only need style editing or want to keep the layout.")
|
||||
|
||||
with gr.Accordion("Advanced Options", open=True):
|
||||
with gr.Row():
|
||||
num_steps = gr.Slider(
|
||||
1, 50, 25, step=1, label="Number of steps"
|
||||
)
|
||||
guidance = gr.Slider(
|
||||
1.0, 5.0, 4.0, step=0.1, label="Guidance", interactive=True
|
||||
)
|
||||
content_long_size = gr.Slider(
|
||||
0, 1024, 512, step=16, label="Content reference size"
|
||||
)
|
||||
seed = gr.Number(-1, label="Seed (-1 for random)")
|
||||
|
||||
generate_btn = gr.Button("Generate")
|
||||
gr.Markdown(star)
|
||||
|
||||
with gr.Column():
|
||||
output_image = gr.Image(label="Generated Image")
|
||||
download_btn = gr.File(
|
||||
label="Download full-resolution", type="filepath", interactive=False
|
||||
)
|
||||
|
||||
inputs = [
|
||||
prompt,
|
||||
image_prompt1,
|
||||
image_prompt2,
|
||||
image_prompt3,
|
||||
seed,
|
||||
width,
|
||||
height,
|
||||
guidance,
|
||||
num_steps,
|
||||
keep_size,
|
||||
content_long_size,
|
||||
]
|
||||
generate_btn.click(
|
||||
fn=pipeline.gradio_generate,
|
||||
inputs=inputs,
|
||||
outputs=[output_image, download_btn],
|
||||
)
|
||||
|
||||
# example_text = gr.Text("", visible=False, label="Case For:")
|
||||
examples = get_examples("./assets/gradio_examples")
|
||||
|
||||
gr.Examples(
|
||||
examples=examples,
|
||||
inputs=[
|
||||
prompt,
|
||||
image_prompt1,
|
||||
image_prompt2,
|
||||
image_prompt3,
|
||||
seed,
|
||||
],
|
||||
# cache_examples='lazy',
|
||||
outputs=[output_image, download_btn],
|
||||
fn=pipeline.gradio_generate,
|
||||
label='row 1-4: identity/subject-driven; row 5-7: style-subject-driven; row 8-9: style-driven; row 10-12: multi-style-driven task; row 13: txt2img',
|
||||
examples_per_page=15
|
||||
)
|
||||
|
||||
return demo
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from typing import Literal
|
||||
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AppArgs:
|
||||
name: Literal["flux-dev", "flux-dev-fp8", "flux-schnell", "flux-krea-dev"] = "flux-dev"
|
||||
device: Literal["cuda", "cpu"] = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
offload: bool = dataclasses.field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "If True, sequantial offload the models(ae, dit, text encoder) to CPU if not used."
|
||||
},
|
||||
)
|
||||
port: int = 7860
|
||||
|
||||
parser = HfArgumentParser([AppArgs])
|
||||
args_tuple = parser.parse_args_into_dataclasses() # type: tuple[AppArgs]
|
||||
args = args_tuple[0]
|
||||
|
||||
demo = create_demo(args.name, args.device, args.offload)
|
||||
demo.launch(server_port=args.port)
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import dataclasses
|
||||
from typing import Literal
|
||||
|
||||
from accelerate import Accelerator
|
||||
from transformers import HfArgumentParser
|
||||
from PIL import Image
|
||||
import json
|
||||
import itertools
|
||||
import torch
|
||||
|
||||
from uso.flux.pipeline import USOPipeline, preprocess_ref
|
||||
from transformers import SiglipVisionModel, SiglipImageProcessor
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def horizontal_concat(images):
|
||||
widths, heights = zip(*(img.size for img in images))
|
||||
|
||||
total_width = sum(widths)
|
||||
max_height = max(heights)
|
||||
|
||||
new_im = Image.new("RGB", (total_width, max_height))
|
||||
|
||||
x_offset = 0
|
||||
for img in images:
|
||||
new_im.paste(img, (x_offset, 0))
|
||||
x_offset += img.size[0]
|
||||
|
||||
return new_im
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class InferenceArgs:
|
||||
prompt: str | None = None
|
||||
image_paths: list[str] | None = None
|
||||
eval_json_path: str | None = None
|
||||
# offload: bool = False
|
||||
offload: bool = True
|
||||
num_images_per_prompt: int = 1
|
||||
model_type: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev-fp8"
|
||||
width: int = 1024
|
||||
height: int = 1024
|
||||
num_steps: int = 25
|
||||
guidance: float = 4
|
||||
seed: int = 3407
|
||||
save_path: str = "output/inference"
|
||||
only_lora: bool = True
|
||||
concat_refs: bool = False
|
||||
lora_rank: int = 128
|
||||
pe: Literal["d", "h", "w", "o"] = "d"
|
||||
content_ref: int = 512
|
||||
ckpt_path: str | None = None
|
||||
use_siglip: bool = True
|
||||
instruct_edit: bool = False
|
||||
hf_download: bool = True
|
||||
|
||||
|
||||
def main(args: InferenceArgs):
|
||||
accelerator = Accelerator()
|
||||
|
||||
# init SigLIP model
|
||||
siglip_processor = None
|
||||
siglip_model = None
|
||||
|
||||
siglip_path = '/workspace/comfyui/models/clip/siglip-so400m-patch14-384'
|
||||
if args.use_siglip:
|
||||
siglip_processor = SiglipImageProcessor.from_pretrained(
|
||||
# "google/siglip-so400m-patch14-384"
|
||||
siglip_path
|
||||
)
|
||||
siglip_model = SiglipVisionModel.from_pretrained(
|
||||
# "google/siglip-so400m-patch14-384"
|
||||
siglip_path
|
||||
)
|
||||
siglip_model.eval()
|
||||
siglip_model.to(accelerator.device)
|
||||
print("SigLIP model loaded successfully")
|
||||
|
||||
pipeline = USOPipeline(
|
||||
args.model_type,
|
||||
accelerator.device,
|
||||
args.offload,
|
||||
only_lora=args.only_lora,
|
||||
lora_rank=args.lora_rank,
|
||||
hf_download=args.hf_download,
|
||||
)
|
||||
if args.use_siglip and siglip_model is not None:
|
||||
pipeline.model.vision_encoder = siglip_model
|
||||
print('-----> hook siglip encoder')
|
||||
|
||||
assert (
|
||||
args.prompt is not None or args.eval_json_path is not None
|
||||
), "Please provide either prompt or eval_json_path"
|
||||
|
||||
if args.eval_json_path is not None:
|
||||
with open(args.eval_json_path, "rt") as f:
|
||||
data_dicts = json.load(f)
|
||||
data_root = os.path.dirname(args.eval_json_path)
|
||||
else:
|
||||
data_root = ""
|
||||
data_dicts = [{"prompt": args.prompt, "image_paths": args.image_paths}]
|
||||
|
||||
print(
|
||||
f"process: {accelerator.num_processes}/{accelerator.process_index}, \
|
||||
process images: {len(data_dicts)}/{len(data_dicts[accelerator.process_index::accelerator.num_processes])}"
|
||||
)
|
||||
|
||||
data_dicts = data_dicts[accelerator.process_index :: accelerator.num_processes]
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
local_task_count = len(data_dicts) * args.num_images_per_prompt
|
||||
if accelerator.is_main_process:
|
||||
progress_bar = tqdm(total=local_task_count, desc="Generating Images")
|
||||
|
||||
for (i, data_dict), j in itertools.product(
|
||||
enumerate(data_dicts), range(args.num_images_per_prompt)
|
||||
):
|
||||
ref_imgs = []
|
||||
for _, img_path in enumerate(data_dict["image_paths"]):
|
||||
if img_path != "":
|
||||
img = Image.open(os.path.join(data_root, img_path)).convert("RGB")
|
||||
ref_imgs.append(img)
|
||||
else:
|
||||
ref_imgs.append(None)
|
||||
siglip_inputs = None
|
||||
if args.use_siglip and siglip_processor is not None:
|
||||
with torch.no_grad():
|
||||
siglip_inputs = [
|
||||
siglip_processor(img, return_tensors="pt").to(pipeline.device)
|
||||
for img in ref_imgs[1:] if isinstance(img, Image.Image)
|
||||
]
|
||||
|
||||
ref_imgs_pil = [
|
||||
preprocess_ref(img, args.content_ref) for img in ref_imgs[:1] if isinstance(img, Image.Image)
|
||||
]
|
||||
|
||||
if args.instruct_edit:
|
||||
args.width, args.height = ref_imgs_pil[0].size
|
||||
args.width, args.height = args.width * (1024 / args.content_ref), args.height * (1024 / args.content_ref)
|
||||
image_gen = pipeline(
|
||||
prompt=data_dict["prompt"],
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
guidance=args.guidance,
|
||||
num_steps=args.num_steps,
|
||||
seed=args.seed + j,
|
||||
ref_imgs=ref_imgs_pil,
|
||||
pe=args.pe,
|
||||
siglip_inputs=siglip_inputs,
|
||||
)
|
||||
if args.concat_refs:
|
||||
image_gen = horizontal_concat([image_gen, *ref_imgs])
|
||||
|
||||
if "save_dir" in data_dict:
|
||||
config_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.json")
|
||||
image_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.png")
|
||||
else:
|
||||
os.makedirs(args.save_path, exist_ok=True)
|
||||
config_save_path = os.path.join(args.save_path, f"{i}_{j}.json")
|
||||
image_save_path = os.path.join(args.save_path, f"{i}_{j}.png")
|
||||
|
||||
# save config and image
|
||||
os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
|
||||
image_gen.save(image_save_path)
|
||||
# ensure the prompt and image_paths are saved in the config file
|
||||
args.prompt = data_dict["prompt"]
|
||||
args.image_paths = data_dict["image_paths"]
|
||||
args_dict = vars(args)
|
||||
with open(config_save_path, "w") as f:
|
||||
json.dump(args_dict, f, indent=4)
|
||||
|
||||
if accelerator.is_main_process:
|
||||
progress_bar.update(1)
|
||||
if accelerator.is_main_process:
|
||||
progress_bar.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser([InferenceArgs])
|
||||
args = parser.parse_args_into_dataclasses()[0]
|
||||
main(args)
|
||||
@@ -0,0 +1,19 @@
|
||||
accelerate==1.1.1
|
||||
deepspeed==0.14.4
|
||||
einops==0.8.0
|
||||
transformers==4.43.3
|
||||
huggingface-hub
|
||||
diffusers==0.30.1
|
||||
sentencepiece==0.2.0
|
||||
gradio==5.22.0
|
||||
opencv-python
|
||||
matplotlib
|
||||
safetensors==0.4.5
|
||||
scipy==1.10.1
|
||||
numpy==1.24.4
|
||||
onnxruntime-gpu
|
||||
# httpx==0.23.3
|
||||
git+https://github.com/openai/CLIP.git
|
||||
--extra-index-url https://download.pytorch.org/whl/cu124
|
||||
torch==2.4.0
|
||||
torchvision==0.19.0
|
||||
@@ -0,0 +1 @@
|
||||
{"enable": true, "untracked_paths": []}
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
import os
|
||||
import dataclasses
|
||||
from typing import Literal
|
||||
|
||||
from accelerate import Accelerator
|
||||
from transformers import HfArgumentParser
|
||||
from PIL import Image
|
||||
import json
|
||||
import itertools
|
||||
import torch
|
||||
|
||||
from .uso.flux.pipeline import USOPipeline, preprocess_ref
|
||||
from transformers import SiglipVisionModel, SiglipImageProcessor
|
||||
from tqdm import tqdm
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
import comfy.utils
|
||||
|
||||
class RH_USO_Loader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("RHUSOMudules",)
|
||||
RETURN_NAMES = ("USO Modules",)
|
||||
FUNCTION = "load"
|
||||
|
||||
CATEGORY = "Runninghub/USO"
|
||||
|
||||
def load(self, **kwargs):
|
||||
# accelerator = Accelerator()
|
||||
device = 'cuda'
|
||||
siglip_path = os.path.join(folder_paths.models_dir, 'clip', 'siglip-so400m-patch14-384')
|
||||
siglip_processor = SiglipImageProcessor.from_pretrained(
|
||||
siglip_path
|
||||
)
|
||||
siglip_model = SiglipVisionModel.from_pretrained(
|
||||
siglip_path
|
||||
)
|
||||
siglip_model.eval()
|
||||
siglip_model.to(device)
|
||||
print("SigLIP model loaded successfully")
|
||||
|
||||
# hardcode hyperparamters -kiki
|
||||
model_type = 'flux-dev-fp8'
|
||||
lora_rank = 128
|
||||
|
||||
pipeline = USOPipeline(
|
||||
model_type,
|
||||
device,
|
||||
True, #args.offload,
|
||||
only_lora=True,
|
||||
lora_rank=lora_rank,
|
||||
hf_download=False,
|
||||
)
|
||||
if siglip_model is not None:
|
||||
pipeline.model.vision_encoder = siglip_model
|
||||
print('-----> hook siglip encoder')
|
||||
return ({'siglip_processor':siglip_processor, 'pipeline':pipeline}, )
|
||||
|
||||
class RH_USO_Sampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"uso": ("RHUSOMudules", ),
|
||||
"prompt": ("STRING", {"multiline": True,
|
||||
'default': ''}),
|
||||
"width": ("INT", {"default": 1024}),
|
||||
"height": ("INT", {"default": 1024}),
|
||||
"num_inference_steps": ("INT", {"default": 25}),
|
||||
"guidance": ("FLOAT", {"default": 4.0}),
|
||||
"seed": ("INT", {"default": 20, "min": 0, "max": 0xffffffffffffffff,
|
||||
"tooltip": "The random seed used for creating the noise."}),
|
||||
},
|
||||
"optional": {
|
||||
"content_image": ("IMAGE", ),
|
||||
"style_image": ("IMAGE", ),
|
||||
"style2_image": ("IMAGE", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "sample"
|
||||
|
||||
CATEGORY = "Runninghub/USO"
|
||||
|
||||
def tensor_2_pil(self, img_tensor):
|
||||
if img_tensor is not None:
|
||||
i = 255. * img_tensor.squeeze().cpu().numpy()
|
||||
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
||||
return img
|
||||
else:
|
||||
return None
|
||||
|
||||
def preprocess_ref(self, raw_image: Image.Image, long_size: int = 512, scale_ratio: int = 1):
|
||||
# 获取原始图像的宽度和高度
|
||||
image_w, image_h = raw_image.size
|
||||
if image_w == image_h and image_w == 16:
|
||||
return raw_image
|
||||
|
||||
# 计算长边和短边
|
||||
if image_w >= image_h:
|
||||
new_w = long_size
|
||||
new_h = int((long_size / image_w) * image_h)
|
||||
else:
|
||||
new_h = long_size
|
||||
new_w = int((long_size / image_h) * image_w)
|
||||
|
||||
# 按新的宽高进行等比例缩放
|
||||
raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
|
||||
|
||||
# 为了能让canny img进行scale
|
||||
scale_ratio = int(scale_ratio)
|
||||
target_w = new_w // (16 * scale_ratio) * (16 * scale_ratio)
|
||||
target_h = new_h // (16 * scale_ratio) * (16 * scale_ratio)
|
||||
|
||||
# 计算裁剪的起始坐标以实现中心裁剪
|
||||
left = (new_w - target_w) // 2
|
||||
top = (new_h - target_h) // 2
|
||||
right = left + target_w
|
||||
bottom = top + target_h
|
||||
|
||||
# 进行中心裁剪
|
||||
raw_image = raw_image.crop((left, top, right, bottom))
|
||||
|
||||
# 转换为 RGB 模式
|
||||
raw_image = raw_image.convert("RGB")
|
||||
return raw_image
|
||||
|
||||
def sample(self, **kwargs):
|
||||
ref_imgs = []
|
||||
content_image = self.tensor_2_pil(kwargs.get('content_image', None))
|
||||
style_image = self.tensor_2_pil(kwargs.get('style_image', None))
|
||||
style2_image = self.tensor_2_pil(kwargs.get('style2_image', None))
|
||||
print(f'conds-c/s1/s2:{content_image is not None} {style_image is not None} {style2_image is not None}')
|
||||
ref_imgs.append(content_image)
|
||||
if style_image is not None:
|
||||
ref_imgs.append(style_image)
|
||||
if style2_image is not None:
|
||||
ref_imgs.append(style2_image)
|
||||
siglip_inputs = None
|
||||
|
||||
width = kwargs.get('width')
|
||||
height = kwargs.get('height')
|
||||
prompt = kwargs.get('prompt')
|
||||
guidance = kwargs.get('guidance')
|
||||
num_steps = kwargs.get('num_inference_steps')
|
||||
seed = kwargs.get('seed') % (2 ** 32)
|
||||
|
||||
# hardcode hyperparameters -kiki
|
||||
content_ref = 512
|
||||
pe = 'd'
|
||||
|
||||
uso = kwargs.get('uso')
|
||||
siglip_processor = uso['siglip_processor']
|
||||
pipeline = uso['pipeline']
|
||||
with torch.no_grad():
|
||||
siglip_inputs = [
|
||||
siglip_processor(img, return_tensors="pt").to(pipeline.device)
|
||||
for img in ref_imgs[1:] if isinstance(img, Image.Image)
|
||||
]
|
||||
|
||||
ref_imgs_pil = [
|
||||
self.preprocess_ref(img, content_ref) for img in ref_imgs[:1] if isinstance(img, Image.Image)
|
||||
]
|
||||
self.pbar = comfy.utils.ProgressBar(num_steps)
|
||||
|
||||
image_gen = pipeline(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
guidance=guidance,
|
||||
num_steps=num_steps,
|
||||
seed=seed,
|
||||
ref_imgs=ref_imgs_pil,
|
||||
pe=pe,
|
||||
siglip_inputs=siglip_inputs,
|
||||
update_func=self.update,
|
||||
)
|
||||
|
||||
image = np.array(image_gen).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
|
||||
return (image, )
|
||||
|
||||
def update(self):
|
||||
self.pbar.update(1)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RunningHub USO Loader": RH_USO_Loader,
|
||||
"RunningHub USO Sampler":RH_USO_Sampler,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RunningHub USO Loader": "RunningHub USO Loader",
|
||||
"RunningHub USO Sampler": "RunningHub USO Sampler",
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
|
||||
q, k = apply_rope(q, k, pe)
|
||||
|
||||
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
x = rearrange(x, "B H L D -> B L (H D)")
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
|
||||
assert dim % 2 == 0
|
||||
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
|
||||
omega = 1.0 / (theta**scale)
|
||||
out = torch.einsum("...n,d->...nd", pos, omega)
|
||||
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
|
||||
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
|
||||
return out.float()
|
||||
|
||||
|
||||
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
|
||||
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
||||
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
||||
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
|
||||
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
|
||||
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from .modules.layers import (
|
||||
DoubleStreamBlock,
|
||||
EmbedND,
|
||||
LastLayer,
|
||||
MLPEmbedder,
|
||||
SingleStreamBlock,
|
||||
timestep_embedding,
|
||||
SigLIPMultiFeatProjModel,
|
||||
)
|
||||
import os
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxParams:
|
||||
in_channels: int
|
||||
vec_in_dim: int
|
||||
context_in_dim: int
|
||||
hidden_size: int
|
||||
mlp_ratio: float
|
||||
num_heads: int
|
||||
depth: int
|
||||
depth_single_blocks: int
|
||||
axes_dim: list[int]
|
||||
theta: int
|
||||
qkv_bias: bool
|
||||
guidance_embed: bool
|
||||
|
||||
|
||||
class Flux(nn.Module):
|
||||
"""
|
||||
Transformer model for flow matching on sequences.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
def __init__(self, params: FluxParams):
|
||||
super().__init__()
|
||||
|
||||
self.params = params
|
||||
self.in_channels = params.in_channels
|
||||
self.out_channels = self.in_channels
|
||||
if params.hidden_size % params.num_heads != 0:
|
||||
raise ValueError(
|
||||
f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
|
||||
)
|
||||
pe_dim = params.hidden_size // params.num_heads
|
||||
if sum(params.axes_dim) != pe_dim:
|
||||
raise ValueError(
|
||||
f"Got {params.axes_dim} but expected positional dim {pe_dim}"
|
||||
)
|
||||
self.hidden_size = params.hidden_size
|
||||
self.num_heads = params.num_heads
|
||||
self.pe_embedder = EmbedND(
|
||||
dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim
|
||||
)
|
||||
self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
|
||||
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
|
||||
self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
|
||||
self.guidance_in = (
|
||||
MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
|
||||
if params.guidance_embed
|
||||
else nn.Identity()
|
||||
)
|
||||
self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
|
||||
|
||||
self.double_blocks = nn.ModuleList(
|
||||
[
|
||||
DoubleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.num_heads,
|
||||
mlp_ratio=params.mlp_ratio,
|
||||
qkv_bias=params.qkv_bias,
|
||||
)
|
||||
for _ in range(params.depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.single_blocks = nn.ModuleList(
|
||||
[
|
||||
SingleStreamBlock(
|
||||
self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio
|
||||
)
|
||||
for _ in range(params.depth_single_blocks)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# feature embedder for siglip multi-feat inputs
|
||||
self.feature_embedder = SigLIPMultiFeatProjModel(
|
||||
siglip_token_nums=729,
|
||||
style_token_nums=64,
|
||||
siglip_token_dims=1152,
|
||||
hidden_size=self.hidden_size,
|
||||
context_layer_norm=True,
|
||||
)
|
||||
print("use semantic encoder siglip multi-feat to encode style image")
|
||||
|
||||
self.vision_encoder = None
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
@property
|
||||
def attn_processors(self):
|
||||
# set recursively
|
||||
processors = {} # type: dict[str, nn.Module]
|
||||
|
||||
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors):
|
||||
if hasattr(module, "set_processor"):
|
||||
processors[f"{name}.processor"] = module.processor
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
def set_attn_processor(self, processor):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
guidance: Tensor | None = None,
|
||||
ref_img: Tensor | None = None,
|
||||
ref_img_ids: Tensor | None = None,
|
||||
siglip_inputs: list[Tensor] | None = None,
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
|
||||
# running on sequences img
|
||||
img = self.img_in(img)
|
||||
vec = self.time_in(timestep_embedding(timesteps, 256))
|
||||
if self.params.guidance_embed:
|
||||
if guidance is None:
|
||||
raise ValueError(
|
||||
"Didn't get guidance strength for guidance distilled model."
|
||||
)
|
||||
vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
|
||||
vec = vec + self.vector_in(y)
|
||||
txt = self.txt_in(txt)
|
||||
if self.feature_embedder is not None and siglip_inputs is not None and len(siglip_inputs) > 0 and self.vision_encoder is not None:
|
||||
# processing style feat into textural hidden space
|
||||
siglip_embedding = [self.vision_encoder(**emb, output_hidden_states=True) for emb in siglip_inputs]
|
||||
# siglip_embedding = [self.vision_encoder(**(emb.to(torch.bfloat16)), output_hidden_states=True) for emb in siglip_inputs]
|
||||
siglip_embedding = torch.cat([self.feature_embedder(emb) for emb in siglip_embedding], dim=1)
|
||||
txt = torch.cat((siglip_embedding, txt), dim=1)
|
||||
siglip_embedding_ids = torch.zeros(
|
||||
siglip_embedding.shape[0], siglip_embedding.shape[1], 3
|
||||
).to(txt_ids.device)
|
||||
txt_ids = torch.cat((siglip_embedding_ids, txt_ids), dim=1)
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=1)
|
||||
|
||||
# concat ref_img/img
|
||||
img_end = img.shape[1]
|
||||
if ref_img is not None:
|
||||
if isinstance(ref_img, tuple) or isinstance(ref_img, list):
|
||||
img_in = [img] + [self.img_in(ref) for ref in ref_img]
|
||||
img_ids = [ids] + [ref_ids for ref_ids in ref_img_ids]
|
||||
img = torch.cat(img_in, dim=1)
|
||||
ids = torch.cat(img_ids, dim=1)
|
||||
else:
|
||||
img = torch.cat((img, self.img_in(ref_img)), dim=1)
|
||||
ids = torch.cat((ids, ref_img_ids), dim=1)
|
||||
pe = self.pe_embedder(ids)
|
||||
|
||||
for index_block, block in enumerate(self.double_blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
img, txt = torch.utils.checkpoint.checkpoint(
|
||||
block,
|
||||
img=img,
|
||||
txt=txt,
|
||||
vec=vec,
|
||||
pe=pe,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
|
||||
|
||||
img = torch.cat((txt, img), 1)
|
||||
for block in self.single_blocks:
|
||||
if self.training and self.gradient_checkpointing:
|
||||
img = torch.utils.checkpoint.checkpoint(
|
||||
block, img, vec=vec, pe=pe, use_reentrant=False
|
||||
)
|
||||
else:
|
||||
img = block(img, vec=vec, pe=pe)
|
||||
img = img[:, txt.shape[1] :, ...]
|
||||
# index img
|
||||
img = img[:, :img_end, ...]
|
||||
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
return img
|
||||
@@ -0,0 +1,327 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEncoderParams:
|
||||
resolution: int
|
||||
in_channels: int
|
||||
ch: int
|
||||
out_ch: int
|
||||
ch_mult: list[int]
|
||||
num_res_blocks: int
|
||||
z_channels: int
|
||||
scale_factor: float
|
||||
shift_factor: float
|
||||
|
||||
|
||||
def swish(x: Tensor) -> Tensor:
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
|
||||
def attention(self, h_: Tensor) -> Tensor:
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
b, c, h, w = q.shape
|
||||
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
|
||||
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
|
||||
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
|
||||
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x + self.proj_out(self.attention(x))
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
# downsampling
|
||||
self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1,) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
block_in = self.ch
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1])
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.ffactor = 2 ** (self.num_resolutions - 1)
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks + 1):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class DiagonalGaussian(nn.Module):
|
||||
def __init__(self, sample: bool = True, chunk_dim: int = 1):
|
||||
super().__init__()
|
||||
self.sample = sample
|
||||
self.chunk_dim = chunk_dim
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
|
||||
if self.sample:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mean + std * torch.randn_like(mean)
|
||||
else:
|
||||
return mean
|
||||
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, params: AutoEncoderParams):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
out_ch=params.out_ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.reg = DiagonalGaussian()
|
||||
|
||||
self.scale_factor = params.scale_factor
|
||||
self.shift_factor = params.shift_factor
|
||||
|
||||
def encode(self, x: Tensor) -> Tensor:
|
||||
z = self.reg(self.encoder(x))
|
||||
z = self.scale_factor * (z - self.shift_factor)
|
||||
return z
|
||||
|
||||
def decode(self, z: Tensor) -> Tensor:
|
||||
z = z / self.scale_factor + self.shift_factor
|
||||
return self.decoder(z)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.decode(self.encode(x))
|
||||
@@ -0,0 +1,55 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from torch import Tensor, nn
|
||||
from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
|
||||
T5Tokenizer)
|
||||
|
||||
|
||||
class HFEmbedder(nn.Module):
|
||||
def __init__(self, version: str, max_length: int, **hf_kwargs):
|
||||
super().__init__()
|
||||
# self.is_clip = "clip" in version.lower()
|
||||
#kiki
|
||||
self.is_clip = hf_kwargs.pop('is_clip', False)
|
||||
self.max_length = max_length
|
||||
self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
|
||||
|
||||
if self.is_clip:
|
||||
self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
|
||||
else:
|
||||
self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
|
||||
|
||||
self.hf_module = self.hf_module.eval().requires_grad_(False)
|
||||
|
||||
def forward(self, text: list[str]) -> Tensor:
|
||||
batch_encoding = self.tokenizer(
|
||||
text,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_length=False,
|
||||
return_overflowing_tokens=False,
|
||||
padding="max_length",
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
outputs = self.hf_module(
|
||||
input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
|
||||
attention_mask=None,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
return outputs[self.output_key]
|
||||
@@ -0,0 +1,631 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import Tensor, nn
|
||||
|
||||
from ..math import attention, rope
|
||||
|
||||
|
||||
class EmbedND(nn.Module):
|
||||
def __init__(self, dim: int, theta: int, axes_dim: list[int]):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
|
||||
def forward(self, ids: Tensor) -> Tensor:
|
||||
n_axes = ids.shape[-1]
|
||||
emb = torch.cat(
|
||||
[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
|
||||
dim=-3,
|
||||
)
|
||||
|
||||
return emb.unsqueeze(1)
|
||||
|
||||
|
||||
def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
t = time_factor * t
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(t.device)
|
||||
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
if torch.is_floating_point(t):
|
||||
embedding = embedding.to(t)
|
||||
return embedding
|
||||
|
||||
|
||||
class MLPEmbedder(nn.Module):
|
||||
def __init__(self, in_dim: int, hidden_dim: int):
|
||||
super().__init__()
|
||||
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
|
||||
self.silu = nn.SiLU()
|
||||
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.out_layer(self.silu(self.in_layer(x)))
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
|
||||
return ((x * rrms) * self.scale.float()).to(dtype=x_dtype)
|
||||
|
||||
|
||||
class QKNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.query_norm = RMSNorm(dim)
|
||||
self.key_norm = RMSNorm(dim)
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
|
||||
q = self.query_norm(q)
|
||||
k = self.key_norm(k)
|
||||
return q.to(v), k.to(v)
|
||||
|
||||
|
||||
class LoRALinearLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
rank=4,
|
||||
network_alpha=None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
|
||||
self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
|
||||
# This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
|
||||
# See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
|
||||
self.network_alpha = network_alpha
|
||||
self.rank = rank
|
||||
|
||||
nn.init.normal_(self.down.weight, std=1 / rank)
|
||||
nn.init.zeros_(self.up.weight)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
orig_dtype = hidden_states.dtype
|
||||
dtype = self.down.weight.dtype
|
||||
|
||||
down_hidden_states = self.down(hidden_states.to(dtype))
|
||||
up_hidden_states = self.up(down_hidden_states)
|
||||
|
||||
if self.network_alpha is not None:
|
||||
up_hidden_states *= self.network_alpha / self.rank
|
||||
|
||||
return up_hidden_states.to(orig_dtype)
|
||||
|
||||
|
||||
class FLuxSelfAttnProcessor:
|
||||
def __call__(self, attn, x, pe, **attention_kwargs):
|
||||
qkv = attn.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
x = attention(q, k, v, pe=pe)
|
||||
x = attn.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class LoraFluxAttnProcessor(nn.Module):
|
||||
|
||||
def __init__(self, dim: int, rank=4, network_alpha=None, lora_weight=1):
|
||||
super().__init__()
|
||||
self.qkv_lora = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
def __call__(self, attn, x, pe, **attention_kwargs):
|
||||
qkv = attn.qkv(x) + self.qkv_lora(x) * self.lora_weight
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
x = attention(q, k, v, pe=pe)
|
||||
x = attn.proj(x) + self.proj_lora(x) * self.lora_weight
|
||||
return x
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.norm = QKNorm(head_dim)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
def forward():
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModulationOut:
|
||||
shift: Tensor
|
||||
scale: Tensor
|
||||
gate: Tensor
|
||||
|
||||
|
||||
class Modulation(nn.Module):
|
||||
def __init__(self, dim: int, double: bool):
|
||||
super().__init__()
|
||||
self.is_double = double
|
||||
self.multiplier = 6 if double else 3
|
||||
self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
|
||||
|
||||
def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
|
||||
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(
|
||||
self.multiplier, dim=-1
|
||||
)
|
||||
|
||||
return (
|
||||
ModulationOut(*out[:3]),
|
||||
ModulationOut(*out[3:]) if self.is_double else None,
|
||||
)
|
||||
|
||||
|
||||
class DoubleStreamBlockLoraProcessor(nn.Module):
|
||||
def __init__(self, dim: int, rank=4, network_alpha=None, lora_weight=1):
|
||||
super().__init__()
|
||||
self.qkv_lora1 = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora1 = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.qkv_lora2 = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora2 = LoRALinearLayer(dim, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
def forward(self, attn, img, txt, vec, pe, **attention_kwargs):
|
||||
img_mod1, img_mod2 = attn.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = attn.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = attn.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = (
|
||||
attn.img_attn.qkv(img_modulated)
|
||||
+ self.qkv_lora1(img_modulated) * self.lora_weight
|
||||
)
|
||||
img_q, img_k, img_v = rearrange(
|
||||
img_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads
|
||||
)
|
||||
img_q, img_k = attn.img_attn.norm(img_q, img_k, img_v)
|
||||
|
||||
# prepare txt for attention
|
||||
txt_modulated = attn.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = (
|
||||
attn.txt_attn.qkv(txt_modulated)
|
||||
+ self.qkv_lora2(txt_modulated) * self.lora_weight
|
||||
)
|
||||
txt_q, txt_k, txt_v = rearrange(
|
||||
txt_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads
|
||||
)
|
||||
txt_q, txt_k = attn.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
|
||||
attn1 = attention(q, k, v, pe=pe)
|
||||
txt_attn, img_attn = attn1[:, : txt.shape[1]], attn1[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * (
|
||||
attn.img_attn.proj(img_attn) + self.proj_lora1(img_attn) * self.lora_weight
|
||||
)
|
||||
img = img + img_mod2.gate * attn.img_mlp(
|
||||
(1 + img_mod2.scale) * attn.img_norm2(img) + img_mod2.shift
|
||||
)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * (
|
||||
attn.txt_attn.proj(txt_attn) + self.proj_lora2(txt_attn) * self.lora_weight
|
||||
)
|
||||
txt = txt + txt_mod2.gate * attn.txt_mlp(
|
||||
(1 + txt_mod2.scale) * attn.txt_norm2(txt) + txt_mod2.shift
|
||||
)
|
||||
return img, txt
|
||||
|
||||
|
||||
class DoubleStreamBlockProcessor:
|
||||
def __call__(self, attn, img, txt, vec, pe, **attention_kwargs):
|
||||
img_mod1, img_mod2 = attn.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = attn.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = attn.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = attn.img_attn.qkv(img_modulated)
|
||||
img_q, img_k, img_v = rearrange(
|
||||
img_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim
|
||||
)
|
||||
img_q, img_k = attn.img_attn.norm(img_q, img_k, img_v)
|
||||
|
||||
# prepare txt for attention
|
||||
txt_modulated = attn.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = attn.txt_attn.qkv(txt_modulated)
|
||||
txt_q, txt_k, txt_v = rearrange(
|
||||
txt_qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads, D=attn.head_dim
|
||||
)
|
||||
txt_q, txt_k = attn.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
|
||||
attn1 = attention(q, k, v, pe=pe)
|
||||
txt_attn, img_attn = attn1[:, : txt.shape[1]], attn1[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * attn.img_attn.proj(img_attn)
|
||||
img = img + img_mod2.gate * attn.img_mlp(
|
||||
(1 + img_mod2.scale) * attn.img_norm2(img) + img_mod2.shift
|
||||
)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * attn.txt_attn.proj(txt_attn)
|
||||
txt = txt + txt_mod2.gate * attn.txt_mlp(
|
||||
(1 + txt_mod2.scale) * attn.txt_norm2(txt) + txt_mod2.shift
|
||||
)
|
||||
return img, txt
|
||||
|
||||
|
||||
class DoubleStreamBlock(nn.Module):
|
||||
def __init__(
|
||||
self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
self.num_heads = num_heads
|
||||
self.hidden_size = hidden_size
|
||||
self.head_dim = hidden_size // num_heads
|
||||
|
||||
self.img_mod = Modulation(hidden_size, double=True)
|
||||
self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_attn = SelfAttention(
|
||||
dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
|
||||
)
|
||||
|
||||
self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
|
||||
self.txt_mod = Modulation(hidden_size, double=True)
|
||||
self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_attn = SelfAttention(
|
||||
dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
|
||||
)
|
||||
|
||||
self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
processor = DoubleStreamBlockProcessor()
|
||||
self.set_processor(processor)
|
||||
|
||||
def set_processor(self, processor) -> None:
|
||||
self.processor = processor
|
||||
|
||||
def get_processor(self):
|
||||
return self.processor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
txt: Tensor,
|
||||
vec: Tensor,
|
||||
pe: Tensor,
|
||||
image_proj: Tensor = None,
|
||||
ip_scale: float = 1.0,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if image_proj is None:
|
||||
return self.processor(self, img, txt, vec, pe)
|
||||
else:
|
||||
return self.processor(self, img, txt, vec, pe, image_proj, ip_scale)
|
||||
|
||||
|
||||
class SingleStreamBlockLoraProcessor(nn.Module):
|
||||
def __init__(
|
||||
self, dim: int, rank: int = 4, network_alpha=None, lora_weight: float = 1
|
||||
):
|
||||
super().__init__()
|
||||
self.qkv_lora = LoRALinearLayer(dim, dim * 3, rank, network_alpha)
|
||||
self.proj_lora = LoRALinearLayer(15360, dim, rank, network_alpha)
|
||||
self.lora_weight = lora_weight
|
||||
|
||||
def forward(self, attn: nn.Module, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
|
||||
|
||||
mod, _ = attn.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * attn.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(
|
||||
attn.linear1(x_mod), [3 * attn.hidden_size, attn.mlp_hidden_dim], dim=-1
|
||||
)
|
||||
qkv = qkv + self.qkv_lora(x_mod) * self.lora_weight
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn_1 = attention(q, k, v, pe=pe)
|
||||
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = attn.linear2(torch.cat((attn_1, attn.mlp_act(mlp)), 2))
|
||||
output = (
|
||||
output
|
||||
+ self.proj_lora(torch.cat((attn_1, attn.mlp_act(mlp)), 2))
|
||||
* self.lora_weight
|
||||
)
|
||||
output = x + mod.gate * output
|
||||
return output
|
||||
|
||||
|
||||
class SingleStreamBlockProcessor:
|
||||
def __call__(
|
||||
self, attn: nn.Module, x: Tensor, vec: Tensor, pe: Tensor, **attention_kwargs
|
||||
) -> Tensor:
|
||||
|
||||
mod, _ = attn.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * attn.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(
|
||||
attn.linear1(x_mod), [3 * attn.hidden_size, attn.mlp_hidden_dim], dim=-1
|
||||
)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=attn.num_heads)
|
||||
q, k = attn.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn_1 = attention(q, k, v, pe=pe)
|
||||
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = attn.linear2(torch.cat((attn_1, attn.mlp_act(mlp)), 2))
|
||||
output = x + mod.gate * output
|
||||
return output
|
||||
|
||||
|
||||
class SingleStreamBlock(nn.Module):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_scale: float | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_size
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
self.scale = qk_scale or self.head_dim**-0.5
|
||||
|
||||
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
# qkv and mlp_in
|
||||
self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
|
||||
# proj and mlp_out
|
||||
self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
|
||||
|
||||
self.norm = QKNorm(self.head_dim)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
|
||||
self.mlp_act = nn.GELU(approximate="tanh")
|
||||
self.modulation = Modulation(hidden_size, double=False)
|
||||
|
||||
processor = SingleStreamBlockProcessor()
|
||||
self.set_processor(processor)
|
||||
|
||||
def set_processor(self, processor) -> None:
|
||||
self.processor = processor
|
||||
|
||||
def get_processor(self):
|
||||
return self.processor
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
vec: Tensor,
|
||||
pe: Tensor,
|
||||
image_proj: Tensor | None = None,
|
||||
ip_scale: float = 1.0,
|
||||
) -> Tensor:
|
||||
if image_proj is None:
|
||||
return self.processor(self, x, vec, pe)
|
||||
else:
|
||||
return self.processor(self, x, vec, pe, image_proj, ip_scale)
|
||||
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(
|
||||
hidden_size, patch_size * patch_size * out_channels, bias=True
|
||||
)
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor) -> Tensor:
|
||||
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
|
||||
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class SigLIPMultiFeatProjModel(torch.nn.Module):
|
||||
"""
|
||||
SigLIP Multi-Feature Projection Model for processing style features from different layers
|
||||
and projecting them into a unified hidden space.
|
||||
|
||||
Args:
|
||||
siglip_token_nums (int): Number of SigLIP tokens, default 257
|
||||
style_token_nums (int): Number of style tokens, default 256
|
||||
siglip_token_dims (int): Dimension of SigLIP tokens, default 1536
|
||||
hidden_size (int): Hidden layer size, default 3072
|
||||
context_layer_norm (bool): Whether to use context layer normalization, default False
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
siglip_token_nums: int = 257,
|
||||
style_token_nums: int = 256,
|
||||
siglip_token_dims: int = 1536,
|
||||
hidden_size: int = 3072,
|
||||
context_layer_norm: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# High-level feature processing (layer -2)
|
||||
self.high_embedding_linear = nn.Sequential(
|
||||
nn.Linear(siglip_token_nums, style_token_nums),
|
||||
nn.SiLU()
|
||||
)
|
||||
self.high_layer_norm = (
|
||||
nn.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
|
||||
)
|
||||
self.high_projection = nn.Linear(siglip_token_dims, hidden_size, bias=True)
|
||||
|
||||
# Mid-level feature processing (layer -11)
|
||||
self.mid_embedding_linear = nn.Sequential(
|
||||
nn.Linear(siglip_token_nums, style_token_nums),
|
||||
nn.SiLU()
|
||||
)
|
||||
self.mid_layer_norm = (
|
||||
nn.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
|
||||
)
|
||||
self.mid_projection = nn.Linear(siglip_token_dims, hidden_size, bias=True)
|
||||
|
||||
# Low-level feature processing (layer -20)
|
||||
self.low_embedding_linear = nn.Sequential(
|
||||
nn.Linear(siglip_token_nums, style_token_nums),
|
||||
nn.SiLU()
|
||||
)
|
||||
self.low_layer_norm = (
|
||||
nn.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
|
||||
)
|
||||
self.low_projection = nn.Linear(siglip_token_dims, hidden_size, bias=True)
|
||||
|
||||
def forward(self, siglip_outputs):
|
||||
"""
|
||||
Forward pass function
|
||||
|
||||
Args:
|
||||
siglip_outputs: Output from SigLIP model, containing hidden_states
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Concatenated multi-layer features with shape [bs, 3*style_token_nums, hidden_size]
|
||||
"""
|
||||
dtype = next(self.high_embedding_linear.parameters()).dtype
|
||||
|
||||
# Process high-level features (layer -2)
|
||||
high_embedding = self._process_layer_features(
|
||||
siglip_outputs.hidden_states[-2],
|
||||
self.high_embedding_linear,
|
||||
self.high_layer_norm,
|
||||
self.high_projection,
|
||||
dtype
|
||||
)
|
||||
|
||||
# Process mid-level features (layer -11)
|
||||
mid_embedding = self._process_layer_features(
|
||||
siglip_outputs.hidden_states[-11],
|
||||
self.mid_embedding_linear,
|
||||
self.mid_layer_norm,
|
||||
self.mid_projection,
|
||||
dtype
|
||||
)
|
||||
|
||||
# Process low-level features (layer -20)
|
||||
low_embedding = self._process_layer_features(
|
||||
siglip_outputs.hidden_states[-20],
|
||||
self.low_embedding_linear,
|
||||
self.low_layer_norm,
|
||||
self.low_projection,
|
||||
dtype
|
||||
)
|
||||
|
||||
# Concatenate features from all layers
|
||||
return torch.cat((high_embedding, mid_embedding, low_embedding), dim=1)
|
||||
|
||||
def _process_layer_features(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
embedding_linear: nn.Module,
|
||||
layer_norm: nn.Module,
|
||||
projection: nn.Module,
|
||||
dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Helper function to process features from a single layer
|
||||
|
||||
Args:
|
||||
hidden_states: Input hidden states [bs, seq_len, dim]
|
||||
embedding_linear: Embedding linear layer
|
||||
layer_norm: Layer normalization
|
||||
projection: Projection layer
|
||||
dtype: Target data type
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Processed features [bs, style_token_nums, hidden_size]
|
||||
"""
|
||||
# Transform dimensions: [bs, seq_len, dim] -> [bs, dim, seq_len] -> [bs, dim, style_token_nums] -> [bs, style_token_nums, dim]
|
||||
embedding = embedding_linear(
|
||||
hidden_states.to(dtype).transpose(1, 2)
|
||||
).transpose(1, 2)
|
||||
|
||||
# Apply layer normalization
|
||||
embedding = layer_norm(embedding)
|
||||
|
||||
# Project to target hidden space
|
||||
embedding = projection(embedding)
|
||||
|
||||
return embedding
|
||||
@@ -0,0 +1,398 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import math
|
||||
from typing import Literal, Optional
|
||||
from torch import Tensor
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import ExifTags, Image
|
||||
import torchvision.transforms.functional as TVF
|
||||
|
||||
from .modules.layers import (
|
||||
DoubleStreamBlockLoraProcessor,
|
||||
DoubleStreamBlockProcessor,
|
||||
SingleStreamBlockLoraProcessor,
|
||||
SingleStreamBlockProcessor,
|
||||
)
|
||||
from .sampling import denoise, get_noise, get_schedule, prepare_multi_ip, unpack
|
||||
from .util import (
|
||||
get_lora_rank,
|
||||
load_ae,
|
||||
load_checkpoint,
|
||||
load_clip,
|
||||
load_flow_model,
|
||||
load_flow_model_only_lora,
|
||||
load_t5,
|
||||
)
|
||||
|
||||
|
||||
def find_nearest_scale(image_h, image_w, predefined_scales):
|
||||
"""
|
||||
根据图片的高度和宽度，找到最近的预定义尺度。
|
||||
|
||||
:param image_h: 图片的高度
|
||||
:param image_w: 图片的宽度
|
||||
:param predefined_scales: 预定义尺度列表 [(h1, w1), (h2, w2), ...]
|
||||
:return: 最近的预定义尺度 (h, w)
|
||||
"""
|
||||
# 计算输入图片的长宽比
|
||||
image_ratio = image_h / image_w
|
||||
|
||||
# 初始化变量以存储最小差异和最近的尺度
|
||||
min_diff = float("inf")
|
||||
nearest_scale = None
|
||||
|
||||
# 遍历所有预定义尺度，找到与输入图片长宽比最接近的尺度
|
||||
for scale_h, scale_w in predefined_scales:
|
||||
predefined_ratio = scale_h / scale_w
|
||||
diff = abs(predefined_ratio - image_ratio)
|
||||
|
||||
if diff < min_diff:
|
||||
min_diff = diff
|
||||
nearest_scale = (scale_h, scale_w)
|
||||
|
||||
return nearest_scale
|
||||
|
||||
|
||||
def preprocess_ref(raw_image: Image.Image, long_size: int = 512, scale_ratio: int = 1):
|
||||
# 获取原始图像的宽度和高度
|
||||
image_w, image_h = raw_image.size
|
||||
if image_w == image_h and image_w == 16:
|
||||
return raw_image
|
||||
|
||||
# 计算长边和短边
|
||||
if image_w >= image_h:
|
||||
new_w = long_size
|
||||
new_h = int((long_size / image_w) * image_h)
|
||||
else:
|
||||
new_h = long_size
|
||||
new_w = int((long_size / image_h) * image_w)
|
||||
|
||||
# 按新的宽高进行等比例缩放
|
||||
raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
|
||||
|
||||
# 为了能让canny img进行scale
|
||||
scale_ratio = int(scale_ratio)
|
||||
target_w = new_w // (16 * scale_ratio) * (16 * scale_ratio)
|
||||
target_h = new_h // (16 * scale_ratio) * (16 * scale_ratio)
|
||||
|
||||
# 计算裁剪的起始坐标以实现中心裁剪
|
||||
left = (new_w - target_w) // 2
|
||||
top = (new_h - target_h) // 2
|
||||
right = left + target_w
|
||||
bottom = top + target_h
|
||||
|
||||
# 进行中心裁剪
|
||||
raw_image = raw_image.crop((left, top, right, bottom))
|
||||
|
||||
# 转换为 RGB 模式
|
||||
raw_image = raw_image.convert("RGB")
|
||||
return raw_image
|
||||
|
||||
|
||||
def resize_and_centercrop_image(image, target_height_ref1, target_width_ref1):
|
||||
target_height_ref1 = int(target_height_ref1 // 64 * 64)
|
||||
target_width_ref1 = int(target_width_ref1 // 64 * 64)
|
||||
h, w = image.shape[-2:]
|
||||
if h < target_height_ref1 or w < target_width_ref1:
|
||||
# 计算长宽比
|
||||
aspect_ratio = w / h
|
||||
if h < target_height_ref1:
|
||||
new_h = target_height_ref1
|
||||
new_w = new_h * aspect_ratio
|
||||
if new_w < target_width_ref1:
|
||||
new_w = target_width_ref1
|
||||
new_h = new_w / aspect_ratio
|
||||
else:
|
||||
new_w = target_width_ref1
|
||||
new_h = new_w / aspect_ratio
|
||||
if new_h < target_height_ref1:
|
||||
new_h = target_height_ref1
|
||||
new_w = new_h * aspect_ratio
|
||||
else:
|
||||
aspect_ratio = w / h
|
||||
tgt_aspect_ratio = target_width_ref1 / target_height_ref1
|
||||
if aspect_ratio > tgt_aspect_ratio:
|
||||
new_h = target_height_ref1
|
||||
new_w = new_h * aspect_ratio
|
||||
else:
|
||||
new_w = target_width_ref1
|
||||
new_h = new_w / aspect_ratio
|
||||
# 使用 TVF.resize 进行图像缩放
|
||||
image = TVF.resize(image, (math.ceil(new_h), math.ceil(new_w)))
|
||||
# 计算中心裁剪的参数
|
||||
top = (image.shape[-2] - target_height_ref1) // 2
|
||||
left = (image.shape[-1] - target_width_ref1) // 2
|
||||
# 使用 TVF.crop 进行中心裁剪
|
||||
image = TVF.crop(image, top, left, target_height_ref1, target_width_ref1)
|
||||
return image
|
||||
|
||||
|
||||
class USOPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
model_type: str,
|
||||
device: torch.device,
|
||||
offload: bool = False,
|
||||
only_lora: bool = False,
|
||||
lora_rank: int = 16,
|
||||
hf_download: bool = True,
|
||||
):
|
||||
self.device = device
|
||||
self.offload = offload
|
||||
self.model_type = model_type
|
||||
|
||||
print(f'----> model type is {model_type}({only_lora})')
|
||||
self.clip = load_clip(self.device)
|
||||
print('----> load clip completely')
|
||||
self.t5 = load_t5(self.device, max_length=512)
|
||||
print('----> load t5 completely')
|
||||
self.ae = load_ae(model_type, device="cpu" if offload else self.device)
|
||||
print('----> load ae completely')
|
||||
self.use_fp8 = "fp8" in model_type
|
||||
if only_lora:
|
||||
self.model = load_flow_model_only_lora(
|
||||
model_type,
|
||||
device="cpu" if offload else self.device,
|
||||
lora_rank=lora_rank,
|
||||
use_fp8=self.use_fp8,
|
||||
hf_download=hf_download,
|
||||
)
|
||||
else:
|
||||
self.model = load_flow_model(
|
||||
model_type, device="cpu" if offload else self.device
|
||||
)
|
||||
|
||||
def load_ckpt(self, ckpt_path):
|
||||
if ckpt_path is not None:
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
print("Loading checkpoint to replace old keys")
|
||||
# load_sft doesn't support torch.device
|
||||
if ckpt_path.endswith("safetensors"):
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
missing, unexpected = self.model.load_state_dict(
|
||||
sd, strict=False, assign=True
|
||||
)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path, map_location="cpu")
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace("module.", "")] = dit_state[k]
|
||||
missing, unexpected = self.model.load_state_dict(
|
||||
sd, strict=False, assign=True
|
||||
)
|
||||
self.model.to(str(self.device))
|
||||
print(f"missing keys: {missing}\n\n\n\n\nunexpected keys: {unexpected}")
|
||||
|
||||
def set_lora(
|
||||
self,
|
||||
local_path: str = None,
|
||||
repo_id: str = None,
|
||||
name: str = None,
|
||||
lora_weight: int = 0.7,
|
||||
):
|
||||
checkpoint = load_checkpoint(local_path, repo_id, name)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def set_lora_from_collection(
|
||||
self, lora_type: str = "realism", lora_weight: int = 0.7
|
||||
):
|
||||
checkpoint = load_checkpoint(
|
||||
None, self.hf_lora_collection, self.lora_types_to_names[lora_type]
|
||||
)
|
||||
self.update_model_with_lora(checkpoint, lora_weight)
|
||||
|
||||
def update_model_with_lora(self, checkpoint, lora_weight):
|
||||
rank = get_lora_rank(checkpoint)
|
||||
lora_attn_procs = {}
|
||||
|
||||
for name, _ in self.model.attn_processors.items():
|
||||
lora_state_dict = {}
|
||||
for k in checkpoint.keys():
|
||||
if name in k:
|
||||
lora_state_dict[k[len(name) + 1 :]] = checkpoint[k] * lora_weight
|
||||
|
||||
if len(lora_state_dict):
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockLoraProcessor(
|
||||
dim=3072, rank=rank
|
||||
)
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(
|
||||
dim=3072, rank=rank
|
||||
)
|
||||
lora_attn_procs[name].load_state_dict(lora_state_dict)
|
||||
lora_attn_procs[name].to(self.device)
|
||||
else:
|
||||
if name.startswith("single_blocks"):
|
||||
lora_attn_procs[name] = SingleStreamBlockProcessor()
|
||||
else:
|
||||
lora_attn_procs[name] = DoubleStreamBlockProcessor()
|
||||
|
||||
self.model.set_attn_processor(lora_attn_procs)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
guidance: float = 4,
|
||||
num_steps: int = 50,
|
||||
seed: int = 123456789,
|
||||
**kwargs,
|
||||
):
|
||||
width = 16 * (width // 16)
|
||||
height = 16 * (height // 16)
|
||||
|
||||
device_type = self.device if isinstance(self.device, str) else self.device.type
|
||||
with torch.autocast(
|
||||
enabled=self.use_fp8, device_type=device_type, dtype=torch.bfloat16
|
||||
):
|
||||
return self.forward(
|
||||
prompt, width, height, guidance, num_steps, seed, **kwargs
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def gradio_generate(
|
||||
self,
|
||||
prompt: str,
|
||||
image_prompt1: Image.Image,
|
||||
image_prompt2: Image.Image,
|
||||
image_prompt3: Image.Image,
|
||||
seed: int,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
guidance: float = 4,
|
||||
num_steps: int = 25,
|
||||
keep_size: bool = False,
|
||||
content_long_size: int = 512,
|
||||
):
|
||||
ref_content_imgs = [image_prompt1]
|
||||
ref_content_imgs = [img for img in ref_content_imgs if isinstance(img, Image.Image)]
|
||||
ref_content_imgs = [preprocess_ref(img, content_long_size) for img in ref_content_imgs]
|
||||
|
||||
ref_style_imgs = [image_prompt2, image_prompt3]
|
||||
ref_style_imgs = [img for img in ref_style_imgs if isinstance(img, Image.Image)]
|
||||
ref_style_imgs = [self.model.vision_encoder_processor(img, return_tensors="pt").to(self.device) for img in ref_style_imgs]
|
||||
|
||||
seed = seed if seed != -1 else torch.randint(0, 10**8, (1,)).item()
|
||||
|
||||
# whether keep input image size
|
||||
if keep_size and len(ref_content_imgs)>0:
|
||||
width, height = ref_content_imgs[0].size
|
||||
width, height = int(width * (1024 / content_long_size)), int(height * (1024 / content_long_size))
|
||||
img = self(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
guidance=guidance,
|
||||
num_steps=num_steps,
|
||||
seed=seed,
|
||||
ref_imgs=ref_content_imgs,
|
||||
siglip_inputs=ref_style_imgs,
|
||||
)
|
||||
|
||||
filename = f"output/gradio/{seed}_{prompt[:20]}.png"
|
||||
os.makedirs(os.path.dirname(filename), exist_ok=True)
|
||||
exif_data = Image.Exif()
|
||||
exif_data[ExifTags.Base.Make] = "USO"
|
||||
exif_data[ExifTags.Base.Model] = self.model_type
|
||||
info = f"{prompt=}, {seed=}, {width=}, {height=}, {guidance=}, {num_steps=}"
|
||||
exif_data[ExifTags.Base.ImageDescription] = info
|
||||
img.save(filename, format="png", exif=exif_data)
|
||||
return img, filename
|
||||
|
||||
@torch.inference_mode
|
||||
def forward(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
guidance: float,
|
||||
num_steps: int,
|
||||
seed: int,
|
||||
ref_imgs: list[Image.Image] | None = None,
|
||||
pe: Literal["d", "h", "w", "o"] = "d",
|
||||
siglip_inputs: list[Tensor] | None = None,
|
||||
**kwargs
|
||||
):
|
||||
|
||||
update_func = kwargs.get('update_func', lambda *args, **kwargs: None)
|
||||
x = get_noise(
|
||||
1, height, width, device=self.device, dtype=torch.bfloat16, seed=seed
|
||||
)
|
||||
timesteps = get_schedule(
|
||||
num_steps,
|
||||
(width // 8) * (height // 8) // (16 * 16),
|
||||
shift=True,
|
||||
)
|
||||
if self.offload:
|
||||
self.ae.encoder = self.ae.encoder.to(self.device)
|
||||
x_1_refs = [
|
||||
self.ae.encode(
|
||||
(TVF.to_tensor(ref_img) * 2.0 - 1.0)
|
||||
.unsqueeze(0)
|
||||
.to(self.device, torch.float32)
|
||||
).to(torch.bfloat16)
|
||||
for ref_img in ref_imgs
|
||||
]
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.ae.encoder)
|
||||
self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
|
||||
inp_cond = prepare_multi_ip(
|
||||
t5=self.t5,
|
||||
clip=self.clip,
|
||||
img=x,
|
||||
prompt=prompt,
|
||||
ref_imgs=x_1_refs,
|
||||
pe=pe,
|
||||
)
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.t5, self.clip)
|
||||
self.model = self.model.to(self.device)
|
||||
|
||||
x = denoise(
|
||||
self.model,
|
||||
**inp_cond,
|
||||
timesteps=timesteps,
|
||||
guidance=guidance,
|
||||
siglip_inputs=siglip_inputs,
|
||||
update_func=update_func,
|
||||
)
|
||||
|
||||
if self.offload:
|
||||
self.offload_model_to_cpu(self.model)
|
||||
self.ae.decoder.to(x.device)
|
||||
x = unpack(x.float(), height, width)
|
||||
x = self.ae.decode(x)
|
||||
self.offload_model_to_cpu(self.ae.decoder)
|
||||
|
||||
x1 = x.clamp(-1, 1)
|
||||
x1 = rearrange(x1[-1], "c h w -> h w c")
|
||||
output_img = Image.fromarray((127.5 * (x1 + 1.0)).cpu().byte().numpy())
|
||||
return output_img
|
||||
|
||||
def offload_model_to_cpu(self, *models):
|
||||
if not self.offload:
|
||||
return
|
||||
for model in models:
|
||||
model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,278 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import Tensor
|
||||
from tqdm import tqdm
|
||||
|
||||
from .model import Flux
|
||||
from .modules.conditioner import HFEmbedder
|
||||
|
||||
|
||||
def get_noise(
|
||||
num_samples: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
):
|
||||
return torch.randn(
|
||||
num_samples,
|
||||
16,
|
||||
# allow for packing
|
||||
2 * math.ceil(height / 16),
|
||||
2 * math.ceil(width / 16),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
generator=torch.Generator(device=device).manual_seed(seed),
|
||||
)
|
||||
|
||||
|
||||
def prepare(
|
||||
t5: HFEmbedder,
|
||||
clip: HFEmbedder,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_img: None | Tensor = None,
|
||||
pe: Literal["d", "h", "w", "o"] = "d",
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ["d", "h", "w", "o"]
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
_, _, ref_h, ref_w = ref_img.shape
|
||||
ref_img = rearrange(
|
||||
ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2
|
||||
)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids = torch.zeros(ref_h // 2, ref_w // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = h // 2 if pe in {"d", "h"} else 0
|
||||
w_offset = w // 2 if pe in {"d", "w"} else 0
|
||||
ref_img_ids[..., 1] = (
|
||||
ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None] + h_offset
|
||||
)
|
||||
ref_img_ids[..., 2] = (
|
||||
ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :] + w_offset
|
||||
)
|
||||
ref_img_ids = repeat(ref_img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
txt = t5(prompt)
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = clip(prompt)
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
if ref_img is not None:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"ref_img": ref_img,
|
||||
"ref_img_ids": ref_img_ids.to(img.device),
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
|
||||
|
||||
def prepare_multi_ip(
|
||||
t5: HFEmbedder,
|
||||
clip: HFEmbedder,
|
||||
img: Tensor,
|
||||
prompt: str | list[str],
|
||||
ref_imgs: list[Tensor] | None = None,
|
||||
pe: Literal["d", "h", "w", "o"] = "d",
|
||||
) -> dict[str, Tensor]:
|
||||
assert pe in ["d", "h", "w", "o"]
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
# tgt img
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
ref_img_ids = []
|
||||
ref_imgs_list = []
|
||||
|
||||
pe_shift_w, pe_shift_h = w // 2, h // 2
|
||||
for ref_img in ref_imgs:
|
||||
_, _, ref_h1, ref_w1 = ref_img.shape
|
||||
ref_img = rearrange(
|
||||
ref_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2
|
||||
)
|
||||
if ref_img.shape[0] == 1 and bs > 1:
|
||||
ref_img = repeat(ref_img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img_ids1 = torch.zeros(ref_h1 // 2, ref_w1 // 2, 3)
|
||||
# img id分别在宽高偏移各自最大值
|
||||
h_offset = pe_shift_h if pe in {"d", "h"} else 0
|
||||
w_offset = pe_shift_w if pe in {"d", "w"} else 0
|
||||
ref_img_ids1[..., 1] = (
|
||||
ref_img_ids1[..., 1] + torch.arange(ref_h1 // 2)[:, None] + h_offset
|
||||
)
|
||||
ref_img_ids1[..., 2] = (
|
||||
ref_img_ids1[..., 2] + torch.arange(ref_w1 // 2)[None, :] + w_offset
|
||||
)
|
||||
ref_img_ids1 = repeat(ref_img_ids1, "h w c -> b (h w) c", b=bs)
|
||||
ref_img_ids.append(ref_img_ids1)
|
||||
ref_imgs_list.append(ref_img)
|
||||
|
||||
# 更新pe shift
|
||||
pe_shift_h += ref_h1 // 2
|
||||
pe_shift_w += ref_w1 // 2
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
txt = t5(prompt)
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = clip(prompt)
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
return {
|
||||
"img": img,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"ref_img": tuple(ref_imgs_list),
|
||||
"ref_img_ids": [ref_img_id.to(img.device) for ref_img_id in ref_img_ids],
|
||||
"txt": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
"vec": vec.to(img.device),
|
||||
}
|
||||
|
||||
|
||||
def time_shift(mu: float, sigma: float, t: Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
def get_lin_function(
|
||||
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
||||
):
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
|
||||
def get_schedule(
|
||||
num_steps: int,
|
||||
image_seq_len: int,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
shift: bool = True,
|
||||
) -> list[float]:
|
||||
# extra step for zero
|
||||
timesteps = torch.linspace(1, 0, num_steps + 1)
|
||||
|
||||
# shifting the schedule to favor high timesteps for higher signal images
|
||||
if shift:
|
||||
# eastimate mu based on linear estimation between two points
|
||||
mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
|
||||
return timesteps.tolist()
|
||||
|
||||
|
||||
def denoise(
|
||||
model: Flux,
|
||||
# model input
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
vec: Tensor,
|
||||
# sampling parameters
|
||||
timesteps: list[float],
|
||||
guidance: float = 4.0,
|
||||
ref_img: Tensor = None,
|
||||
ref_img_ids: Tensor = None,
|
||||
siglip_inputs: list[Tensor] | None = None,
|
||||
#kiki
|
||||
update_func = None,
|
||||
):
|
||||
i = 0
|
||||
guidance_vec = torch.full(
|
||||
(img.shape[0],), guidance, device=img.device, dtype=img.dtype
|
||||
)
|
||||
for t_curr, t_prev in tqdm(
|
||||
zip(timesteps[:-1], timesteps[1:]), total=len(timesteps) - 1
|
||||
):
|
||||
if update_func is not None:
|
||||
update_func()
|
||||
# for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):
|
||||
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
||||
pred = model(
|
||||
img=img,
|
||||
img_ids=img_ids,
|
||||
ref_img=ref_img,
|
||||
ref_img_ids=ref_img_ids,
|
||||
txt=txt,
|
||||
txt_ids=txt_ids,
|
||||
y=vec,
|
||||
timesteps=t_vec,
|
||||
guidance=guidance_vec,
|
||||
siglip_inputs=siglip_inputs,
|
||||
)
|
||||
img = img + (t_prev - t_curr) * pred
|
||||
i += 1
|
||||
return img
|
||||
|
||||
|
||||
def unpack(x: Tensor, height: int, width: int) -> Tensor:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
||||
h=math.ceil(height / 16),
|
||||
w=math.ceil(width / 16),
|
||||
ph=2,
|
||||
pw=2,
|
||||
)
|
||||
@@ -0,0 +1,535 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
|
||||
# Copyright (c) 2024 Black Forest Labs and The XLabs-AI Team. All rights reserved.
|
||||
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import json
|
||||
import numpy as np
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
from .model import Flux, FluxParams
|
||||
from .modules.autoencoder import AutoEncoder, AutoEncoderParams
|
||||
from .modules.conditioner import HFEmbedder
|
||||
|
||||
import re
|
||||
from .modules.layers import (
|
||||
DoubleStreamBlockLoraProcessor,
|
||||
SingleStreamBlockLoraProcessor,
|
||||
)
|
||||
|
||||
import os
|
||||
try:
|
||||
import folder_paths
|
||||
print('run in comfyui')
|
||||
except:
|
||||
print('not run in comfyui')
|
||||
from types import SimpleNamespace
|
||||
folder_paths = SimpleNamespace()
|
||||
folder_paths.models_dir = '/workspace/comfyui/models/'
|
||||
|
||||
|
||||
def load_model(ckpt, device="cpu"):
|
||||
if ckpt.endswith("safetensors"):
|
||||
from safetensors import safe_open
|
||||
|
||||
pl_sd = {}
|
||||
with safe_open(ckpt, framework="pt", device=device) as f:
|
||||
for k in f.keys():
|
||||
pl_sd[k] = f.get_tensor(k)
|
||||
else:
|
||||
pl_sd = torch.load(ckpt, map_location=device)
|
||||
return pl_sd
|
||||
|
||||
|
||||
def load_safetensors(path):
|
||||
tensors = {}
|
||||
with safe_open(path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
tensors[key] = f.get_tensor(key)
|
||||
return tensors
|
||||
|
||||
|
||||
def get_lora_rank(checkpoint):
|
||||
for k in checkpoint.keys():
|
||||
if k.endswith(".down.weight"):
|
||||
return checkpoint[k].shape[0]
|
||||
|
||||
|
||||
def load_checkpoint(local_path, repo_id, name):
|
||||
if local_path is not None:
|
||||
if ".safetensors" in local_path:
|
||||
print(f"Loading .safetensors checkpoint from {local_path}")
|
||||
checkpoint = load_safetensors(local_path)
|
||||
else:
|
||||
print(f"Loading checkpoint from {local_path}")
|
||||
checkpoint = torch.load(local_path, map_location="cpu")
|
||||
elif repo_id is not None and name is not None:
|
||||
print(f"Loading checkpoint {name} from repo id {repo_id}")
|
||||
checkpoint = load_from_repo_id(repo_id, name)
|
||||
else:
|
||||
raise ValueError(
|
||||
"LOADING ERROR: you must specify local_path or repo_id with name in HF to download"
|
||||
)
|
||||
return checkpoint
|
||||
|
||||
|
||||
def c_crop(image):
|
||||
width, height = image.size
|
||||
new_size = min(width, height)
|
||||
left = (width - new_size) / 2
|
||||
top = (height - new_size) / 2
|
||||
right = (width + new_size) / 2
|
||||
bottom = (height + new_size) / 2
|
||||
return image.crop((left, top, right, bottom))
|
||||
|
||||
|
||||
def pad64(x):
|
||||
return int(np.ceil(float(x) / 64.0) * 64 - x)
|
||||
|
||||
|
||||
def HWC3(x):
|
||||
assert x.dtype == np.uint8
|
||||
if x.ndim == 2:
|
||||
x = x[:, :, None]
|
||||
assert x.ndim == 3
|
||||
H, W, C = x.shape
|
||||
assert C == 1 or C == 3 or C == 4
|
||||
if C == 3:
|
||||
return x
|
||||
if C == 1:
|
||||
return np.concatenate([x, x, x], axis=2)
|
||||
if C == 4:
|
||||
color = x[:, :, 0:3].astype(np.float32)
|
||||
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
||||
y = color * alpha + 255.0 * (1.0 - alpha)
|
||||
y = y.clip(0, 255).astype(np.uint8)
|
||||
return y
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
params: FluxParams
|
||||
ae_params: AutoEncoderParams
|
||||
ckpt_path: str | None
|
||||
ae_path: str | None
|
||||
repo_id: str | None
|
||||
repo_flow: str | None
|
||||
repo_ae: str | None
|
||||
repo_id_ae: str | None
|
||||
|
||||
|
||||
configs = {
|
||||
"flux-dev": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-dev",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-dev.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_DEV"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-dev-fp8": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-dev",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-dev.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_DEV_FP8"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-krea-dev": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-Krea-dev",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-Krea-dev",
|
||||
repo_flow="flux1-krea-dev.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_KREA_DEV"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-schnell": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-schnell",
|
||||
repo_id_ae="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-schnell.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_SCHNELL"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=False,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
|
||||
if len(missing) > 0 and len(unexpected) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
print("\n" + "-" * 79 + "\n")
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
elif len(missing) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
elif len(unexpected) > 0:
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
|
||||
|
||||
def load_from_repo_id(repo_id, checkpoint_name):
|
||||
ckpt_path = hf_hub_download(repo_id, checkpoint_name)
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
return sd
|
||||
|
||||
|
||||
def load_flow_model(
|
||||
name: str, device: str | torch.device = "cuda", hf_download: bool = True
|
||||
):
|
||||
# Loading Flux
|
||||
print("Init model")
|
||||
ckpt_path = configs[name].ckpt_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
|
||||
# with torch.device("meta" if ckpt_path is not None else device):
|
||||
with torch.device(device):
|
||||
model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
if ckpt_path is not None:
|
||||
print("Loading main checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_model(ckpt_path, device="cpu")
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return model.to(str(device))
|
||||
|
||||
|
||||
def load_flow_model_only_lora(
|
||||
name: str,
|
||||
device: str | torch.device = "cuda",
|
||||
hf_download: bool = True,
|
||||
lora_rank: int = 16,
|
||||
use_fp8: bool = False,
|
||||
):
|
||||
# Loading Flux
|
||||
# ckpt_path = configs[name].ckpt_path
|
||||
# kiki
|
||||
ckpt_path = os.path.join(folder_paths.models_dir, 'diffusers', 'FLUX.1-dev', 'flux1-dev.safetensors')
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
):
|
||||
ckpt_path = hf_hub_download(
|
||||
configs[name].repo_id, configs[name].repo_flow.replace("sft", "safetensors")
|
||||
)
|
||||
|
||||
# if hf_download:
|
||||
# try:
|
||||
# lora_ckpt_path = hf_hub_download(
|
||||
# "bytedance-research/USO", "uso_flux_v1.0/dit_lora.safetensors"
|
||||
# )
|
||||
# except Exception as e:
|
||||
# print(f"Failed to download lora checkpoint: {e}")
|
||||
# print("Trying to load lora from local")
|
||||
# lora_ckpt_path = os.environ.get("LORA", None)
|
||||
# try:
|
||||
# proj_ckpt_path = hf_hub_download(
|
||||
# "bytedance-research/USO", "uso_flux_v1.0/projector.safetensors"
|
||||
# )
|
||||
# except Exception as e:
|
||||
# print(f"Failed to download projection_model checkpoint: {e}")
|
||||
# print("Trying to load projection_model from local")
|
||||
# proj_ckpt_path = os.environ.get("PROJECTION_MODEL", None)
|
||||
# else:
|
||||
# lora_ckpt_path = os.environ.get("LORA", None)
|
||||
# proj_ckpt_path = os.environ.get("PROJECTION_MODEL", None)
|
||||
# print(lora_ckpt_path)
|
||||
# print(proj_ckpt_path)
|
||||
base_ckpt_path = os.path.join(folder_paths.models_dir, 'uso', 'uso_flux_v1.0')
|
||||
lora_ckpt_path = os.path.join(base_ckpt_path, 'dit_lora.safetensors')
|
||||
proj_ckpt_path = os.path.join(base_ckpt_path, 'projector.safetensors')
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = Flux(configs[name].params)
|
||||
|
||||
model = set_lora(
|
||||
model, lora_rank, device="meta" if lora_ckpt_path is not None else device
|
||||
)
|
||||
|
||||
if ckpt_path is not None:
|
||||
print(f"Loading lora from {lora_ckpt_path}")
|
||||
lora_sd = (
|
||||
load_sft(lora_ckpt_path, device=str(device))
|
||||
if lora_ckpt_path.endswith("safetensors")
|
||||
else torch.load(lora_ckpt_path, map_location="cpu")
|
||||
)
|
||||
proj_sd = (
|
||||
load_sft(proj_ckpt_path, device=str(device))
|
||||
if proj_ckpt_path.endswith("safetensors")
|
||||
else torch.load(proj_ckpt_path, map_location="cpu")
|
||||
)
|
||||
lora_sd.update(proj_sd)
|
||||
|
||||
print("Loading main checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
|
||||
if ckpt_path.endswith("safetensors"):
|
||||
if use_fp8:
|
||||
print(
|
||||
"####\n"
|
||||
"We are in fp8 mode right now, since the fp8 checkpoint of XLabs-AI/flux-dev-fp8 seems broken\n"
|
||||
"we convert the fp8 checkpoint on flight from bf16 checkpoint\n"
|
||||
"If your storage is constrained"
|
||||
"you can save the fp8 checkpoint and replace the bf16 checkpoint by yourself\n"
|
||||
)
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
sd = {
|
||||
k: v.to(dtype=torch.float8_e4m3fn, device=device)
|
||||
for k, v in sd.items()
|
||||
}
|
||||
else:
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
else:
|
||||
dit_state = torch.load(ckpt_path, map_location="cpu")
|
||||
sd = {}
|
||||
for k in dit_state.keys():
|
||||
sd[k.replace("module.", "")] = dit_state[k]
|
||||
sd.update(lora_sd)
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
model.to(str(device))
|
||||
print_load_warning(missing, unexpected)
|
||||
return model
|
||||
|
||||
|
||||
def set_lora(
|
||||
model: Flux,
|
||||
lora_rank: int,
|
||||
double_blocks_indices: list[int] | None = None,
|
||||
single_blocks_indices: list[int] | None = None,
|
||||
device: str | torch.device = "cpu",
|
||||
) -> Flux:
|
||||
double_blocks_indices = (
|
||||
list(range(model.params.depth))
|
||||
if double_blocks_indices is None
|
||||
else double_blocks_indices
|
||||
)
|
||||
single_blocks_indices = (
|
||||
list(range(model.params.depth_single_blocks))
|
||||
if single_blocks_indices is None
|
||||
else single_blocks_indices
|
||||
)
|
||||
|
||||
lora_attn_procs = {}
|
||||
with torch.device(device):
|
||||
for name, attn_processor in model.attn_processors.items():
|
||||
match = re.search(r"\.(\d+)\.", name)
|
||||
if match:
|
||||
layer_index = int(match.group(1))
|
||||
|
||||
if (
|
||||
name.startswith("double_blocks")
|
||||
and layer_index in double_blocks_indices
|
||||
):
|
||||
lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(
|
||||
dim=model.params.hidden_size, rank=lora_rank
|
||||
)
|
||||
elif (
|
||||
name.startswith("single_blocks")
|
||||
and layer_index in single_blocks_indices
|
||||
):
|
||||
lora_attn_procs[name] = SingleStreamBlockLoraProcessor(
|
||||
dim=model.params.hidden_size, rank=lora_rank
|
||||
)
|
||||
else:
|
||||
lora_attn_procs[name] = attn_processor
|
||||
model.set_attn_processor(lora_attn_procs)
|
||||
return model
|
||||
|
||||
|
||||
def load_flow_model_quintized(
|
||||
name: str, device: str | torch.device = "cuda", hf_download: bool = True
|
||||
):
|
||||
# Loading Flux
|
||||
from optimum.quanto import requantize
|
||||
|
||||
print("Init model")
|
||||
ckpt_path = configs[name].ckpt_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
json_path = hf_hub_download(configs[name].repo_id, "flux_dev_quantization_map.json")
|
||||
|
||||
model = Flux(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
print("Loading checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
sd = {k: v.to(dtype=torch.float8_e4m3fn, device=device) for k, v in sd.items()}
|
||||
model.load_state_dict(sd, assign=True)
|
||||
return model
|
||||
with open(json_path, "r") as f:
|
||||
quantization_map = json.load(f)
|
||||
print("Start a quantization process...")
|
||||
requantize(model, sd, quantization_map, device=device)
|
||||
print("Model is quantized!")
|
||||
return model
|
||||
|
||||
|
||||
def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
|
||||
# max length 64, 128, 256 and 512 should work (if your sequence is short enough)
|
||||
#version = os.environ.get("T5", "xlabs-ai/xflux_text_encoders")
|
||||
# version = '/workspace/comfyui/models/clip/xflux_text_encoders'
|
||||
version = os.path.join(folder_paths.models_dir, 'clip', 'xflux_text_encoders')
|
||||
return HFEmbedder(version, max_length=max_length, torch_dtype=torch.bfloat16).to(
|
||||
device
|
||||
)
|
||||
|
||||
|
||||
def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
|
||||
# version = os.environ.get("CLIP", "openai/clip-vit-large-patch14")
|
||||
#kiki
|
||||
# version = '/workspace/comfyui/models/clip_vision/clip-vit-large-patch14'
|
||||
version = os.path.join(folder_paths.models_dir, 'clip_vision', 'clip-vit-large-patch14')
|
||||
return HFEmbedder(version, max_length=77, torch_dtype=torch.bfloat16, is_clip=True).to(device)
|
||||
|
||||
|
||||
def load_ae(
|
||||
name: str, device: str | torch.device = "cuda", hf_download: bool = True
|
||||
) -> AutoEncoder:
|
||||
# ckpt_path = configs[name].ae_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_ae is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id_ae, configs[name].repo_ae)
|
||||
#kiki
|
||||
ckpt_path = os.path.join(folder_paths.models_dir, 'diffusers', 'FLUX.1-dev', 'ae.safetensors')
|
||||
|
||||
# Loading the autoencoder
|
||||
print("Init AE")
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
ae = AutoEncoder(configs[name].ae_params)
|
||||
|
||||
if ckpt_path is not None:
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
print_load_warning(missing, unexpected)
|
||||
return ae
|
||||
Reference in New Issue
Block a user