feat: init
This commit is contained in:
raykindle
+34
@@ -0,0 +1,34 @@
|
||||
# Python cache files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# Distribution / packaging
|
||||
dist/
|
||||
build/
|
||||
*.egg-info/
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Models and data
|
||||
models/
|
||||
outputs/
|
||||
|
||||
# IDE specific files
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# OS specific files
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
refer_src/
|
||||
@@ -0,0 +1,135 @@
|
||||
# ComfyUI_Step1X-Edit
|
||||
|
||||
[English](README.md) | [中文文档](README_CN.md)
|
||||
|
||||
This custom node integrates the [Step1X-Edit](https://github.com/stepfun-ai/Step1X-Edit) image editing model into [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Step1X-Edit is a state-of-the-art image editing model that processes a reference image and user's editing instruction to generate a new image.
|
||||
|
||||
## Features
|
||||
|
||||
- Load and use Step1X-Edit models in ComfyUI
|
||||
- Edit images with natural language instructions
|
||||
- Control the editing process with parameters like steps and guidance scale
|
||||
|
||||
## Examples
|
||||
|
||||
Here are some examples of what you can achieve with ComfyUI_Step1X-Edit:
|
||||
|
||||
| Example 1 | Example 2 |
|
||||
|:-----------:|:------------:|
|
||||
|  |  |
|
||||
|
||||
- Example 1: The corresponding prompt was: "给这个女生的脖子上戴一个带有红宝石的吊坠。" ("Add pendant with a ruby around this girl's neck.")
|
||||
- Example 2: The corresponding prompt was: "让她哭。" ("Let her cry.")
|
||||
|
||||
You can find the example workflow in the [examples directory](examples/step1x_edit_example.json). Load it directly into ComfyUI to see how it works.
|
||||
|
||||
## Installation
|
||||
|
||||
1. Clone this repository into your ComfyUI's `custom_nodes` directory:
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/raykindle/ComfyUI_Step1X-Edit.git
|
||||
```
|
||||
|
||||
2. Install the required dependencies:
|
||||
Step 1: Install ComfyUI_Step1X-Edit
|
||||
```bash
|
||||
cd ComfyUI_Step1X-Edit
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Step 2: Install [`flash-attn`](https://github.com/Dao-AILab/flash-attention), here we provide a script to help find the pre-built wheel suitable for your system.
|
||||
|
||||
```bash
|
||||
python utils/get_flash_attn.py
|
||||
```
|
||||
|
||||
The script will generate a wheel name like `flash_attn-2.7.2.post1+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl`, which could be found in [the release page of flash-attn](https://github.com/Dao-AILab/flash-attention/releases).
|
||||
|
||||
Then you can download the corresponding pre-built wheel and install it following the instructions in [`flash-attn`](https://github.com/Dao-AILab/flash-attention).
|
||||
|
||||
|
||||
3. Download the Step1X-Edit-FP8 model:
|
||||
- Step1X-Edit diffusion model: Download `step1x-edit-i1258-FP8.safetensors` from [HuggingFace](https://huggingface.co/meimeilook/Step1X-Edit-FP8/tree/main) and place it in ComfyUI's `models/diffusion_models` directory
|
||||
- Step1X-Edit VAE: Download `vae.safetensors` from [HuggingFace](https://huggingface.co/meimeilook/Step1X-Edit-FP8/tree/main) and place it in ComfyUI's `models/vae` directory
|
||||
- Qwen2.5-VL model: Download [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct/tree/main) and place it in ComfyUI's `models/text_encoders/Qwen2.5-VL-7B-Instruct` directory
|
||||
|
||||
## Usage
|
||||
|
||||
1. Start ComfyUI and create a new workflow.
|
||||
2. Add the "Step1X-Edit Model Loader" node to your workflow.
|
||||
3. Configure the model parameters:
|
||||
- Select `step1x-edit-i1258-FP8.safetensors` as the diffusion model
|
||||
- Select `vae.safetensors` as the VAE
|
||||
- Set "Qwen2.5-VL-7B-Instruct" as the text encoder
|
||||
- Set additional parameters (dtype, quantized, offload) as needed
|
||||
4. Connect an "Step1X-Edit Generate" node to the model node.
|
||||
5. Provide an input image and an editing prompt.
|
||||
6. Run the workflow to generate edited images.
|
||||
|
||||
## Example Workflow
|
||||
|
||||
Here's a simple example workflow:
|
||||
|
||||
1. Load Image → Step1X-Edit Generate
|
||||
2. Step1X-Edit Model Loader → Step1X-Edit Generate
|
||||
3. Step1X-Edit Generate → Preview Image
|
||||
|
||||
## Parameters
|
||||
|
||||
### Step1X-Edit Model Loader
|
||||
|
||||
- `diffusion_model`: The Step1X-Edit diffusion model file (select from the diffusion_models dropdown)
|
||||
- `vae`: The Step1X-Edit VAE file (select from the vae dropdown)
|
||||
- `text_encoder`: The path to the Qwen2.5-VL model directory name (e.g., "Qwen2.5-VL-7B-Instruct")
|
||||
- `dtype`: Model precision (bfloat16, float16, or float32)
|
||||
- `quantized`: Whether to use FP8 quantized weights (true recommended)
|
||||
- `offload`: Whether to offload models to CPU when not in use
|
||||
|
||||
### Step1X-Edit Generate
|
||||
|
||||
- `model`: The Step1X-Edit model bundle
|
||||
- `image`: The input image to edit
|
||||
- `prompt`: Text instructions describing the desired edit
|
||||
- `negative_prompt`: Text describing what to avoid
|
||||
- `steps`: Number of denoising steps (more steps = better quality but slower)
|
||||
- `cfg_scale`: Guidance scale (how closely to follow the prompt)
|
||||
- `image_size`: Size of the output image (512 recommended)
|
||||
- `seed`: Random seed for reproducibility
|
||||
|
||||
## Memory Requirements
|
||||
|
||||
The Step1X-Edit model requires significant GPU memory:
|
||||
|
||||
| Model Version | Peak GPU Memory (512px / 768px / 1024px) |
|
||||
|:-------------:|:----------------------------------------:|
|
||||
| Standard | 42.5GB / 46.5GB / 49.8GB |
|
||||
| FP8 Quantized | 31GB / 31.5GB / 34GB |
|
||||
| With Offload | 25.9GB / 27.3GB / 29.1GB |
|
||||
| FP8 + Offload | 18GB / 18GB / 18GB |
|
||||
|
||||
For lower memory usage, enable the `quantized` and/or `offload` options in the Model Loader node.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- If you encounter CUDA out of memory errors, try:
|
||||
- Enabling the `offload` option
|
||||
- Using the FP8 quantized model
|
||||
- Reducing the image size
|
||||
- Closing other GPU-intensive applications
|
||||
- If you get errors about missing files:
|
||||
- Make sure your model paths are correct
|
||||
- The diffusion model should be in `models/diffusion_models`
|
||||
- The VAE should be in `models/vae`
|
||||
- The text encoder should be in `models/text_encoders/Qwen2.5-VL-7B-Instruct`
|
||||
- If you get import errors, ensure all dependencies are installed correctly
|
||||
|
||||
## License
|
||||
|
||||
This custom node is licensed under the Apache License 2.0, consistent with the original Step1X-Edit repository.
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
- The Step1X-Edit team for creating the original model
|
||||
- ComfyUI for the extensible UI framework
|
||||
+135
@@ -0,0 +1,135 @@
|
||||
# ComfyUI_Step1X-Edit
|
||||
|
||||
[English](README.md) | [中文文档](README_CN.md)
|
||||
|
||||
此自定义节点将 [Step1X-Edit](https://github.com/stepfun-ai/Step1X-Edit) 图像编辑模型集成到 [ComfyUI](https://github.com/comfyanonymous/ComfyUI) 中。Step1X-Edit 是一个先进的图像编辑模型，它接收参考图像和用户的编辑指令，生成新的图像。
|
||||
|
||||
## 功能特点
|
||||
|
||||
- 在 ComfyUI 中加载和使用 Step1X-Edit 模型
|
||||
- 使用自然语言指令编辑图像
|
||||
- 通过步数和引导系数等参数控制编辑过程
|
||||
|
||||
## 示例展示
|
||||
|
||||
以下是使用 ComfyUI_Step1X-Edit 可以实现的效果示例：
|
||||
|
||||
| 示例 1 | 示例 2 |
|
||||
|:-----------:|:------------:|
|
||||
|  |  |
|
||||
|
||||
- 示例 1：对应的提示词是："给这个女生的脖子上戴一个带有红宝石的吊坠。" ("Add pendant with a ruby around this girl's neck.")
|
||||
- 示例 2：对应的提示词是："让她哭。" ("Let her cry.")
|
||||
|
||||
您可以在[示例目录](examples/step1x_edit_example.json)中找到示例工作流。直接将其加载到 ComfyUI 中即可查看其工作原理。
|
||||
|
||||
## 安装方法
|
||||
|
||||
1. 将此仓库克隆到 ComfyUI 的 `custom_nodes` 目录中：
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/raykindle/ComfyUI_Step1X-Edit.git
|
||||
```
|
||||
|
||||
2. 安装所需依赖：
|
||||
步骤 1: 安装 ComfyUI_Step1X-Edit
|
||||
```bash
|
||||
cd ComfyUI_Step1X-Edit
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
步骤 2: 安装 [`flash-attn`](https://github.com/Dao-AILab/flash-attention)，我们提供了一个脚本，帮助您找到适合您系统的预编译轮子。
|
||||
|
||||
```bash
|
||||
python utils/get_flash_attn.py
|
||||
```
|
||||
|
||||
该脚本将生成一个轮子名称，例如：`flash_attn-2.7.2.post1+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl`，该名称可在 [flash-attn 的发布页面](https://github.com/Dao-AILab/flash-attention/releases) 中找到。
|
||||
|
||||
然后，您可以下载相应的预编译轮子，并按照 [`flash-attn`](https://github.com/Dao-AILab/flash-attention) 中的说明进行安装。
|
||||
|
||||
|
||||
3. 下载 Step1X-Edit-FP8 模型：
|
||||
- Step1X-Edit 扩散模型：从 [HuggingFace](https://huggingface.co/meimeilook/Step1X-Edit-FP8/tree/main) 下载 `step1x-edit-i1258-FP8.safetensors` 并放置在 ComfyUI 的 `models/diffusion_models` 目录中
|
||||
- Step1X-Edit VAE：从 [HuggingFace](https://huggingface.co/meimeilook/Step1X-Edit-FP8/tree/main) 下载 `vae.safetensors` 并放置在 ComfyUI 的 `models/vae` 目录中
|
||||
- Qwen2.5-VL 模型：下载 [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct/tree/main) 并放置在 ComfyUI 的 `models/text_encoders/Qwen2.5-VL-7B-Instruct` 目录中
|
||||
|
||||
## 使用方法
|
||||
|
||||
1. 启动 ComfyUI 并创建新的工作流。
|
||||
2. 添加 "Step1X-Edit Model Loader"（Step1X-Edit 模型加载器）节点到工作流中。
|
||||
3. 配置模型参数：
|
||||
- 选择 `step1x-edit-i1258-FP8.safetensors` 作为扩散模型
|
||||
- 选择 `vae.safetensors` 作为 VAE
|
||||
- 设置 "Qwen2.5-VL-7B-Instruct" 作为文本编码器
|
||||
- 根据需要设置其他参数（dtype、quantized、offload）
|
||||
4. 连接 "Step1X-Edit Generate"（Step1X-Edit 生成）节点到模型节点。
|
||||
5. 提供输入图像和编辑提示。
|
||||
6. 运行工作流生成编辑后的图像。
|
||||
|
||||
## 工作流示例
|
||||
|
||||
以下是一个简单的工作流示例：
|
||||
|
||||
1. Load Image（加载图像）→ Step1X-Edit Generate（Step1X-Edit 生成）
|
||||
2. Step1X-Edit Model Loader（Step1X-Edit 模型加载器）→ Step1X-Edit Generate（Step1X-Edit 生成）
|
||||
3. Step1X-Edit Generate（Step1X-Edit 生成）→ Preview Image（预览图像）
|
||||
|
||||
## 参数说明
|
||||
|
||||
### Step1X-Edit Model Loader（Step1X-Edit 模型加载器）
|
||||
|
||||
- `diffusion_model`：Step1X-Edit 扩散模型文件（从 diffusion_models 下拉菜单中选择）
|
||||
- `vae`：Step1X-Edit VAE 文件（从 vae 下拉菜单中选择）
|
||||
- `text_encoder`：Qwen2.5-VL 模型目录名称（例如 "Qwen2.5-VL-7B-Instruct"）
|
||||
- `dtype`：模型精度（bfloat16、float16 或 float32）
|
||||
- `quantized`：是否使用 FP8 量化权重（推荐 开启）
|
||||
- `offload`：在不使用时是否将模型卸载到 CPU
|
||||
|
||||
### Step1X-Edit Generate（Step1X-Edit 生成）
|
||||
|
||||
- `model_bundle`：Step1X-Edit 模型包
|
||||
- `input_image`：要编辑的输入图像
|
||||
- `prompt`：描述所需编辑的文本指令
|
||||
- `negative_prompt`：描述要避免的内容的文本
|
||||
- `num_steps`：去噪步数（更多步数 = 更好的质量但更慢）
|
||||
- `cfg_guidance`：引导系数（控制对提示的遵循程度）
|
||||
- `size_level`：输出图像大小（推荐 512）
|
||||
- `seed`：随机种子，用于可重现性
|
||||
|
||||
## 内存需求
|
||||
|
||||
Step1X-Edit 模型需要相当大的 GPU 内存：
|
||||
|
||||
| 模型版本 | 峰值 GPU 内存（512px / 768px / 1024px） |
|
||||
|:-------------:|:----------------------------------------:|
|
||||
| 标准版 | 42.5GB / 46.5GB / 49.8GB |
|
||||
| FP8 量化版 | 31GB / 31.5GB / 34GB |
|
||||
| 使用 Offload | 25.9GB / 27.3GB / 29.1GB |
|
||||
| FP8 + Offload | 18GB / 18GB / 18GB |
|
||||
|
||||
为了降低内存使用，请在模型加载器节点中启用 `quantized` 和/或 `offload` 选项。
|
||||
|
||||
## 故障排除
|
||||
|
||||
- 如果遇到 CUDA 内存不足错误，请尝试：
|
||||
- 启用 `offload` 选项
|
||||
- 使用 FP8 量化模型
|
||||
- 减小图像大小
|
||||
- 关闭其他占用 GPU 的应用程序
|
||||
- 如果遇到文件缺失错误：
|
||||
- 确保模型路径正确
|
||||
- 扩散模型应位于 `models/diffusion_models`
|
||||
- VAE 应位于 `models/vae`
|
||||
- 文本编码器应位于 `models/text_encoders/Qwen2.5-VL-7B-Instruct`
|
||||
- 如果遇到导入错误，请确保所有依赖项都正确安装
|
||||
|
||||
## 许可证
|
||||
|
||||
此自定义节点采用 Apache License 2.0 许可，与原始 Step1X-Edit 仓库一致。
|
||||
|
||||
## 致谢
|
||||
|
||||
- 感谢 Step1X-Edit 团队创建原始模型
|
||||
- 感谢 ComfyUI 提供可扩展的 UI 框架
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
from .nodes.step1x_edit_nodes import Step1XEditModelLoader, Step1XEditGenerateNode
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Step1XEditModelLoader": Step1XEditModelLoader,
|
||||
"Step1XEditGenerate": Step1XEditGenerateNode
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Step1XEditModelLoader": "Step1X-Edit Model Loader",
|
||||
"Step1XEditGenerate": "Step1X-Edit Generate"
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 238 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 245 KiB |
@@ -0,0 +1,196 @@
|
||||
{
|
||||
"last_node_id": 10,
|
||||
"last_link_id": 17,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 2,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
745.9241333007812,
|
||||
676.44091796875
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
299.4796447753906
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
15
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"0001.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
961.9817504882812,
|
||||
358.0769958496094
|
||||
],
|
||||
"size": [
|
||||
397.47607421875,
|
||||
618.7960815429688
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 17
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "Step1XEditGenerate",
|
||||
"pos": [
|
||||
745.9320068359375,
|
||||
358.9280700683594
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
282
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model_bundle",
|
||||
"type": "STEP1X_MODEL_BUNDLE",
|
||||
"link": 16
|
||||
},
|
||||
{
|
||||
"name": "input_image",
|
||||
"type": "IMAGE",
|
||||
"link": 15
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
17
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "Step1XEditGenerate"
|
||||
},
|
||||
"widgets_values": [
|
||||
"让她哭。",
|
||||
"",
|
||||
10,
|
||||
6,
|
||||
318,
|
||||
"randomize",
|
||||
768
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"type": "Step1XEditModelLoader",
|
||||
"pos": [
|
||||
744.1864624023438,
|
||||
127.32762908935547
|
||||
],
|
||||
"size": [
|
||||
617.0538330078125,
|
||||
198.91787719726562
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "STEP1X_MODEL_BUNDLE",
|
||||
"type": "STEP1X_MODEL_BUNDLE",
|
||||
"links": [
|
||||
16
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "Step1XEditModelLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"step1x-edit-i1258-FP8.safetensors",
|
||||
"vae.safetensors",
|
||||
"Qwen2.5-VL-7B-Instruct",
|
||||
"bfloat16",
|
||||
true,
|
||||
false
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
15,
|
||||
2,
|
||||
0,
|
||||
9,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
16,
|
||||
10,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"STEP1X_MODEL_BUNDLE"
|
||||
],
|
||||
[
|
||||
17,
|
||||
9,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.7283845674106205,
|
||||
"offset": [
|
||||
-159.39467322432967,
|
||||
-30.630752842927603
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
try:
|
||||
import flash_attn
|
||||
from flash_attn.flash_attn_interface import (
|
||||
_flash_attn_forward,
|
||||
flash_attn_func,
|
||||
flash_attn_varlen_func,
|
||||
)
|
||||
except ImportError:
|
||||
flash_attn = None
|
||||
flash_attn_varlen_func = None
|
||||
_flash_attn_forward = None
|
||||
flash_attn_func = None
|
||||
|
||||
MEMORY_LAYOUT = {
|
||||
# flash模式:
|
||||
# 预处理: 输入 [batch_size, seq_len, num_heads, head_dim]
|
||||
# 后处理: 保持形状不变
|
||||
"flash": (
|
||||
lambda x: x, # 保持形状
|
||||
lambda x: x, # 保持形状
|
||||
),
|
||||
# torch/vanilla模式:
|
||||
# 预处理: 交换序列和注意力头的维度 [B,S,A,D] -> [B,A,S,D]
|
||||
# 后处理: 交换回原始维度 [B,A,S,D] -> [B,S,A,D]
|
||||
"torch": (
|
||||
lambda x: x.transpose(1, 2), # (B,S,A,D) -> (B,A,S,D)
|
||||
lambda x: x.transpose(1, 2), # (B,A,S,D) -> (B,S,A,D)
|
||||
),
|
||||
"vanilla": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mode="flash",
|
||||
drop_rate=0,
|
||||
attn_mask=None,
|
||||
causal=False,
|
||||
):
|
||||
"""
|
||||
执行QKV自注意力计算
|
||||
|
||||
Args:
|
||||
q (torch.Tensor): 查询张量，形状 [batch_size, seq_len, num_heads, head_dim]
|
||||
k (torch.Tensor): 键张量，形状 [batch_size, seq_len_kv, num_heads, head_dim]
|
||||
v (torch.Tensor): 值张量，形状 [batch_size, seq_len_kv, num_heads, head_dim]
|
||||
mode (str): 注意力模式，可选 'flash', 'torch', 'vanilla'
|
||||
drop_rate (float): 注意力矩阵的dropout概率
|
||||
attn_mask (torch.Tensor): 注意力掩码，形状根据模式不同而变化
|
||||
causal (bool): 是否使用因果注意力（仅关注前面位置）
|
||||
|
||||
Returns:
|
||||
torch.Tensor: 注意力输出，形状 [batch_size, seq_len, num_heads * head_dim]
|
||||
"""
|
||||
# 获取预处理和后处理函数
|
||||
pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
|
||||
|
||||
# 应用预处理变换
|
||||
q = pre_attn_layout(q) # 形状根据模式变化
|
||||
k = pre_attn_layout(k)
|
||||
v = pre_attn_layout(v)
|
||||
|
||||
if mode == "torch":
|
||||
# 使用PyTorch原生的scaled_dot_product_attention
|
||||
if attn_mask is not None and attn_mask.dtype != torch.bool:
|
||||
attn_mask = attn_mask.to(q.dtype)
|
||||
x = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal
|
||||
)
|
||||
elif mode == "flash":
|
||||
assert flash_attn_func is not None, "flash_attn_func未定义"
|
||||
assert attn_mask is None, "不支持的注意力掩码"
|
||||
x: torch.Tensor = flash_attn_func(
|
||||
q, k, v, dropout_p=drop_rate, causal=causal, softmax_scale=None
|
||||
) # type: ignore
|
||||
elif mode == "vanilla":
|
||||
# 手动实现注意力机制
|
||||
scale_factor = 1 / math.sqrt(q.size(-1)) # 缩放因子 1/sqrt(d_k)
|
||||
|
||||
b, a, s, _ = q.shape # 获取形状参数
|
||||
s1 = k.size(2) # 键值序列长度
|
||||
|
||||
# 初始化注意力偏置
|
||||
attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
|
||||
|
||||
# 处理因果掩码
|
||||
if causal:
|
||||
assert attn_mask is None, "因果掩码和注意力掩码不能同时使用"
|
||||
# 生成下三角因果掩码
|
||||
temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(
|
||||
diagonal=0
|
||||
)
|
||||
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
||||
attn_bias = attn_bias.to(q.dtype)
|
||||
|
||||
# 处理自定义注意力掩码
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
|
||||
else:
|
||||
attn_bias += attn_mask # 允许类似ALiBi的位置偏置
|
||||
|
||||
# 计算注意力矩阵
|
||||
attn = (q @ k.transpose(-2, -1)) * scale_factor # [B,A,S,S1]
|
||||
attn += attn_bias
|
||||
|
||||
# softmax和dropout
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = torch.dropout(attn, p=drop_rate, train=True)
|
||||
|
||||
# 计算输出
|
||||
x = attn @ v # [B,A,S,D]
|
||||
else:
|
||||
raise NotImplementedError(f"不支持的注意力模式: {mode}")
|
||||
|
||||
# 应用后处理变换
|
||||
x = post_attn_layout(x) # 恢复原始维度顺序
|
||||
|
||||
# 合并注意力头维度
|
||||
b, s, a, d = x.shape
|
||||
out = x.reshape(b, s, -1) # [B,S,A*D]
|
||||
return out
|
||||
@@ -0,0 +1,326 @@
|
||||
# Modified from Flux
|
||||
#
|
||||
# Copyright 2024 Black Forest Labs
|
||||
|
||||
# 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.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
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,
|
||||
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,
|
||||
):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(
|
||||
resolution=resolution,
|
||||
in_channels=in_channels,
|
||||
ch=ch,
|
||||
ch_mult=ch_mult,
|
||||
num_res_blocks=num_res_blocks,
|
||||
z_channels=z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=resolution,
|
||||
in_channels=in_channels,
|
||||
ch=ch,
|
||||
out_ch=out_ch,
|
||||
ch_mult=ch_mult,
|
||||
num_res_blocks=num_res_blocks,
|
||||
z_channels=z_channels,
|
||||
)
|
||||
self.reg = DiagonalGaussian()
|
||||
|
||||
self.scale_factor = scale_factor
|
||||
self.shift_factor = 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,217 @@
|
||||
import torch
|
||||
from qwen_vl_utils import process_vision_info
|
||||
from transformers import (
|
||||
AutoProcessor,
|
||||
Qwen2VLForConditionalGeneration,
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
)
|
||||
from torchvision.transforms import ToPILImage
|
||||
|
||||
to_pil = ToPILImage()
|
||||
|
||||
Qwen25VL_7b_PREFIX = '''Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:
|
||||
- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.
|
||||
- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n
|
||||
Here are examples of how to transform or refine prompts:
|
||||
- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.
|
||||
- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n
|
||||
Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:
|
||||
User Prompt:'''
|
||||
|
||||
|
||||
def split_string(s):
|
||||
# 将中文引号替换为英文引号
|
||||
s = s.replace("“", '"').replace("”", '"') # use english quotes
|
||||
result = []
|
||||
# 标记是否在引号内
|
||||
in_quotes = False
|
||||
temp = ""
|
||||
|
||||
# 遍历字符串中的每个字符及其索引
|
||||
for idx, char in enumerate(s):
|
||||
# 如果字符是引号且索引大于 155
|
||||
if char == '"' and idx > 155:
|
||||
# 将引号添加到临时字符串
|
||||
temp += char
|
||||
# 如果不在引号内
|
||||
if not in_quotes:
|
||||
# 将临时字符串添加到结果列表
|
||||
result.append(temp)
|
||||
# 清空临时字符串
|
||||
temp = ""
|
||||
|
||||
# 切换引号状态
|
||||
in_quotes = not in_quotes
|
||||
continue
|
||||
# 如果在引号内
|
||||
if in_quotes:
|
||||
# 如果字符是空格
|
||||
if char.isspace():
|
||||
pass # have space token
|
||||
|
||||
# 将字符用中文引号包裹后添加到结果列表
|
||||
result.append("“" + char + "”")
|
||||
else:
|
||||
# 将字符添加到临时字符串
|
||||
temp += char
|
||||
|
||||
# 如果临时字符串不为空
|
||||
if temp:
|
||||
# 将临时字符串添加到结果列表
|
||||
result.append(temp)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Qwen25VL_7b_Embedder(torch.nn.Module):
|
||||
def __init__(self, model_path, max_length=640, dtype=torch.bfloat16, device="cuda"):
|
||||
super(Qwen25VL_7b_Embedder, self).__init__()
|
||||
self.max_length = max_length
|
||||
self.dtype = dtype
|
||||
self.device = device
|
||||
|
||||
print(f"Qwen25VL model run on current device: {torch.cuda.current_device()}")
|
||||
self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=dtype,
|
||||
attn_implementation="flash_attention_2"
|
||||
).to(torch.cuda.current_device())
|
||||
|
||||
self.model.requires_grad_(False)
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
model_path, min_pixels=256 * 28 * 28, max_pixels=324 * 28 * 28, use_fast=True
|
||||
)
|
||||
|
||||
self.prefix = Qwen25VL_7b_PREFIX
|
||||
|
||||
def forward(self, caption, ref_images):
|
||||
text_list = caption
|
||||
embs = torch.zeros(
|
||||
len(text_list),
|
||||
self.max_length,
|
||||
self.model.config.hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
hidden_states = torch.zeros(
|
||||
len(text_list),
|
||||
self.max_length,
|
||||
self.model.config.hidden_size,
|
||||
dtype=torch.bfloat16,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
masks = torch.zeros(
|
||||
len(text_list),
|
||||
self.max_length,
|
||||
dtype=torch.long,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
input_ids_list = []
|
||||
attention_mask_list = []
|
||||
emb_list = []
|
||||
|
||||
def split_string(s):
|
||||
s = s.replace("“", '"').replace("”", '"').replace("'", '''"''') # use english quotes
|
||||
result = []
|
||||
in_quotes = False
|
||||
temp = ""
|
||||
|
||||
for idx,char in enumerate(s):
|
||||
if char == '"' and idx>155:
|
||||
temp += char
|
||||
if not in_quotes:
|
||||
result.append(temp)
|
||||
temp = ""
|
||||
|
||||
in_quotes = not in_quotes
|
||||
continue
|
||||
if in_quotes:
|
||||
if char.isspace():
|
||||
pass # have space token
|
||||
|
||||
result.append("“" + char + "”")
|
||||
else:
|
||||
temp += char
|
||||
|
||||
if temp:
|
||||
result.append(temp)
|
||||
|
||||
return result
|
||||
|
||||
for idx, (txt, imgs) in enumerate(zip(text_list, ref_images)):
|
||||
|
||||
messages = [{"role": "user", "content": []}]
|
||||
|
||||
messages[0]["content"].append({"type": "text", "text": f"{self.prefix}"})
|
||||
|
||||
messages[0]["content"].append({"type": "image", "image": to_pil(imgs)})
|
||||
|
||||
# 再添加 text
|
||||
messages[0]["content"].append({"type": "text", "text": f"{txt}"})
|
||||
|
||||
# Preparation for inference
|
||||
text = self.processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True, add_vision_id=True
|
||||
)
|
||||
|
||||
image_inputs, video_inputs = process_vision_info(messages)
|
||||
|
||||
inputs = self.processor(
|
||||
text=[text],
|
||||
images=image_inputs,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
old_inputs_ids = inputs.input_ids
|
||||
text_split_list = split_string(text)
|
||||
|
||||
token_list = []
|
||||
for text_each in text_split_list:
|
||||
txt_inputs = self.processor(
|
||||
text=text_each,
|
||||
images=None,
|
||||
videos=None,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
token_each = txt_inputs.input_ids
|
||||
if token_each[0][0] == 2073 and token_each[0][-1] == 854:
|
||||
token_each = token_each[:, 1:-1]
|
||||
token_list.append(token_each)
|
||||
else:
|
||||
token_list.append(token_each)
|
||||
|
||||
new_txt_ids = torch.cat(token_list, dim=1).to("cuda")
|
||||
|
||||
new_txt_ids = new_txt_ids.to(old_inputs_ids.device)
|
||||
|
||||
idx1 = (old_inputs_ids == 151653).nonzero(as_tuple=True)[1][0]
|
||||
idx2 = (new_txt_ids == 151653).nonzero(as_tuple=True)[1][0]
|
||||
inputs.input_ids = (
|
||||
torch.cat([old_inputs_ids[0, :idx1], new_txt_ids[0, idx2:]], dim=0)
|
||||
.unsqueeze(0)
|
||||
.to("cuda")
|
||||
)
|
||||
inputs.attention_mask = (inputs.input_ids > 0).long().to("cuda")
|
||||
outputs = self.model(
|
||||
input_ids=inputs.input_ids,
|
||||
attention_mask=inputs.attention_mask,
|
||||
pixel_values=inputs.pixel_values.to("cuda"),
|
||||
image_grid_thw=inputs.image_grid_thw.to("cuda"),
|
||||
output_hidden_states=True,
|
||||
)
|
||||
|
||||
emb = outputs["hidden_states"][-1]
|
||||
|
||||
embs[idx, : min(self.max_length, emb.shape[1] - 217)] = emb[0, 217:][
|
||||
: self.max_length
|
||||
]
|
||||
|
||||
masks[idx, : min(self.max_length, emb.shape[1] - 217)] = torch.ones(
|
||||
(min(self.max_length, emb.shape[1] - 217)),
|
||||
dtype=torch.long,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
|
||||
return embs, masks
|
||||
@@ -0,0 +1,486 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
from .layers import MLP, TextProjection, TimestepEmbedder, apply_gate, attention
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
elementwise_affine=True,
|
||||
eps: float = 1e-6,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
"""
|
||||
Initialize the RMSNorm normalization layer.
|
||||
|
||||
Args:
|
||||
dim (int): The dimension of the input tensor.
|
||||
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
|
||||
|
||||
Attributes:
|
||||
eps (float): A small value added to the denominator for numerical stability.
|
||||
weight (nn.Parameter): Learnable scaling parameter.
|
||||
|
||||
"""
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
|
||||
|
||||
def _norm(self, x):
|
||||
"""
|
||||
Apply the RMSNorm normalization to the input tensor.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The normalized tensor.
|
||||
|
||||
"""
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass through the RMSNorm layer.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The output tensor after applying RMSNorm.
|
||||
|
||||
"""
|
||||
output = self._norm(x.float()).type_as(x)
|
||||
if hasattr(self, "weight"):
|
||||
output = output * self.weight
|
||||
return output
|
||||
|
||||
|
||||
def get_norm_layer(norm_layer):
|
||||
"""
|
||||
Get the normalization layer.
|
||||
|
||||
Args:
|
||||
norm_layer (str): The type of normalization layer.
|
||||
|
||||
Returns:
|
||||
norm_layer (nn.Module): The normalization layer.
|
||||
"""
|
||||
if norm_layer == "layer":
|
||||
return nn.LayerNorm
|
||||
elif norm_layer == "rms":
|
||||
return RMSNorm
|
||||
else:
|
||||
raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
|
||||
|
||||
|
||||
def get_activation_layer(act_type):
|
||||
"""get activation layer
|
||||
|
||||
Args:
|
||||
act_type (str): the activation type
|
||||
|
||||
Returns:
|
||||
torch.nn.functional: the activation layer
|
||||
"""
|
||||
if act_type == "gelu":
|
||||
return lambda: nn.GELU()
|
||||
elif act_type == "gelu_tanh":
|
||||
return lambda: nn.GELU(approximate="tanh")
|
||||
elif act_type == "relu":
|
||||
return nn.ReLU
|
||||
elif act_type == "silu":
|
||||
return nn.SiLU
|
||||
else:
|
||||
raise ValueError(f"Unknown activation type: {act_type}")
|
||||
|
||||
class IndividualTokenRefinerBlock(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
mlp_width_ratio: str = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
need_CA: bool = False,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.need_CA = need_CA
|
||||
self.heads_num = heads_num
|
||||
head_dim = hidden_size // heads_num
|
||||
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
|
||||
|
||||
self.norm1 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.self_attn_qkv = nn.Linear(
|
||||
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
qk_norm_layer = get_norm_layer(qk_norm_type)
|
||||
self.self_attn_q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_proj = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
act_layer = get_activation_layer(act_type)
|
||||
self.mlp = MLP(
|
||||
in_channels=hidden_size,
|
||||
hidden_channels=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
drop=mlp_drop_rate,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
|
||||
if self.need_CA:
|
||||
self.cross_attnblock=CrossAttnBlock(hidden_size=hidden_size,
|
||||
heads_num=heads_num,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_drop_rate=mlp_drop_rate,
|
||||
act_type=act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
**factory_kwargs,)
|
||||
# Zero-initialize the modulation
|
||||
nn.init.zeros_(self.adaLN_modulation[1].weight)
|
||||
nn.init.zeros_(self.adaLN_modulation[1].bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
c: torch.Tensor, # timestep_aware_representations + context_aware_representations
|
||||
attn_mask: torch.Tensor = None,
|
||||
y: torch.Tensor = None,
|
||||
):
|
||||
gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
|
||||
norm_x = self.norm1(x)
|
||||
qkv = self.self_attn_qkv(norm_x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
|
||||
# Apply QK-Norm if needed
|
||||
q = self.self_attn_q_norm(q).to(v)
|
||||
k = self.self_attn_k_norm(k).to(v)
|
||||
|
||||
# Self-Attention
|
||||
attn = attention(q, k, v, mode="torch", attn_mask=attn_mask)
|
||||
|
||||
x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
|
||||
|
||||
if self.need_CA:
|
||||
x = self.cross_attnblock(x, c, attn_mask, y)
|
||||
|
||||
# FFN Layer
|
||||
x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
class CrossAttnBlock(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
mlp_width_ratio: str = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.heads_num = heads_num
|
||||
head_dim = hidden_size // heads_num
|
||||
|
||||
self.norm1 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.norm1_2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.self_attn_q = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
self.self_attn_kv = nn.Linear(
|
||||
hidden_size, hidden_size*2, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
qk_norm_layer = get_norm_layer(qk_norm_type)
|
||||
self.self_attn_q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_proj = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
act_layer = get_activation_layer(act_type)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
# Zero-initialize the modulation
|
||||
nn.init.zeros_(self.adaLN_modulation[1].weight)
|
||||
nn.init.zeros_(self.adaLN_modulation[1].bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
c: torch.Tensor, # timestep_aware_representations + context_aware_representations
|
||||
attn_mask: torch.Tensor = None,
|
||||
y: torch.Tensor=None,
|
||||
|
||||
):
|
||||
gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
|
||||
norm_x = self.norm1(x)
|
||||
norm_y = self.norm1_2(y)
|
||||
q = self.self_attn_q(norm_x)
|
||||
q = rearrange(q, "B L (H D) -> B L H D", H=self.heads_num)
|
||||
kv = self.self_attn_kv(norm_y)
|
||||
k, v = rearrange(kv, "B L (K H D) -> K B L H D", K=2, H=self.heads_num)
|
||||
# Apply QK-Norm if needed
|
||||
q = self.self_attn_q_norm(q).to(v)
|
||||
k = self.self_attn_k_norm(k).to(v)
|
||||
|
||||
# Self-Attention
|
||||
attn = attention(q, k, v, mode="torch", attn_mask=attn_mask)
|
||||
|
||||
x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class IndividualTokenRefiner(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
depth,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
need_CA:bool=False,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.need_CA = need_CA
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
IndividualTokenRefinerBlock(
|
||||
hidden_size=hidden_size,
|
||||
heads_num=heads_num,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_drop_rate=mlp_drop_rate,
|
||||
act_type=act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
need_CA=self.need_CA,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(depth)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
c: torch.LongTensor,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
y:torch.Tensor=None,
|
||||
):
|
||||
self_attn_mask = None
|
||||
if mask is not None:
|
||||
batch_size = mask.shape[0]
|
||||
seq_len = mask.shape[1]
|
||||
mask = mask.to(x.device)
|
||||
# batch_size x 1 x seq_len x seq_len
|
||||
self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat(
|
||||
1, 1, seq_len, 1
|
||||
)
|
||||
# batch_size x 1 x seq_len x seq_len
|
||||
self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
|
||||
# batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num
|
||||
self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
|
||||
# avoids self-attention weight being NaN for padding tokens
|
||||
self_attn_mask[:, :, :, 0] = True
|
||||
|
||||
|
||||
for block in self.blocks:
|
||||
x = block(x, c, self_attn_mask,y)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SingleTokenRefiner(torch.nn.Module):
|
||||
"""
|
||||
A single token refiner block for llm text embedding refine.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
depth,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
need_CA:bool=False,
|
||||
attn_mode: str = "torch",
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.attn_mode = attn_mode
|
||||
self.need_CA = need_CA
|
||||
assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
|
||||
|
||||
self.input_embedder = nn.Linear(
|
||||
in_channels, hidden_size, bias=True, **factory_kwargs
|
||||
)
|
||||
if self.need_CA:
|
||||
self.input_embedder_CA = nn.Linear(
|
||||
in_channels, hidden_size, bias=True, **factory_kwargs
|
||||
)
|
||||
|
||||
act_layer = get_activation_layer(act_type)
|
||||
# Build timestep embedding layer
|
||||
self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs)
|
||||
# Build context embedding layer
|
||||
self.c_embedder = TextProjection(
|
||||
in_channels, hidden_size, act_layer, **factory_kwargs
|
||||
)
|
||||
|
||||
self.individual_token_refiner = IndividualTokenRefiner(
|
||||
hidden_size=hidden_size,
|
||||
heads_num=heads_num,
|
||||
depth=depth,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_drop_rate=mlp_drop_rate,
|
||||
act_type=act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
need_CA=need_CA,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
t: torch.LongTensor,
|
||||
mask: Optional[torch.LongTensor] = None,
|
||||
y: torch.LongTensor=None,
|
||||
):
|
||||
timestep_aware_representations = self.t_embedder(t)
|
||||
|
||||
if mask is None:
|
||||
context_aware_representations = x.mean(dim=1)
|
||||
else:
|
||||
mask_float = mask.unsqueeze(-1) # [b, s1, 1]
|
||||
context_aware_representations = (x * mask_float).sum(
|
||||
dim=1
|
||||
) / mask_float.sum(dim=1)
|
||||
context_aware_representations = self.c_embedder(context_aware_representations)
|
||||
c = timestep_aware_representations + context_aware_representations
|
||||
|
||||
x = self.input_embedder(x)
|
||||
if self.need_CA:
|
||||
y = self.input_embedder_CA(y)
|
||||
x = self.individual_token_refiner(x, c, mask, y)
|
||||
else:
|
||||
x = self.individual_token_refiner(x, c, mask)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class Qwen2Connector(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
# biclip_dim=1024,
|
||||
in_channels=3584,
|
||||
hidden_size=4096,
|
||||
heads_num=32,
|
||||
depth=2,
|
||||
need_CA=False,
|
||||
device=None,
|
||||
dtype=torch.bfloat16,
|
||||
):
|
||||
super().__init__()
|
||||
factory_kwargs = {"device": device, "dtype":dtype}
|
||||
|
||||
self.S =SingleTokenRefiner(in_channels=in_channels,hidden_size=hidden_size,heads_num=heads_num,depth=depth,need_CA=need_CA,**factory_kwargs)
|
||||
self.global_proj_out=nn.Linear(in_channels,768)
|
||||
|
||||
self.scale_factor = nn.Parameter(torch.zeros(1))
|
||||
with torch.no_grad():
|
||||
self.scale_factor.data += -(1 - 0.09)
|
||||
|
||||
def forward(self, x,t,mask):
|
||||
mask_float = mask.unsqueeze(-1) # [b, s1, 1]
|
||||
x_mean = (x * mask_float).sum(
|
||||
dim=1
|
||||
) / mask_float.sum(dim=1) * (1 + self.scale_factor.to(x.dtype))
|
||||
|
||||
global_out=self.global_proj_out(x_mean)
|
||||
encoder_hidden_states = self.S(x,t,mask)
|
||||
return encoder_hidden_states,global_out
|
||||
@@ -0,0 +1,639 @@
|
||||
# Modified from Flux
|
||||
#
|
||||
# Copyright 2024 Black Forest Labs
|
||||
|
||||
# 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.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math # noqa: I001
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from liger_kernel.ops.rms_norm import LigerRMSNormFunction
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
try:
|
||||
import flash_attn
|
||||
from flash_attn.flash_attn_interface import (
|
||||
_flash_attn_forward,
|
||||
flash_attn_varlen_func,
|
||||
)
|
||||
except ImportError:
|
||||
flash_attn = None
|
||||
flash_attn_varlen_func = None
|
||||
_flash_attn_forward = None
|
||||
|
||||
|
||||
MEMORY_LAYOUT = {
|
||||
"flash": (
|
||||
lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
|
||||
lambda x: x,
|
||||
),
|
||||
"torch": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
"vanilla": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mode="flash",
|
||||
drop_rate=0,
|
||||
attn_mask=None,
|
||||
causal=False,
|
||||
cu_seqlens_q=None,
|
||||
cu_seqlens_kv=None,
|
||||
max_seqlen_q=None,
|
||||
max_seqlen_kv=None,
|
||||
batch_size=1,
|
||||
):
|
||||
"""
|
||||
Perform QKV self attention.
|
||||
|
||||
Args:
|
||||
q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads.
|
||||
k (torch.Tensor): Key tensor with shape [b, s1, a, d]
|
||||
v (torch.Tensor): Value tensor with shape [b, s1, a, d]
|
||||
mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'.
|
||||
drop_rate (float): Dropout rate in attention map. (default: 0)
|
||||
attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla).
|
||||
(default: None)
|
||||
causal (bool): Whether to use causal attention. (default: False)
|
||||
cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
|
||||
used to index into q.
|
||||
cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
|
||||
used to index into kv.
|
||||
max_seqlen_q (int): The maximum sequence length in the batch of q.
|
||||
max_seqlen_kv (int): The maximum sequence length in the batch of k and v.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor after self attention with shape [b, s, ad]
|
||||
"""
|
||||
pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
|
||||
q = pre_attn_layout(q)
|
||||
k = pre_attn_layout(k)
|
||||
v = pre_attn_layout(v)
|
||||
|
||||
if mode == "torch":
|
||||
if attn_mask is not None and attn_mask.dtype != torch.bool:
|
||||
attn_mask = attn_mask.to(q.dtype)
|
||||
x = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal
|
||||
)
|
||||
elif mode == "flash":
|
||||
assert flash_attn_varlen_func is not None
|
||||
x: torch.Tensor = flash_attn_varlen_func(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
) # type: ignore
|
||||
# x with shape [(bxs), a, d]
|
||||
x = x.view(batch_size, max_seqlen_q, x.shape[-2], x.shape[-1]) # type: ignore # reshape x to [b, s, a, d]
|
||||
elif mode == "vanilla":
|
||||
scale_factor = 1 / math.sqrt(q.size(-1))
|
||||
|
||||
b, a, s, _ = q.shape
|
||||
s1 = k.size(2)
|
||||
attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
|
||||
if causal:
|
||||
# Only applied to self attention
|
||||
assert attn_mask is None, (
|
||||
"Causal mask and attn_mask cannot be used together"
|
||||
)
|
||||
temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(
|
||||
diagonal=0
|
||||
)
|
||||
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
||||
attn_bias.to(q.dtype)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
|
||||
else:
|
||||
attn_bias += attn_mask
|
||||
|
||||
# TODO: Maybe force q and k to be float32 to avoid numerical overflow
|
||||
attn = (q @ k.transpose(-2, -1)) * scale_factor
|
||||
attn += attn_bias
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = torch.dropout(attn, p=drop_rate, train=True)
|
||||
x = attn @ v
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported attention mode: {mode}")
|
||||
|
||||
x = post_attn_layout(x)
|
||||
b, s, a, d = x.shape
|
||||
out = x.reshape(b, s, -1)
|
||||
return out
|
||||
|
||||
|
||||
def apply_gate(x, gate=None, tanh=False):
|
||||
"""AI is creating summary for apply_gate
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): input tensor.
|
||||
gate (torch.Tensor, optional): gate tensor. Defaults to None.
|
||||
tanh (bool, optional): whether to use tanh function. Defaults to False.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the output tensor after apply gate.
|
||||
"""
|
||||
if gate is None:
|
||||
return x
|
||||
if tanh:
|
||||
return x * gate.unsqueeze(1).tanh()
|
||||
else:
|
||||
return x * gate.unsqueeze(1)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_channels=None,
|
||||
out_features=None,
|
||||
act_layer=nn.GELU,
|
||||
norm_layer=None,
|
||||
bias=True,
|
||||
drop=0.0,
|
||||
use_conv=False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
super().__init__()
|
||||
out_features = out_features or in_channels
|
||||
hidden_channels = hidden_channels or in_channels
|
||||
bias = (bias, bias)
|
||||
drop_probs = (drop, drop)
|
||||
linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
|
||||
|
||||
self.fc1 = linear_layer(
|
||||
in_channels, hidden_channels, bias=bias[0], device=device, dtype=dtype
|
||||
)
|
||||
self.act = act_layer()
|
||||
self.drop1 = nn.Dropout(drop_probs[0])
|
||||
self.norm = (
|
||||
norm_layer(hidden_channels, device=device, dtype=dtype)
|
||||
if norm_layer is not None
|
||||
else nn.Identity()
|
||||
)
|
||||
self.fc2 = linear_layer(
|
||||
hidden_channels, out_features, bias=bias[1], device=device, dtype=dtype
|
||||
)
|
||||
self.drop2 = nn.Dropout(drop_probs[1])
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop1(x)
|
||||
x = self.norm(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop2(x)
|
||||
return x
|
||||
|
||||
|
||||
class TextProjection(nn.Module):
|
||||
"""
|
||||
Projects text embeddings. Also handles dropout for classifier-free guidance.
|
||||
|
||||
Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
self.linear_1 = nn.Linear(
|
||||
in_features=in_channels,
|
||||
out_features=hidden_size,
|
||||
bias=True,
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.act_1 = act_layer()
|
||||
self.linear_2 = nn.Linear(
|
||||
in_features=hidden_size,
|
||||
out_features=hidden_size,
|
||||
bias=True,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def forward(self, caption):
|
||||
hidden_states = self.linear_1(caption)
|
||||
hidden_states = self.act_1(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
act_layer,
|
||||
frequency_embedding_size=256,
|
||||
max_period=10000,
|
||||
out_size=None,
|
||||
dtype=None,
|
||||
device=None,
|
||||
):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
self.max_period = max_period
|
||||
if out_size is None:
|
||||
out_size = hidden_size
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(
|
||||
frequency_embedding_size, hidden_size, bias=True, **factory_kwargs
|
||||
),
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
nn.init.normal_(self.mlp[0].weight, std=0.02) # type: ignore
|
||||
nn.init.normal_(self.mlp[2].weight, std=0.02) # type: ignore
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
|
||||
Args:
|
||||
t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
dim (int): the dimension of the output.
|
||||
max_period (int): controls the minimum frequency of the embeddings.
|
||||
|
||||
Returns:
|
||||
embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
|
||||
|
||||
.. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
"""
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(device=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
|
||||
)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(
|
||||
t, self.frequency_embedding_size, self.max_period
|
||||
).type(self.mlp[0].weight.dtype) # type: ignore
|
||||
t_emb = self.mlp(t_freq)
|
||||
return t_emb
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)))
|
||||
|
||||
|
||||
def rope(pos, dim: int, theta: int):
|
||||
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 attention_after_rope(q, k, v, pe):
|
||||
q, k = apply_rope(q, k, pe)
|
||||
|
||||
from .attention import attention
|
||||
|
||||
x = attention(q, k, v, mode="flash")
|
||||
return x
|
||||
|
||||
|
||||
@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True)
|
||||
def apply_rope(xq, xk, freqs_cis):
|
||||
# 将 num_heads 和 seq_len 的维度交换回原函数的处理顺序
|
||||
xq = xq.transpose(1, 2) # [batch, num_heads, seq_len, head_dim]
|
||||
xk = xk.transpose(1, 2)
|
||||
|
||||
# 将 head_dim 拆分为复数部分（实部和虚部）
|
||||
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]
|
||||
|
||||
# 恢复张量形状并转置回目标维度顺序
|
||||
xq_out = xq_out.reshape(*xq.shape).type_as(xq).transpose(1, 2)
|
||||
xk_out = xk_out.reshape(*xk.shape).type_as(xk).transpose(1, 2)
|
||||
|
||||
return xq_out, xk_out
|
||||
|
||||
|
||||
@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True)
|
||||
def scale_add_residual(
|
||||
x: torch.Tensor, scale: torch.Tensor, residual: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
return x * scale + residual
|
||||
|
||||
|
||||
@torch.compile(mode="max-autotune-no-cudagraphs", dynamic=True)
|
||||
def layernorm_and_scale_shift(
|
||||
x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
return torch.nn.functional.layer_norm(x, (x.size(-1),)) * (scale + 1) + shift
|
||||
|
||||
|
||||
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(self, x: Tensor, pe: Tensor) -> Tensor:
|
||||
qkv = self.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
x = attention_after_rope(q, k, v, pe=pe)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModulationOut:
|
||||
shift: Tensor
|
||||
scale: Tensor
|
||||
gate: Tensor
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
@staticmethod
|
||||
def rms_norm_fast(x, weight, eps):
|
||||
return LigerRMSNormFunction.apply(
|
||||
x,
|
||||
weight,
|
||||
eps,
|
||||
0.0,
|
||||
"gemma",
|
||||
True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def rms_norm(x, weight, eps):
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)
|
||||
return (x * rrms).to(dtype=x_dtype) * weight
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
return self.rms_norm_fast(x, self.scale.to(x.dtype), 1e-6)
|
||||
|
||||
|
||||
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 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 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.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),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
img_mod1, img_mod2 = self.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = self.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = self.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = self.img_attn.qkv(img_modulated)
|
||||
img_q, img_k, img_v = rearrange(
|
||||
img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads
|
||||
)
|
||||
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
|
||||
|
||||
# prepare txt for attention
|
||||
txt_modulated = self.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = self.txt_attn.qkv(txt_modulated)
|
||||
txt_q, txt_k, txt_v = rearrange(
|
||||
txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads
|
||||
)
|
||||
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=1)
|
||||
k = torch.cat((txt_k, img_k), dim=1)
|
||||
v = torch.cat((txt_v, img_v), dim=1)
|
||||
|
||||
attn = attention_after_rope(q, k, v, pe=pe)
|
||||
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
|
||||
img_mlp = self.img_mlp(
|
||||
(1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift
|
||||
)
|
||||
img = scale_add_residual(img_mlp, img_mod2.gate, img)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
|
||||
txt_mlp = self.txt_mlp(
|
||||
(1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift
|
||||
)
|
||||
txt = scale_add_residual(txt_mlp, txt_mod2.gate, txt)
|
||||
return img, txt
|
||||
|
||||
|
||||
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
|
||||
head_dim = hidden_size // num_heads
|
||||
self.scale = qk_scale or head_dim**-0.5
|
||||
|
||||
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
# qkv and mlp_in
|
||||
self.linear1 = 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(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)
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
|
||||
mod, _ = self.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(
|
||||
self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
|
||||
)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn = attention_after_rope(q, k, v, pe=pe)
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
|
||||
return scale_add_residual(output, mod.gate, x)
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,143 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from .connector_edit import Qwen2Connector
|
||||
from .layers import DoubleStreamBlock, EmbedND, LastLayer, MLPEmbedder, SingleStreamBlock
|
||||
|
||||
|
||||
@dataclass
|
||||
class Step1XParams:
|
||||
in_channels: int
|
||||
out_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
|
||||
|
||||
|
||||
class Step1XEdit(nn.Module):
|
||||
"""
|
||||
Transformer model for flow matching on sequences.
|
||||
"""
|
||||
|
||||
def __init__(self, params: Step1XParams):
|
||||
super().__init__()
|
||||
|
||||
self.params = params
|
||||
self.in_channels = params.in_channels
|
||||
self.out_channels = params.out_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.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.connector = Qwen2Connector()
|
||||
|
||||
@staticmethod
|
||||
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
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
|
||||
img = self.img_in(img)
|
||||
vec = self.time_in(self.timestep_embedding(timesteps, 256))
|
||||
|
||||
vec = vec + self.vector_in(y)
|
||||
txt = self.txt_in(txt)
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=1)
|
||||
pe = self.pe_embedder(ids)
|
||||
|
||||
for block in self.double_blocks:
|
||||
img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
|
||||
|
||||
img = torch.cat((txt, img), 1)
|
||||
for block in self.single_blocks:
|
||||
img = block(img, vec=vec, pe=pe)
|
||||
img = img[:, txt.shape[1] :, ...]
|
||||
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
return img
|
||||
@@ -0,0 +1 @@
|
||||
# Step1X-Edit nodes package
|
||||
@@ -0,0 +1,338 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
import itertools
|
||||
import math
|
||||
from einops import rearrange, repeat
|
||||
from torchvision.transforms import functional as F
|
||||
|
||||
from ..utils.model_loader import Step1XEditModelBundle
|
||||
from ..utils.sampling import get_schedule
|
||||
import folder_paths
|
||||
|
||||
|
||||
class Step1XEditModelLoader:
|
||||
"""Node for loading the Step1X-Edit model."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"diffusion_model": (folder_paths.get_filename_list("diffusion_models"), {"default": "step1x-edit-i1258-FP8.safetensors"}),
|
||||
"vae": (folder_paths.get_filename_list("vae"), {"default": "vae.safetensors"}),
|
||||
"text_encoder": (os.listdir(folder_paths.get_folder_paths("text_encoders")[0]), {"default": "Qwen2.5-VL-7B-Instruct"}),
|
||||
"dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
|
||||
"quantized": ("BOOLEAN", {"default": True}),
|
||||
"offload": ("BOOLEAN", {"default": False})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STEP1X_MODEL_BUNDLE",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "Step1X-Edit"
|
||||
|
||||
def load_model(self, diffusion_model, vae, text_encoder, dtype, quantized, offload):
|
||||
"""Load the Step1X-Edit model components."""
|
||||
dit_path = folder_paths.get_full_path("diffusion_models", diffusion_model)
|
||||
ae_path = folder_paths.get_full_path("vae", vae)
|
||||
qwen2vl_model_path = os.path.join(folder_paths.get_folder_paths("text_encoders")[0], text_encoder)
|
||||
|
||||
# Convert dtype string to torch dtype
|
||||
if dtype == "bfloat16":
|
||||
torch_dtype = torch.bfloat16
|
||||
elif dtype == "float16":
|
||||
torch_dtype = torch.float16
|
||||
else:
|
||||
torch_dtype = torch.float32
|
||||
|
||||
# Determine device
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# Create model bundle
|
||||
model_bundle = Step1XEditModelBundle(device=device, dtype=torch_dtype)
|
||||
model_bundle.set_quantized(quantized)
|
||||
model_bundle.set_offload(offload)
|
||||
|
||||
# Load models
|
||||
success = model_bundle.load_models(
|
||||
dit_path=dit_path,
|
||||
ae_path=ae_path,
|
||||
qwen2vl_model_path=qwen2vl_model_path,
|
||||
max_length=640
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise RuntimeError("Failed to load Step1X-Edit models. Please check paths and ensure all dependencies are installed.")
|
||||
|
||||
return (model_bundle,)
|
||||
|
||||
|
||||
class Step1XEditGenerateNode:
|
||||
"""Node for generating images using Step1X-Edit."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model_bundle": ("STEP1X_MODEL_BUNDLE",),
|
||||
"input_image": ("IMAGE",),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"num_steps": ("INT", {"default": 28, "min": 1, "max": 100}),
|
||||
"cfg_guidance": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
||||
"seed": ("INT", {"default": 42}),
|
||||
"size_level": ("INT", {"default": 512, "min": 256, "max": 1024, "step": 64})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "Step1X-Edit"
|
||||
|
||||
def process_diff_norm(self, diff_norm, k):
|
||||
pow_result = torch.pow(diff_norm, k)
|
||||
result = torch.where(
|
||||
diff_norm > 1.0,
|
||||
pow_result,
|
||||
torch.where(diff_norm < 1.0, torch.ones_like(diff_norm), diff_norm),
|
||||
)
|
||||
return result
|
||||
|
||||
def unpack(self, x, height, width):
|
||||
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,
|
||||
)
|
||||
|
||||
def prepare(self, prompt, img, ref_image, ref_image_raw, model_bundle, device):
|
||||
bs, _, h, w = img.shape
|
||||
bs, _, ref_h, ref_w = ref_image.shape
|
||||
|
||||
assert h == ref_h and w == ref_w
|
||||
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
elif bs >= 1 and isinstance(prompt, str):
|
||||
prompt = [prompt] * bs
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
ref_img = rearrange(ref_image, "b c (ref_h ph) (ref_w pw) -> b (ref_h ref_w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
ref_img = repeat(ref_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 = torch.zeros(ref_h // 2, ref_w // 2, 3)
|
||||
ref_img_ids[..., 1] = ref_img_ids[..., 1] + torch.arange(ref_h // 2)[:, None]
|
||||
ref_img_ids[..., 2] = ref_img_ids[..., 2] + torch.arange(ref_w // 2)[None, :]
|
||||
ref_img_ids = repeat(ref_img_ids, "ref_h ref_w c -> b (ref_h ref_w) c", b=bs)
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
|
||||
llm_encoder = model_bundle.llm_encoder
|
||||
if model_bundle.offload:
|
||||
llm_encoder = llm_encoder.to(device)
|
||||
txt, mask = llm_encoder(prompt, ref_image_raw)
|
||||
if model_bundle.offload:
|
||||
llm_encoder = llm_encoder.cpu()
|
||||
self.cudagc()
|
||||
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
img = torch.cat([img, ref_img.to(device=img.device, dtype=img.dtype)], dim=-2)
|
||||
img_ids = torch.cat([img_ids, ref_img_ids], dim=-2)
|
||||
|
||||
return {
|
||||
"img": img,
|
||||
"mask": mask,
|
||||
"img_ids": img_ids.to(img.device),
|
||||
"llm_embedding": txt.to(img.device),
|
||||
"txt_ids": txt_ids.to(img.device),
|
||||
}
|
||||
|
||||
def cudagc(self):
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
def denoise(self, model_bundle, img, img_ids, llm_embedding, txt_ids, timesteps, cfg_guidance=4.5, mask=None, timesteps_truncate=1.0):
|
||||
device = img.device
|
||||
dit = model_bundle.dit
|
||||
if model_bundle.offload:
|
||||
dit = dit.to(device)
|
||||
|
||||
for t_curr, t_prev in itertools.pairwise(timesteps):
|
||||
if img.shape[0] == 1 and cfg_guidance != -1:
|
||||
img = torch.cat([img, img], dim=0)
|
||||
|
||||
t_vec = torch.full(
|
||||
(img.shape[0],), t_curr, dtype=img.dtype, device=img.device
|
||||
)
|
||||
txt, vec = dit.connector(llm_embedding, t_vec, mask)
|
||||
|
||||
pred = dit(
|
||||
img=img,
|
||||
img_ids=img_ids,
|
||||
txt=txt,
|
||||
txt_ids=txt_ids,
|
||||
y=vec,
|
||||
timesteps=t_vec,
|
||||
)
|
||||
|
||||
if cfg_guidance != -1:
|
||||
cond, uncond = (
|
||||
pred[0 : pred.shape[0] // 2, :],
|
||||
pred[pred.shape[0] // 2 :, :],
|
||||
)
|
||||
if t_curr > timesteps_truncate:
|
||||
diff = cond - uncond
|
||||
diff_norm = torch.norm(diff, dim=(2), keepdim=True)
|
||||
pred = uncond + cfg_guidance * (
|
||||
cond - uncond
|
||||
) / self.process_diff_norm(diff_norm, k=0.4)
|
||||
else:
|
||||
pred = uncond + cfg_guidance * (cond - uncond)
|
||||
|
||||
tem_img = img[0 : img.shape[0] // 2, :] + (t_prev - t_curr) * pred
|
||||
img_input_length = img.shape[1] // 2
|
||||
img = torch.cat(
|
||||
[
|
||||
tem_img[:, :img_input_length],
|
||||
img[ : img.shape[0] // 2, img_input_length:],
|
||||
], dim=1
|
||||
)
|
||||
|
||||
if model_bundle.offload:
|
||||
dit = dit.cpu()
|
||||
self.cudagc()
|
||||
|
||||
return img[:, :img.shape[1] // 2]
|
||||
|
||||
def input_process_image(self, img_tensor, size_level=512):
|
||||
# Convert tensor to PIL for resizing
|
||||
img = F.to_pil_image(img_tensor.squeeze(0))
|
||||
|
||||
# Calculate dimensions
|
||||
w, h = img.size
|
||||
r = w / h
|
||||
|
||||
if w > h:
|
||||
w_new = math.ceil(math.sqrt(size_level * size_level * r))
|
||||
h_new = math.ceil(w_new / r)
|
||||
else:
|
||||
h_new = math.ceil(math.sqrt(size_level * size_level / r))
|
||||
w_new = math.ceil(h_new * r)
|
||||
|
||||
h_new = math.ceil(h_new) // 16 * 16
|
||||
w_new = math.ceil(w_new) // 16 * 16
|
||||
|
||||
img_resized = img.resize((w_new, h_new))
|
||||
return img_resized, img.size
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, model_bundle, input_image, prompt, negative_prompt, num_steps, cfg_guidance, seed, size_level):
|
||||
# ComfyUI passes images in BHWC format [batch, height, width, channels] with values in range [0, 1]
|
||||
# Convert to BCHW format for processing
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# Convert from ComfyUI format [B, H, W, C] to [B, C, H, W]
|
||||
input_tensor = input_image.permute(0, 3, 1, 2)
|
||||
|
||||
# Process the input image
|
||||
batch_size = input_tensor.shape[0]
|
||||
processed_images = []
|
||||
|
||||
for i in range(batch_size):
|
||||
single_img = input_tensor[i:i+1]
|
||||
|
||||
# Process image to get proper dimensions
|
||||
ref_img_pil, original_size = self.input_process_image(single_img, size_level)
|
||||
width, height = ref_img_pil.width, ref_img_pil.height
|
||||
|
||||
# Convert PIL back to tensor in proper format
|
||||
ref_img_tensor = F.to_tensor(ref_img_pil).unsqueeze(0).to(device) * 2 - 1
|
||||
|
||||
# Get autoencoder from model bundle
|
||||
ae = model_bundle.ae
|
||||
if model_bundle.offload:
|
||||
ae = ae.to(device)
|
||||
|
||||
# Encode the image
|
||||
ref_images = ae.encode(ref_img_tensor)
|
||||
|
||||
if model_bundle.offload:
|
||||
ae = ae.cpu()
|
||||
self.cudagc()
|
||||
|
||||
# Setup seed
|
||||
seed_value = int(seed)
|
||||
if seed_value < 0:
|
||||
seed_value = torch.Generator(device="cpu").seed()
|
||||
|
||||
# Create latent noise
|
||||
x = torch.randn(
|
||||
1, # num_samples
|
||||
16,
|
||||
height // 8,
|
||||
width // 8,
|
||||
device=device,
|
||||
dtype=torch.bfloat16 if model_bundle.dtype == torch.bfloat16 else torch.float32,
|
||||
generator=torch.Generator(device=device).manual_seed(seed_value),
|
||||
)
|
||||
|
||||
# Get timesteps
|
||||
timesteps = get_schedule(
|
||||
num_steps, x.shape[-1] * x.shape[-2] // 4, shift=True
|
||||
)
|
||||
|
||||
# Cat for classifier-free guidance
|
||||
x = torch.cat([x, x], dim=0)
|
||||
ref_images = torch.cat([ref_images, ref_images], dim=0)
|
||||
ref_img_tensor = torch.cat([ref_img_tensor, ref_img_tensor], dim=0)
|
||||
|
||||
# Prepare inputs
|
||||
inputs = self.prepare([prompt, negative_prompt], x, ref_image=ref_images, ref_image_raw=ref_img_tensor, model_bundle=model_bundle, device=device)
|
||||
|
||||
# Denoise
|
||||
with torch.autocast(device_type=device, dtype=torch.bfloat16):
|
||||
x = self.denoise(
|
||||
model_bundle=model_bundle,
|
||||
**inputs,
|
||||
cfg_guidance=cfg_guidance,
|
||||
timesteps=timesteps,
|
||||
timesteps_truncate=1.0,
|
||||
)
|
||||
|
||||
# Unpack and decode
|
||||
x = self.unpack(x.float(), height, width)
|
||||
if model_bundle.offload:
|
||||
ae = model_bundle.ae.to(device)
|
||||
x = ae.decode(x)
|
||||
if model_bundle.offload:
|
||||
ae = ae.cpu()
|
||||
self.cudagc()
|
||||
|
||||
x = x.clamp(-1, 1)
|
||||
x = x.mul(0.5).add(0.5)
|
||||
|
||||
# Convert back to PIL for resizing to original dimensions
|
||||
for img in x.float():
|
||||
pil_img = F.to_pil_image(img)
|
||||
pil_img = pil_img.resize(original_size)
|
||||
|
||||
# Convert back to tensor in ComfyUI format [H, W, C]
|
||||
img_tensor = torch.tensor(np.array(pil_img)).float() / 255.0
|
||||
processed_images.append(img_tensor)
|
||||
|
||||
# Stack all processed images
|
||||
output_tensor = torch.stack(processed_images)
|
||||
|
||||
return (output_tensor,)
|
||||
@@ -0,0 +1,7 @@
|
||||
torch>=2.3.1
|
||||
liger_kernel>=0.5.4
|
||||
einops>=0.8.1
|
||||
transformers>=4.49.0
|
||||
qwen_vl_utils>=0.0.10
|
||||
safetensors>=0.4.5
|
||||
pillow>=11.1.0
|
||||
@@ -0,0 +1 @@
|
||||
# Step1X-Edit utilities package
|
||||
@@ -0,0 +1,59 @@
|
||||
import platform
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def get_cuda_version():
|
||||
if torch.cuda.is_available():
|
||||
cuda_version = torch.version.cuda
|
||||
return f"cu{cuda_version.replace('.', '')[:2]}" # 例如：cu121
|
||||
return "cpu"
|
||||
|
||||
|
||||
def get_torch_version():
|
||||
return f"torch{torch.__version__.split('+')[0]}"[:-2] # 例如：torch2.2
|
||||
|
||||
|
||||
def get_python_version():
|
||||
version = sys.version_info
|
||||
return f"cp{version.major}{version.minor}" # 例如：cp310
|
||||
|
||||
|
||||
def get_abi_flag():
|
||||
return "abiTRUE" if torch._C._GLIBCXX_USE_CXX11_ABI else "abiFALSE"
|
||||
|
||||
|
||||
def get_platform():
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
if system == "linux" and machine == "x86_64":
|
||||
return "linux_x86_64"
|
||||
elif system == "windows" and machine == "amd64":
|
||||
return "win_amd64"
|
||||
elif system == "darwin" and machine == "x86_64":
|
||||
return "macosx_x86_64"
|
||||
else:
|
||||
raise ValueError(f"Unsupported platform: {system}_{machine}")
|
||||
|
||||
|
||||
def generate_flash_attn_filename(flash_attn_version="2.7.2.post1"):
|
||||
cuda_version = get_cuda_version()
|
||||
torch_version = get_torch_version()
|
||||
python_version = get_python_version()
|
||||
abi_flag = get_abi_flag()
|
||||
platform_tag = get_platform()
|
||||
|
||||
filename = (
|
||||
f"flash_attn-{flash_attn_version}+{cuda_version}{torch_version}cxx11{abi_flag}-"
|
||||
f"{python_version}-{python_version}-{platform_tag}.whl"
|
||||
)
|
||||
return filename
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
filename = generate_flash_attn_filename()
|
||||
print("Generated filename:", filename)
|
||||
except Exception as e:
|
||||
print("Error generating filename:", e)
|
||||
@@ -0,0 +1,120 @@
|
||||
import torch
|
||||
from .step1x_edit_utils import load_state_dict
|
||||
|
||||
|
||||
# Try to import Step1X-Edit modules
|
||||
try:
|
||||
from ..modules.autoencoder import AutoEncoder
|
||||
from ..modules.conditioner import Qwen25VL_7b_Embedder as Qwen2VLEmbedder
|
||||
from ..modules.model_edit import Step1XParams, Step1XEdit
|
||||
|
||||
except ImportError:
|
||||
print("Warning: Could not import Step1X-Edit modules. Will attempt to load them dynamically.")
|
||||
|
||||
# Define placeholders for the modules
|
||||
AutoEncoder = None
|
||||
Qwen2VLEmbedder = None
|
||||
Step1XParams = None
|
||||
Step1XEdit = None
|
||||
|
||||
|
||||
class Step1XEditModelBundle:
|
||||
"""Bundle containing all components of the Step1X-Edit model."""
|
||||
|
||||
def __init__(self, ae=None, dit=None, llm_encoder=None, device="cuda", dtype=torch.bfloat16):
|
||||
self.ae = ae
|
||||
self.dit = dit
|
||||
self.llm_encoder = llm_encoder
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
self.quantized = False
|
||||
self.offload = False
|
||||
|
||||
def load_models(self, dit_path, ae_path, qwen2vl_model_path, max_length=256):
|
||||
print(f"Loading models from dit_path: {dit_path}, \n ae_path: {ae_path}, \n qwen2vl_model_path: {qwen2vl_model_path}")
|
||||
"""Load all model components."""
|
||||
if Step1XParams is None or AutoEncoder is None or Qwen2VLEmbedder is None:
|
||||
print("Error: Step1X-Edit modules not found. Please make sure they are properly installed.")
|
||||
return False
|
||||
|
||||
# Initialize LLM encoder
|
||||
self.llm_encoder = Qwen2VLEmbedder(
|
||||
qwen2vl_model_path,
|
||||
device=self.device,
|
||||
max_length=max_length,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
# Initialize autoencoder
|
||||
with torch.device("meta"):
|
||||
self.ae = AutoEncoder(
|
||||
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,
|
||||
)
|
||||
|
||||
# Initialize DIT model
|
||||
step1x_params = Step1XParams(
|
||||
in_channels=64,
|
||||
out_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,
|
||||
)
|
||||
self.dit = Step1XEdit(step1x_params)
|
||||
|
||||
# Load weights
|
||||
self.ae = load_state_dict(self.ae, ae_path, 'cpu')
|
||||
self.dit = load_state_dict(self.dit, dit_path, 'cpu')
|
||||
|
||||
# Set model precision
|
||||
self.ae = self.ae.to(dtype=torch.float32)
|
||||
|
||||
if not self.quantized:
|
||||
self.dit = self.dit.to(dtype=self.dtype)
|
||||
|
||||
if not self.offload:
|
||||
self.dit = self.dit.to(device=self.device)
|
||||
self.ae = self.ae.to(device=self.device)
|
||||
|
||||
return True
|
||||
|
||||
def set_quantized(self, quantized):
|
||||
"""Set whether to use quantized weights."""
|
||||
self.quantized = quantized
|
||||
|
||||
def set_offload(self, offload):
|
||||
"""Set whether to offload models to CPU."""
|
||||
self.offload = offload
|
||||
|
||||
def to_device(self, device):
|
||||
"""Move models to specified device."""
|
||||
self.device = device
|
||||
if not self.offload:
|
||||
self.dit = self.dit.to(device=self.device)
|
||||
self.ae = self.ae.to(device=self.device)
|
||||
|
||||
def get_models(self):
|
||||
"""Return all model components."""
|
||||
return {
|
||||
"ae": self.ae,
|
||||
"dit": self.dit,
|
||||
"llm_encoder": self.llm_encoder,
|
||||
"device": self.device,
|
||||
"dtype": self.dtype,
|
||||
"quantized": self.quantized,
|
||||
"offload": self.offload
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
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 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) -> Callable[[float], float]:
|
||||
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:
|
||||
# estimate 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()
|
||||
@@ -0,0 +1,24 @@
|
||||
import os
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
|
||||
|
||||
def load_state_dict(model, ckpt_path, device="cuda", strict=False, assign=True):
|
||||
"""Load state dict from checkpoint file."""
|
||||
if os.path.splitext(ckpt_path)[1] == ".safetensors":
|
||||
state_dict = load_file(ckpt_path, device)
|
||||
else:
|
||||
state_dict = torch.load(ckpt_path, map_location="cpu")
|
||||
|
||||
missing, unexpected = model.load_state_dict(
|
||||
state_dict, strict=strict, assign=assign
|
||||
)
|
||||
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))
|
||||
return model
|
||||
Reference in New Issue
Block a user