Compare commits
9 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b71269675e | |||
| 36059182f1 | |||
| 9169e81609 | |||
| 8160289373 | |||
| 08782bf3bf | |||
| 0f252be0ed | |||
| e3d4a6b070 | |||
| ad00c565b7 | |||
| f27949dad9 |
@@ -340,8 +340,6 @@
|
||||
title: AllegroTransformer3DModel
|
||||
- local: api/models/aura_flow_transformer2d
|
||||
title: AuraFlowTransformer2DModel
|
||||
- local: api/models/bria_transformer
|
||||
title: BriaTransformer2DModel
|
||||
- local: api/models/chroma_transformer
|
||||
title: ChromaTransformer2DModel
|
||||
- local: api/models/cogvideox_transformer3d
|
||||
@@ -470,8 +468,6 @@
|
||||
title: AutoPipeline
|
||||
- local: api/pipelines/blip_diffusion
|
||||
title: BLIP-Diffusion
|
||||
- local: api/pipelines/bria_3_2
|
||||
title: Bria 3.2
|
||||
- local: api/pipelines/chroma
|
||||
title: Chroma
|
||||
- local: api/pipelines/cogvideox
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
-->
|
||||
|
||||
# BriaTransformer2DModel
|
||||
|
||||
A modified flux Transformer model from [Bria](https://huggingface.co/briaai/BRIA-3.2)
|
||||
|
||||
## BriaTransformer2DModel
|
||||
|
||||
[[autodoc]] BriaTransformer2DModel
|
||||
@@ -1,44 +0,0 @@
|
||||
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
-->
|
||||
|
||||
# Bria 3.2
|
||||
|
||||
Bria 3.2 is the next-generation commercial-ready text-to-image model. With just 4 billion parameters, it provides exceptional aesthetics and text rendering, evaluated to provide on par results to leading open-source models, and outperforming other licensed models.
|
||||
In addition to being built entirely on licensed data, 3.2 provides several advantages for enterprise and commercial use:
|
||||
|
||||
- Efficient Compute - the model is X3 smaller than the equivalent models in the market (4B parameters vs 12B parameters other open source models)
|
||||
- Architecture Consistency: Same architecture as 3.1—ideal for users looking to upgrade without disruption.
|
||||
- Fine-tuning Speedup: 2x faster fine-tuning on L40S and A100.
|
||||
|
||||
Original model checkpoints for Bria 3.2 can be found [here](https://huggingface.co/briaai/BRIA-3.2).
|
||||
Github repo for Bria 3.2 can be found [here](https://github.com/Bria-AI/BRIA-3.2).
|
||||
|
||||
If you want to learn more about the Bria platform, and get free traril access, please visit [bria.ai](https://bria.ai).
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
_As the model is gated, before using it with diffusers you first need to go to the [Bria 3.2 Hugging Face page](https://huggingface.co/briaai/BRIA-3.2), fill in the form and accept the gate. Once you are in, you need to login so that your system knows you’ve accepted the gate._
|
||||
|
||||
Use the command below to log in:
|
||||
|
||||
```bash
|
||||
hf auth login
|
||||
```
|
||||
|
||||
|
||||
## BriaPipeline
|
||||
|
||||
[[autodoc]] BriaPipeline
|
||||
- all
|
||||
- __call__
|
||||
|
||||
@@ -37,7 +37,6 @@ The table below lists all the pipelines currently available in 🤗 Diffusers an
|
||||
| [AudioLDM2](audioldm2) | text2audio |
|
||||
| [AuraFlow](auraflow) | text2image |
|
||||
| [BLIP Diffusion](blip_diffusion) | text2image |
|
||||
| [Bria 3.2](bria_3_2) | text2image |
|
||||
| [CogVideoX](cogvideox) | text2video |
|
||||
| [Consistency Models](consistency_models) | unconditional image generation |
|
||||
| [ControlNet](controlnet) | text2image, image2image, inpainting |
|
||||
|
||||
@@ -15,73 +15,16 @@
|
||||
- local: using-diffusers/schedulers
|
||||
title: Load schedulers and models
|
||||
|
||||
- title: Inference
|
||||
isExpanded: false
|
||||
sections:
|
||||
- local: training/distributed_inference
|
||||
title: Distributed inference
|
||||
|
||||
- title: Inference optimization
|
||||
isExpanded: false
|
||||
sections:
|
||||
- local: optimization/fp16
|
||||
title: Accelerate inference
|
||||
- local: optimization/cache
|
||||
title: Caching
|
||||
- local: optimization/memory
|
||||
title: Reduce memory usage
|
||||
- local: optimization/speed-memory-optims
|
||||
title: Compile and offloading quantized models
|
||||
- title: Community optimizations
|
||||
sections:
|
||||
- local: optimization/pruna
|
||||
title: Pruna
|
||||
- local: optimization/xformers
|
||||
title: xFormers
|
||||
- local: optimization/tome
|
||||
title: Token merging
|
||||
- local: optimization/deepcache
|
||||
title: DeepCache
|
||||
- local: optimization/tgate
|
||||
title: TGATE
|
||||
- local: optimization/xdit
|
||||
title: xDiT
|
||||
- local: optimization/para_attn
|
||||
title: ParaAttention
|
||||
|
||||
- title: Hybrid Inference
|
||||
isExpanded: false
|
||||
sections:
|
||||
- local: hybrid_inference/overview
|
||||
title: Overview
|
||||
- local: hybrid_inference/vae_encode
|
||||
title: VAE Encode
|
||||
- local: hybrid_inference/api_reference
|
||||
title: API Reference
|
||||
|
||||
- title: Modular Diffusers
|
||||
isExpanded: false
|
||||
sections:
|
||||
- local: modular_diffusers/overview
|
||||
title: Overview
|
||||
- local: modular_diffusers/quickstart
|
||||
title: Quickstart
|
||||
- local: modular_diffusers/modular_diffusers_states
|
||||
title: States
|
||||
- local: modular_diffusers/pipeline_block
|
||||
title: ModularPipelineBlocks
|
||||
- local: modular_diffusers/sequential_pipeline_blocks
|
||||
title: SequentialPipelineBlocks
|
||||
- local: modular_diffusers/loop_sequential_pipeline_blocks
|
||||
title: LoopSequentialPipelineBlocks
|
||||
- local: modular_diffusers/auto_pipeline_blocks
|
||||
title: AutoPipelineBlocks
|
||||
- local: modular_diffusers/modular_pipeline
|
||||
title: ModularPipeline
|
||||
- local: modular_diffusers/components_manager
|
||||
title: ComponentsManager
|
||||
- local: modular_diffusers/guiders
|
||||
title: Guiders
|
||||
|
||||
- title: Training
|
||||
isExpanded: false
|
||||
@@ -94,20 +37,12 @@
|
||||
sections:
|
||||
- local: training/text2image
|
||||
title: Text-to-image
|
||||
- local: training/kandinsky
|
||||
title: Kandinsky 2.2
|
||||
- local: training/wuerstchen
|
||||
title: Wuerstchen
|
||||
- local: training/controlnet
|
||||
title: ControlNet
|
||||
- local: training/instructpix2pix
|
||||
title: InstructPix2Pix
|
||||
- title: Methods
|
||||
sections:
|
||||
- local: training/text_inversion
|
||||
title: Textual Inversion
|
||||
- local: training/dreambooth
|
||||
title: DreamBooth
|
||||
- local: training/lora
|
||||
title: LoRA
|
||||
|
||||
@@ -116,16 +51,6 @@
|
||||
sections:
|
||||
- local: optimization/onnx
|
||||
title: ONNX
|
||||
- local: optimization/open_vino
|
||||
title: OpenVINO
|
||||
- local: optimization/coreml
|
||||
title: Core ML
|
||||
- local: optimization/mps
|
||||
title: Metal Performance Shaders (MPS)
|
||||
- local: optimization/habana
|
||||
title: Intel Gaudi
|
||||
- local: optimization/neuron
|
||||
title: AWS Neuron
|
||||
|
||||
- title: Specific pipeline examples
|
||||
isExpanded: false
|
||||
@@ -138,8 +63,6 @@
|
||||
sections:
|
||||
- title: Task recipes
|
||||
sections:
|
||||
- local: community_projects
|
||||
title: Projects built with Diffusers
|
||||
- local: conceptual/philosophy
|
||||
title: Philosophy
|
||||
- local: conceptual/contribution
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
<!--版权 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证,版本2.0("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件是按"原样"分发的,没有任何形式的明示或暗示的担保或条件。有关许可证的特定语言,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 社区项目
|
||||
|
||||
欢迎来到社区项目。这个空间致力于展示我们充满活力的社区使用`diffusers`库创建的令人难以置信的工作和创新应用。
|
||||
|
||||
本节旨在:
|
||||
|
||||
- 突出使用`diffusers`构建的多样化和鼓舞人心的项目
|
||||
- 促进我们社区内的知识共享
|
||||
- 提供如何利用`diffusers`的实际例子
|
||||
|
||||
探索愉快,感谢您成为Diffusers社区的一部分!
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>项目名称</th>
|
||||
<th>描述</th>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/carson-katri/dream-textures"> dream-textures </a></td>
|
||||
<td>Stable Diffusion内置到Blender</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/megvii-research/HiDiffusion"> HiDiffusion </a></td>
|
||||
<td>仅通过添加一行代码即可提高扩散模型的分辨率和速度</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/lllyasviel/IC-Light"> IC-Light </a></td>
|
||||
<td>IC-Light是一个用于操作图像照明的项目</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/InstantID/InstantID"> InstantID </a></td>
|
||||
<td>InstantID:零样本身份保留生成在几秒钟内</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/Sanster/IOPaint"> IOPaint </a></td>
|
||||
<td>由SOTA AI模型驱动的图像修复工具。从您的图片中移除任何不需要的物体、缺陷、人物,或擦除并替换(由stable_diffusion驱动)图片上的任何内容。</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/bmaltais/kohya_ss"> Kohya </a></td>
|
||||
<td>Kohya的Stable Diffusion训练器的Gradio GUI</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/magic-research/magic-animate"> MagicAnimate </a></td>
|
||||
<td>MagicAnimate:使用扩散模型进行时间一致的人体图像动画</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/levihsu/OOTDiffusion"> OOTDiffusion </a></td>
|
||||
<td>基于潜在扩散的虚拟试穿控制</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/vladmandic/automatic"> SD.Next </a></td>
|
||||
<td>SD.Next: Stable Diffusion 和其他基于Diffusion的生成图像模型的高级实现</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/ashawkey/stable-dreamfusion"> stable-dreamfusion </a></td>
|
||||
<td>使用 NeRF + Diffusion 进行文本到3D & 图像到3D & 网格导出</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/HVision-NKU/StoryDiffusion"> StoryDiffusion </a></td>
|
||||
<td>StoryDiffusion 可以通过生成一致的图像和视频来创造一个神奇的故事。</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/cumulo-autumn/StreamDiffusion"> StreamDiffusion </a></td>
|
||||
<td>实时交互生成的管道级解决方案</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/Netwrck/stable-diffusion-server"> Stable Diffusion Server </a></td>
|
||||
<td>配置用于使用一个 stable diffusion 模型进行修复/生成/img2img 的服务器</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/suzukimain/auto_diffusers"> Model Search </a></td>
|
||||
<td>在 Civitai 和 Hugging Face 上搜索模型</td>
|
||||
</tr>
|
||||
<tr style="border-top: 2px solid black">
|
||||
<td><a href="https://github.com/beinsezii/skrample"> Skrample </a></td>
|
||||
<td>完全模块化的调度器功能,具有一流的 diffusers 集成。</td>
|
||||
</tr>
|
||||
</table>
|
||||
@@ -1,9 +0,0 @@
|
||||
# 混合推理 API 参考
|
||||
|
||||
## 远程解码
|
||||
|
||||
[[autodoc]] utils.remote_utils.remote_decode
|
||||
|
||||
## 远程编码
|
||||
|
||||
[[autodoc]] utils.remote_utils.remote_encode
|
||||
@@ -1,55 +0,0 @@
|
||||
<!--版权 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则不得使用此文件。
|
||||
您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,否则根据许可证分发的软件按"原样"分发,不附带任何明示或暗示的担保或条件。请参阅许可证以了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# 混合推理
|
||||
|
||||
**通过混合推理赋能本地 AI 构建者**
|
||||
|
||||
> [!TIP]
|
||||
> 混合推理是一项[实验性功能](https://huggingface.co/blog/remote_vae)。
|
||||
> 可以在此处提供反馈[此处](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml)。
|
||||
|
||||
## 为什么使用混合推理?
|
||||
|
||||
混合推理提供了一种快速简单的方式来卸载本地生成需求。
|
||||
|
||||
- 🚀 **降低要求:** 无需昂贵硬件即可访问强大模型。
|
||||
- 💎 **无妥协:** 在不牺牲性能的情况下实现最高质量。
|
||||
- 💰 **成本效益高:** 它是免费的!🤑
|
||||
- 🎯 **多样化用例:** 与 Diffusers � 和更广泛的社区完全兼容。
|
||||
- 🔧 **开发者友好:** 简单请求,快速响应。
|
||||
|
||||
---
|
||||
|
||||
## 可用模型
|
||||
|
||||
* **VAE 解码 🖼️:** 快速将潜在表示解码为高质量图像,不影响性能或工作流速度。
|
||||
* **VAE 编码 🔢:** 高效将图像编码为潜在表示,用于生成和训练。
|
||||
* **文本编码器 📃(即将推出):** 快速准确地计算提示的文本嵌入,确保流畅高质量的工作流。
|
||||
|
||||
---
|
||||
|
||||
## 集成
|
||||
|
||||
* **[SD.Next](https://github.com/vladmandic/sdnext):** 一体化 UI,直接支持混合推理。
|
||||
* **[ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae):** 用于混合推理的 ComfyUI 节点。
|
||||
|
||||
## 更新日志
|
||||
|
||||
- 2025 年 3 月 10 日:添加了 VAE 编码
|
||||
- 2025 年 3 月 2 日:初始发布,包含 VAE 解码
|
||||
|
||||
## 内容
|
||||
|
||||
文档分为三个部分:
|
||||
|
||||
* **VAE 解码** 学习如何使用混合推理进行 VAE 解码的基础知识。
|
||||
* **VAE 编码** 学习如何使用混合推理进行 VAE 编码的基础知识。
|
||||
* **API 参考** 深入了解任务特定设置和参数。
|
||||
@@ -1,184 +0,0 @@
|
||||
# 入门:使用混合推理进行 VAE 编码
|
||||
|
||||
VAE 编码用于训练、图像到图像和图像到视频——将图像或视频转换为潜在表示。
|
||||
|
||||
## 内存
|
||||
|
||||
这些表格展示了在不同 GPU 上使用 SD v1 和 SD XL 进行 VAE 编码的 VRAM 需求。
|
||||
|
||||
对于这些 GPU 中的大多数,内存使用百分比决定了其他模型(文本编码器、UNet/Transformer)必须被卸载,或者必须使用分块编码,这会增加时间并影响质量。
|
||||
|
||||
<details><summary>SD v1.5</summary>
|
||||
|
||||
| GPU | 分辨率 | 时间(秒) | 内存(%) | 分块时间(秒) | 分块内存(%) |
|
||||
|:------------------------------|:-------------|-----------------:|-------------:|--------------------:|-------------------:|
|
||||
| NVIDIA GeForce RTX 4090 | 512x512 | 0.015 | 3.51901 | 0.015 | 3.51901 |
|
||||
| NVIDIA GeForce RTX 4090 | 256x256 | 0.004 | 1.3154 | 0.005 | 1.3154 |
|
||||
| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.402 | 47.1852 | 0.496 | 3.51901 |
|
||||
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.078 | 12.2658 | 0.094 | 3.51901 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.023 | 5.30105 | 0.023 | 5.30105 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.006 | 1.98152 | 0.006 | 1.98152 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 0.574 | 71.08 | 0.656 | 5.30105 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.111 | 18.4772 | 0.14 | 5.30105 |
|
||||
| NVIDIA GeForce RTX 3090 | 512x512 | 0.032 | 3.52782 | 0.032 | 3.52782 |
|
||||
| NVIDIA GeForce RTX 3090 | 256x256 | 0.01 | 1.31869 | 0.009 | 1.31869 |
|
||||
| NVIDIA GeForce RTX 3090 | 2048x2048 | 0.742 | 47.3033 | 0.954 | 3.52782 |
|
||||
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.136 | 12.2965 | 0.207 | 3.52782 |
|
||||
| NVIDIA GeForce RTX 3080 | 512x512 | 0.036 | 8.51761 | 0.036 | 8.51761 |
|
||||
| NVIDIA GeForce RTX 3080 | 256x256 | 0.01 | 3.18387 | 0.01 | 3.18387 |
|
||||
| NVIDIA GeForce RTX 3080 | 2048x2048 | 0.863 | 86.7424 | 1.191 | 8.51761 |
|
||||
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.157 | 29.6888 | 0.227 | 8.51761 |
|
||||
| NVIDIA GeForce RTX 3070 | 512x512 | 0.051 | 10.6941 | 0.051 | 10.6941 |
|
||||
| NVIDIA GeForce RTX 3070 | 256x256 | 0.015 |
|
||||
| 3.99743 | 0.015 | 3.99743 |
|
||||
| NVIDIA GeForce RTX 3070 | 2048x2048 | 1.217 | 96.054 | 1.482 | 10.6941 |
|
||||
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.223 | 37.2751 | 0.327 | 10.6941 |
|
||||
|
||||
</details>
|
||||
|
||||
<details><summary>SDXL</summary>
|
||||
|
||||
| GPU | Resolution | Time (seconds) | Memory Consumed (%) | Tiled Time (seconds) | Tiled Memory (%) |
|
||||
|:------------------------------|:-------------|-----------------:|----------------------:|-----------------------:|-------------------:|
|
||||
| NVIDIA GeForce RTX 4090 | 512x512 | 0.029 | 4.95707 | 0.029 | 4.95707 |
|
||||
| NVIDIA GeForce RTX 4090 | 256x256 | 0.007 | 2.29666 | 0.007 | 2.29666 |
|
||||
| NVIDIA GeForce RTX 4090 | 2048x2048 | 0.873 | 66.3452 | 0.863 | 15.5649 |
|
||||
| NVIDIA GeForce RTX 4090 | 1024x1024 | 0.142 | 15.5479 | 0.143 | 15.5479 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 512x512 | 0.044 | 7.46735 | 0.044 | 7.46735 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 256x256 | 0.01 | 3.4597 | 0.01 | 3.4597 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 2048x2048 | 1.317 | 87.1615 | 1.291 | 23.447 |
|
||||
| NVIDIA GeForce RTX 4080 SUPER | 1024x1024 | 0.213 | 23.4215 | 0.214 | 23.4215 |
|
||||
| NVIDIA GeForce RTX 3090 | 512x512 | 0.058 | 5.65638 | 0.058 | 5.65638 |
|
||||
| NVIDIA GeForce RTX 3090 | 256x256 | 0.016 | 2.45081 | 0.016 | 2.45081 |
|
||||
| NVIDIA GeForce RTX 3090 | 2048x2048 | 1.755 | 77.8239 | 1.614 | 18.4193 |
|
||||
| NVIDIA GeForce RTX 3090 | 1024x1024 | 0.265 | 18.4023 | 0.265 | 18.4023 |
|
||||
| NVIDIA GeForce RTX 3080 | 512x512 | 0.064 | 13.6568 | 0.064 | 13.6568 |
|
||||
| NVIDIA GeForce RTX 3080 | 256x256 | 0.018 | 5.91728 | 0.018 | 5.91728 |
|
||||
| NVIDIA GeForce RTX 3080 | 2048x2048 | 内存不足 (OOM) | 内存不足 (OOM) | 1.866 | 44.4717 |
|
||||
| NVIDIA GeForce RTX 3080 | 1024x1024 | 0.302 | 44.4308 | 0.302 | 44.4308 |
|
||||
| NVIDIA GeForce RTX 3070 | 512x512 | 0.093 | 17.1465 | 0.093 | 17.1465 |
|
||||
| NVIDIA GeForce R
|
||||
| NVIDIA GeForce RTX 3070 | 256x256 | 0.025 | 7.42931 | 0.026 | 7.42931 |
|
||||
| NVIDIA GeForce RTX 3070 | 2048x2048 | OOM | OOM | 2.674 | 55.8355 |
|
||||
| NVIDIA GeForce RTX 3070 | 1024x1024 | 0.443 | 55.7841 | 0.443 | 55.7841 |
|
||||
|
||||
</details>
|
||||
|
||||
## 可用 VAE
|
||||
|
||||
| | **端点** | **模型** |
|
||||
|:-:|:-----------:|:--------:|
|
||||
| **Stable Diffusion v1** | [https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud](https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud) | [`stabilityai/sd-vae-ft-mse`](https://hf.co/stabilityai/sd-vae-ft-mse) |
|
||||
| **Stable Diffusion XL** | [https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud](https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud) | [`madebyollin/sdxl-vae-fp16-fix`](https://hf.co/madebyollin/sdxl-vae-fp16-fix) |
|
||||
| **Flux** | [https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud](https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud) | [`black-forest-labs/FLUX.1-schnell`](https://hf.co/black-forest-labs/FLUX.1-schnell) |
|
||||
|
||||
|
||||
> [!TIP]
|
||||
> 模型支持可以在此处请求:[这里](https://github.com/huggingface/diffusers/issues/new?template=remote-vae-pilot-feedback.yml)。
|
||||
|
||||
|
||||
## 代码
|
||||
|
||||
> [!TIP]
|
||||
> 从 `main` 安装 `diffusers` 以运行代码:`pip install git+https://github.com/huggingface/diffusers@main`
|
||||
|
||||
|
||||
一个辅助方法简化了与混合推理的交互。
|
||||
|
||||
```python
|
||||
from diffusers.utils.remote_utils import remote_encode
|
||||
```
|
||||
|
||||
### 基本示例
|
||||
|
||||
让我们编码一张图像,然后解码以演示。
|
||||
|
||||
<figure class="image flex flex-col items-center justify-center text-center m-0 w-full">
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"/>
|
||||
</figure>
|
||||
|
||||
<details><summary>代码</summary>
|
||||
|
||||
```python
|
||||
from diffusers.utils import load_image
|
||||
from diffusers.utils.remote_utils import remote_decode
|
||||
|
||||
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg?download=true")
|
||||
|
||||
latent = remote_encode(
|
||||
endpoint="https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud/",
|
||||
scaling_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
)
|
||||
|
||||
decoded = remote_decode(
|
||||
endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
|
||||
tensor=latent,
|
||||
scaling_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<figure class="image flex flex-col items-center justify-center text-center m-0 w-full">
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/remote_vae/decoded.png"/>
|
||||
</figure>
|
||||
|
||||
|
||||
### 生成
|
||||
|
||||
现在让我们看一个生成示例,我们将编码图像,生成,然后远程解码!
|
||||
|
||||
<details><summary>代码</summary>
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import StableDiffusionImg2ImgPip
|
||||
from diffusers.utils import load_image
|
||||
from diffusers.utils.remote_utils import remote_decode, remote_encode
|
||||
|
||||
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16",
|
||||
vae=None,
|
||||
).to("cuda")
|
||||
|
||||
init_image = load_image(
|
||||
"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
|
||||
)
|
||||
init_image = init_image.resize((768, 512))
|
||||
|
||||
init_latent = remote_encode(
|
||||
endpoint="https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud/",
|
||||
image=init_image,
|
||||
scaling_factor=0.18215,
|
||||
)
|
||||
|
||||
prompt = "A fantasy landscape, trending on artstation"
|
||||
latent = pipe(
|
||||
prompt=prompt,
|
||||
image=init_latent,
|
||||
strength=0.75,
|
||||
output_type="latent",
|
||||
).images
|
||||
|
||||
image = remote_decode(
|
||||
endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
|
||||
tensor=latent,
|
||||
scaling_factor=0.18215,
|
||||
)
|
||||
image.save("fantasy_landscape.jpg")
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<figure class="image flex flex-col items-center justify-center text-center m-0 w-full">
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/remote_vae/fantasy_landscape.png"/>
|
||||
</figure>
|
||||
|
||||
## 集成
|
||||
|
||||
* **[SD.Next](https://github.com/vladmandic/sdnext):** 具有直接支持混合推理功能的一体化用户界面。
|
||||
* **[ComfyUI-HFRemoteVae](https://github.com/kijai/ComfyUI-HFRemoteVae):** 用于混合推理的 ComfyUI 节点。
|
||||
@@ -1,156 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证的特定语言管理权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# AutoPipelineBlocks
|
||||
|
||||
[`~modular_pipelines.AutoPipelineBlocks`] 是一种包含支持不同工作流程的块的多块类型。它根据运行时提供的输入自动选择要运行的子块。这通常用于将多个工作流程(文本到图像、图像到图像、修复)打包到一个管道中以便利。
|
||||
|
||||
本指南展示如何创建 [`~modular_pipelines.AutoPipelineBlocks`]。
|
||||
|
||||
创建三个 [`~modular_pipelines.ModularPipelineBlocks`] 用于文本到图像、图像到图像和修复。这些代表了管道中可用的不同工作流程。
|
||||
|
||||
<hfoptions id="auto">
|
||||
<hfoption id="text-to-image">
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import ModularPipelineBlocks, InputParam, OutputParam
|
||||
|
||||
class TextToImageBlock(ModularPipelineBlocks):
|
||||
model_name = "text2img"
|
||||
|
||||
@property
|
||||
def inputs(self):
|
||||
return [InputParam(name="prompt")]
|
||||
|
||||
@property
|
||||
def intermediate_outputs(self):
|
||||
return []
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
return "我是一个文本到图像的工作流程!"
|
||||
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
print("运行文本到图像工作流程")
|
||||
# 在这里添加你的文本到图像逻辑
|
||||
# 例如:根据提示生成图像
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="image-to-image">
|
||||
|
||||
```py
|
||||
class ImageToImageBlock(ModularPipelineBlocks):
|
||||
model_name = "img2img"
|
||||
|
||||
@property
|
||||
def inputs(self):
|
||||
return [InputParam(name="prompt"), InputParam(name="image")]
|
||||
|
||||
@property
|
||||
def intermediate_outputs(self):
|
||||
return []
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
return "我是一个图像到图像的工作流程!"
|
||||
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
print("运行图像到图像工作流程")
|
||||
# 在这里添加你的图像到图像逻辑
|
||||
# 例如:根据提示转换输入图像
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="inpaint">
|
||||
|
||||
```py
|
||||
class InpaintBlock(ModularPipelineBlocks):
|
||||
model_name = "inpaint"
|
||||
|
||||
@property
|
||||
def inputs(self):
|
||||
return [InputParam(name="prompt"), InputParam(name="image"), InputParam(name="mask")]
|
||||
|
||||
@property
|
||||
|
||||
def intermediate_outputs(self):
|
||||
return []
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
return "我是一个修复工作流!"
|
||||
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
print("运行修复工作流")
|
||||
# 在这里添加你的修复逻辑
|
||||
# 例如:根据提示填充被遮罩的区域
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
创建一个包含子块类及其对应块名称列表的[`~modular_pipelines.AutoPipelineBlocks`]类。
|
||||
|
||||
你还需要包括`block_trigger_inputs`,一个触发相应块的输入名称列表。如果在运行时提供了触发输入,则选择该块运行。使用`None`来指定如果未检测到触发输入时运行的默认块。
|
||||
|
||||
最后,重要的是包括一个`description`,清楚地解释哪些输入触发哪些工作流。这有助于用户理解如何运行特定的工作流。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import AutoPipelineBlocks
|
||||
|
||||
class AutoImageBlocks(AutoPipelineBlocks):
|
||||
# 选择子块类的列表
|
||||
block_classes = [block_inpaint_cls, block_i2i_cls, block_t2i_cls]
|
||||
# 每个块的名称,顺序相同
|
||||
block_names = ["inpaint", "img2img", "text2img"]
|
||||
# 决定运行哪个块的触发输入
|
||||
# - "mask" 触发修复工作流
|
||||
# - "image" 触发img2img工作流(但仅在未提供mask时)
|
||||
# - 如果以上都没有,运行text2img工作流(默认)
|
||||
block_trigger_inputs = ["mask", "image", None]
|
||||
# 对于AutoPipelineBlocks来说,描述极其重要
|
||||
|
||||
def description(self):
|
||||
return (
|
||||
"Pipeline generates images given different types of conditions!\n"
|
||||
+ "This is an auto pipeline block that works for text2img, img2img and inpainting tasks.\n"
|
||||
+ " - inpaint workflow is run when `mask` is provided.\n"
|
||||
+ " - img2img workflow is run when `image` is provided (but only when `mask` is not provided).\n"
|
||||
+ " - text2img workflow is run when neither `image` nor `mask` is provided.\n"
|
||||
)
|
||||
```
|
||||
|
||||
包含`description`以避免任何关于如何运行块和需要什么输入的混淆**非常**重要。虽然[`~modular_pipelines.AutoPipelineBlocks`]很方便,但如果它没有正确解释,其条件逻辑可能难以理解。
|
||||
|
||||
创建`AutoImageBlocks`的一个实例。
|
||||
|
||||
```py
|
||||
auto_blocks = AutoImageBlocks()
|
||||
```
|
||||
|
||||
对于更复杂的组合,例如在更大的管道中作为子块使用的嵌套[`~modular_pipelines.AutoPipelineBlocks`]块,使用[`~modular_pipelines.SequentialPipelineBlocks.get_execution_blocks`]方法根据你的输入提取实际运行的块。
|
||||
|
||||
```py
|
||||
auto_blocks.get_execution_blocks("mask")
|
||||
```
|
||||
@@ -1,188 +0,0 @@
|
||||
<!--版权所有 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版("许可证")授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。请参阅许可证以了解特定语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# 组件管理器
|
||||
|
||||
[`ComponentsManager`] 是 Modular Diffusers 的模型注册和管理系统。它添加和跟踪模型,存储有用的元数据(模型大小、设备放置、适配器),防止重复模型实例,并支持卸载。
|
||||
|
||||
本指南将展示如何使用 [`ComponentsManager`] 来管理组件和设备内存。
|
||||
|
||||
## 添加组件
|
||||
|
||||
[`ComponentsManager`] 应与 [`ModularPipeline`] 一起创建,在 [`~ModularPipeline.from_pretrained`] 或 [`~ModularPipelineBlocks.init_pipeline`] 中。
|
||||
|
||||
> [!TIP]
|
||||
> `collection` 参数是可选的,但可以更轻松地组织和管理组件。
|
||||
|
||||
<hfoptions id="create">
|
||||
<hfoption id="from_pretrained">
|
||||
|
||||
```py
|
||||
from diffusers import ModularPipeline, ComponentsManager
|
||||
|
||||
comp = ComponentsManager()
|
||||
pipe = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test1")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="init_pipeline">
|
||||
|
||||
```py
|
||||
from diffusers import ComponentsManager
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
|
||||
|
||||
t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
|
||||
|
||||
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
|
||||
components = ComponentsManager()
|
||||
t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
组件仅在调用 [`~ModularPipeline.load_components`] 或 [`~ModularPipeline.load_default_components`] 时加载和注册。以下示例使用 [`~ModularPipeline.load_default_components`] 创建第二个管道,重用第一个管道的所有组件,并将其分配到不同的集合。
|
||||
|
||||
```py
|
||||
pipe.load_default_components()
|
||||
pipe2 = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test2")
|
||||
```
|
||||
|
||||
使用 [`~ModularPipeline.null_component_names`] 属性来识别需要加载的任何组件,使用 [`~ComponentsManager.get_components_by_names`] 检索它们,然后调用 [`~ModularPipeline.update_components`] 来添加缺失的组件。
|
||||
|
||||
```py
|
||||
pipe2.null_component_names
|
||||
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'image_encoder', 'unet', 'vae', 'scheduler', 'controlnet']
|
||||
|
||||
comp_dict = comp.get_components_by_names(names=pipe2.null_component_names)
|
||||
pipe2.update_components(**comp_dict)
|
||||
```
|
||||
|
||||
要添加单个组件,请使用 [`~ComponentsManager.add`] 方法。这会使用唯一 id 注册一个组件。
|
||||
|
||||
```py
|
||||
from diffusers import AutoModel
|
||||
|
||||
text_encoder = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
|
||||
component_id = comp.add("text_encoder", text_encoder)
|
||||
comp
|
||||
```
|
||||
|
||||
使用 [`~ComponentsManager.remove`] 通过其 id 移除一个组件。
|
||||
|
||||
```py
|
||||
comp.remove("text_encoder_139917733042864")
|
||||
```
|
||||
|
||||
## 检索组件
|
||||
|
||||
[`ComponentsManager`] 提供了几种方法来检索已注册的组件。
|
||||
|
||||
### get_one
|
||||
|
||||
[`~ComponentsManager.get_one`] 方法返回单个组件,并支持对 `name` 参数进行模式匹配。如果多个组件匹配,[`~ComponentsManager.get_one`] 会返回错误。
|
||||
|
||||
| 模式 | 示例 | 描述 |
|
||||
|-------------|----------------------------------|-------------------------------------------|
|
||||
| exact | `comp.get_one(name="unet")` | 精确名称匹配 |
|
||||
| wildcard | `comp.get_one(name="unet*")` | 名称以 "unet" 开头 |
|
||||
| exclusion | `comp.get_one(name="!unet")` | 排除名为 "unet" 的组件 |
|
||||
| or | `comp.get_one(name="unet|vae")` | 名称为 "unet" 或 "vae" |
|
||||
|
||||
[`~ComponentsManager.get_one`] 还通过 `collection` 参数或 `load_id` 参数过滤组件。
|
||||
|
||||
```py
|
||||
comp.get_one(name="unet", collection="sdxl")
|
||||
```
|
||||
|
||||
### get_components_by_names
|
||||
|
||||
[`~ComponentsManager.get_components_by_names`] 方法接受一个名称列表,并返回一个将名称映射到组件的字典。这在 [`ModularPipeline`] 中特别有用,因为它们提供了所需组件名称的列表,并且返回的字典可以直接传递给 [`~ModularPipeline.update_components`]。
|
||||
|
||||
```py
|
||||
component_dict = comp.get_components_by_names(names=["text_encoder", "unet", "vae"])
|
||||
{"text_encoder": component1, "unet": component2, "vae": component3}
|
||||
```
|
||||
|
||||
## 重复检测
|
||||
|
||||
建议使用 [`ComponentSpec`] 加载模型组件,以分配具有唯一 id 的组件,该 id 编码了它们的加载参数。这允许 [`ComponentsManager`] 自动检测并防止重复的模型实例,即使不同的对象代表相同的底层检查点。
|
||||
|
||||
```py
|
||||
from diffusers import ComponentSpec, ComponentsManager
|
||||
from transformers import CLIPTextModel
|
||||
|
||||
comp = ComponentsManager()
|
||||
|
||||
# 为第一个文本编码器创建 ComponentSpec
|
||||
spec = ComponentSpec(name="text_encoder", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=AutoModel)
|
||||
# 为重复的文本编码器创建 ComponentSpec(它是相同的检查点,来自相同的仓库/子文件夹)
|
||||
spec_duplicated = ComponentSpec(name="text_encoder_duplicated", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", ty
|
||||
pe_hint=CLIPTextModel)
|
||||
|
||||
# 加载并添加两个组件 - 管理器会检测到它们是同一个模型
|
||||
comp.add("text_encoder", spec.load())
|
||||
comp.add("text_encoder_duplicated", spec_duplicated.load())
|
||||
```
|
||||
|
||||
这会返回一个警告,附带移除重复项的说明。
|
||||
|
||||
```py
|
||||
ComponentsManager: adding component 'text_encoder_duplicated_139917580682672', but it has duplicate load_id 'stabilityai/stable-diffusion-xl-base-1.0|text_encoder|null|null' with existing components: text_encoder_139918506246832. To remove a duplicate, call `components_manager.remove('<component_id>')`.
|
||||
'text_encoder_duplicated_139917580682672'
|
||||
```
|
||||
|
||||
您也可以不使用 [`ComponentSpec`] 添加组件,并且在大多数情况下,即使您以不同名称添加相同组件,重复检测仍然有效。
|
||||
|
||||
然而,当您将相同组件加载到不同对象时,[`ComponentManager`] 无法检测重复项。在这种情况下,您应该使用 [`ComponentSpec`] 加载模型。
|
||||
|
||||
```py
|
||||
text_encoder_2 = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
|
||||
comp.add("text_encoder", text_encoder_2)
|
||||
'text_encoder_139917732983664'
|
||||
```
|
||||
|
||||
## 集合
|
||||
|
||||
集合是为组件分配的标签,用于更好的组织和管理。使用 [`~ComponentsManager.add`] 中的 `collection` 参数将组件添加到集合中。
|
||||
|
||||
每个集合中只允许每个名称有一个组件。添加第二个同名组件会自动移除第一个组件。
|
||||
|
||||
```py
|
||||
from diffusers import ComponentSpec, ComponentsManager
|
||||
|
||||
comp = ComponentsManager()
|
||||
# 为第一个 UNet 创建 ComponentSpec
|
||||
spec = ComponentSpec(name="unet", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", type_hint=AutoModel)
|
||||
# 为另一个 UNet 创建 ComponentSpec
|
||||
spec2 = ComponentSpec(name="unet", repo="RunDiffusion/Juggernaut-XL-v9", subfolder="unet", type_hint=AutoModel, variant="fp16")
|
||||
|
||||
# 将两个 UNet 添加到同一个集合 - 第二个将替换第一个
|
||||
comp.add("unet", spec.load(), collection="sdxl")
|
||||
comp.add("unet", spec2.load(), collection="sdxl")
|
||||
```
|
||||
|
||||
这使得在基于节点的系统中工作变得方便,因为您可以:
|
||||
|
||||
- 使用 `collection` 标签标记所有从一个节点加载的模型。
|
||||
- 当新检查点以相同名称加载时自动替换模型。
|
||||
- 当节点被移除时批量删除集合中的所有模型。
|
||||
|
||||
## 卸载
|
||||
|
||||
[`~ComponentsManager.enable_auto_cpu_offload`] 方法是一种全局卸载策略,适用于所有模型,无论哪个管道在使用它们。一旦启用,您无需担心设备放置,如果您添加或移除组件。
|
||||
|
||||
```py
|
||||
comp.enable_auto_cpu_offload(device="cuda")
|
||||
```
|
||||
|
||||
所有模型开始时都在 CPU 上,[`ComponentsManager`] 在需要它们之前将它们移动到适当的设备,并在 GPU 内存不足时将其他模型移回 CPU。
|
||||
|
||||
您可以设置自己的规则来决定哪些模型要卸载。
|
||||
@@ -1,173 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版("许可证")授权;除非遵守许可证,否则不得使用此文件。
|
||||
您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,不附带任何明示或暗示的担保或条件。请参阅许可证了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# 引导器
|
||||
|
||||
[Classifier-free guidance](https://huggingface.co/papers/2207.12598) 引导模型生成更好地匹配提示,通常用于提高生成质量、控制和提示的遵循度。有不同类型的引导方法,在 Diffusers 中,它们被称为*引导器*。与块类似,可以轻松切换和使用不同的引导器以适应不同的用例,而无需重写管道。
|
||||
|
||||
本指南将向您展示如何切换引导器、调整引导器参数,以及将它们加载并共享到 Hub。
|
||||
|
||||
## 切换引导器
|
||||
|
||||
[`ClassifierFreeGuidance`] 是默认引导器,在使用 [`~ModularPipelineBlocks.init_pipeline`] 初始化管道时创建。它通过 `from_config` 创建,这意味着它不需要从模块化存储库加载规范。引导器不会列在 `modular_model_index.json` 中。
|
||||
|
||||
使用 [`~ModularPipeline.get_component_spec`] 来检查引导器。
|
||||
|
||||
```py
|
||||
t2i_pipeline.get_component_spec("guider")
|
||||
ComponentSpec(name='guider', type_hint=<class 'diffusers.guiders.classifier_free_guidance.ClassifierFreeGuidance'>, description=None, config=FrozenDict([('guidance_scale', 7.5), ('guidance_rescale', 0.0), ('use_original_formulation', False), ('start', 0.0), ('stop', 1.0), ('_use_default_values', ['start', 'guidance_rescale', 'stop', 'use_original_formulation'])]), repo=None, subfolder=None, variant=None, revision=None, default_creation_method='from_config')
|
||||
```
|
||||
|
||||
通过将新引导器传递给 [`~ModularPipeline.update_components`] 来切换到不同的引导器。
|
||||
|
||||
> [!TIP]
|
||||
> 更改引导器将返回文本,让您知道您正在更改引导器类型。
|
||||
> ```bash
|
||||
> ModularPipeline.update_components: 添加具有新类型的引导器: PerturbedAttentionGuidance, 先前类型: ClassifierFreeGuidance
|
||||
> ```
|
||||
|
||||
```py
|
||||
from diffusers import LayerSkipConfig, PerturbedAttentionGuidance
|
||||
|
||||
config = LayerSkipConfig(indices=[2, 9], fqn="mid_block.attentions.0.transformer_blocks", skip_attention=False, skip_attention_scores=True, skip_ff=False)
|
||||
guider = PerturbedAttentionGuidance(
|
||||
guidance_scale=5.0, perturbed_guidance_scale=2.5, perturbed_guidance_config=config
|
||||
)
|
||||
t2i_pipeline.update_components(guider=guider)
|
||||
```
|
||||
|
||||
再次使用 [`~ModularPipeline.get_component_spec`] 来验证引导器类型是否不同。
|
||||
|
||||
```py
|
||||
t2i_pipeline.get_component_spec("guider")
|
||||
ComponentSpec(name='guider', type_hint=<class 'diffusers.guiders.perturbed_attention_guidance.PerturbedAttentionGuidance'>, description=None, config=FrozenDict([('guidance_scale', 5.0), ('perturbed_guidance_scale', 2.5), ('perturbed_guidance_start', 0.01), ('perturbed_guidance_stop', 0.2), ('perturbed_guidance_layers', None), ('perturbed_guidance_config', LayerSkipConfig(indices=[2, 9], fqn='mid_block.attentions.0.transformer_blocks', skip_attention=False, skip_attention_scores=True, skip_ff=False, dropout=1.0)), ('guidance_rescale', 0.0), ('use_original_formulation', False), ('start', 0.0), ('stop', 1.0), ('_use_default_values', ['perturbed_guidance_start', 'use_original_formulation', 'perturbed_guidance_layers', 'stop', 'start', 'guidance_rescale', 'perturbed_guidance_stop']), ('_class_name', 'PerturbedAttentionGuidance'), ('_diffusers_version', '0.35.0.dev0')]), repo=None, subfolder=None, variant=None, revision=None, default_creation_method='from_config')
|
||||
```
|
||||
|
||||
## 加载自定义引导器
|
||||
|
||||
已经在 Hub 上保存并带有 `modular_model_index.json` 文件的引导器现在被视为 `from_pretrained` 组件,而不是 `from_config` 组件。
|
||||
|
||||
```json
|
||||
{
|
||||
"guider": [
|
||||
null,
|
||||
null,
|
||||
{
|
||||
"repo": "YiYiXu/modular-loader-t2i-guider",
|
||||
"revision": null,
|
||||
"subfolder": "pag_guider",
|
||||
"type_hint": [
|
||||
"diffusers",
|
||||
"PerturbedAttentionGuidance"
|
||||
],
|
||||
"variant": null
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
引导器只有在调用 [`~ModularPipeline.load_default_components`] 之后才会创建,基于 `modular_model_index.json` 中的加载规范。
|
||||
|
||||
```py
|
||||
t2i_pipeline = t2i_blocks.init_pipeline("YiYiXu/modular-doc-guider")
|
||||
# 在初始化时未创建
|
||||
assert t2i_pipeline.guider is None
|
||||
t2i_pipeline.load_default_components()
|
||||
# 加载为 PAG 引导器
|
||||
t2i_pipeline.guider
|
||||
```
|
||||
|
||||
## 更改引导器参数
|
||||
|
||||
引导器参数可以通过 [`~ComponentSpec.create`] 方法或 [`~ModularPipeline.update_components`] 方法进行调整。下面的示例更改了 `guidance_scale` 值。
|
||||
|
||||
<hfoptions id="switch">
|
||||
<hfoption id="create">
|
||||
|
||||
```py
|
||||
guider_spec = t2i_pipeline.get_component_spec("guider")
|
||||
guider = guider_spec.create(guidance_scale=10)
|
||||
t2i_pipeline.update_components(guider=guider)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="update_components">
|
||||
|
||||
```py
|
||||
guider_spec = t2i_pipeline.get_component_spec("guider")
|
||||
guider_spec.config["guidance_scale"] = 10
|
||||
t2i_pipeline.update_components(guider=guider_spec)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 上传自定义引导器
|
||||
|
||||
在自定义引导器上调用 [`~utils.PushToHubMixin.push_to_hub`] 方法,将其分享到 Hub。
|
||||
|
||||
```py
|
||||
guider.push_to_hub("YiYiXu/modular-loader-t2i-guider", subfolder="pag_guider")
|
||||
```
|
||||
|
||||
要使此引导器可用于管道,可以修改 `modular_model_index.json` 文件或使用 [`~ModularPipeline.update_components`] 方法。
|
||||
|
||||
<hfoptions id="upload">
|
||||
<hfoption id="modular_model_index.json">
|
||||
|
||||
编辑 `modular_model_index.json` 文件,并添加引导器的加载规范,指向包含引导器配置的文件夹
|
||||
例如。
|
||||
|
||||
```json
|
||||
{
|
||||
"guider": [
|
||||
"diffusers",
|
||||
"PerturbedAttentionGuidance",
|
||||
{
|
||||
"repo": "YiYiXu/modular-loader-t2i-guider",
|
||||
"revision": null,
|
||||
"subfolder": "pag_guider",
|
||||
"type_hint": [
|
||||
"diffusers",
|
||||
"PerturbedAttentionGuidance"
|
||||
],
|
||||
"variant": null
|
||||
}
|
||||
],
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="update_components">
|
||||
|
||||
将 [`~ComponentSpec.default_creation_method`] 更改为 `from_pretrained` 并使用 [`~ModularPipeline.update_components`] 来更新引导器和组件规范以及管道配置。
|
||||
|
||||
> [!TIP]
|
||||
> 更改创建方法将返回文本,告知您正在将创建类型更改为 `from_pretrained`。
|
||||
> ```bash
|
||||
> ModularPipeline.update_components: 将引导器的 default_creation_method 从 from_config 更改为 from_pretrained。
|
||||
> ```
|
||||
|
||||
```py
|
||||
guider_spec = t2i_pipeline.get_component_spec("guider")
|
||||
guider_spec.default_creation_method="from_pretrained"
|
||||
guider_spec.repo="YiYiXu/modular-loader-t2i-guider"
|
||||
guider_spec.subfolder="pag_guider"
|
||||
pag_guider = guider_spec.load()
|
||||
t2i_pipeline.update_components(guider=pag_guider)
|
||||
```
|
||||
|
||||
要使其成为管道的默认引导器,请调用 [`~utils.PushToHubMixin.push_to_hub`]。这是一个可选步骤,如果您仅在本地进行实验,则不需要。
|
||||
|
||||
```py
|
||||
t2i_pipeline.push_to_hub("YiYiXu/modular-doc-guider")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
@@ -1,93 +0,0 @@
|
||||
<!--版权 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。请参阅许可证了解
|
||||
特定语言下的权限和限制。
|
||||
-->
|
||||
|
||||
# LoopSequentialPipelineBlocks
|
||||
|
||||
[`~modular_pipelines.LoopSequentialPipelineBlocks`] 是一种多块类型,它将其他 [`~modular_pipelines.ModularPipelineBlocks`] 以循环方式组合在一起。数据循环流动,使用 `intermediate_inputs` 和 `intermediate_outputs`,并且每个块都是迭代运行的。这通常用于创建一个默认是迭代的去噪循环。
|
||||
|
||||
本指南向您展示如何创建 [`~modular_pipelines.LoopSequentialPipelineBlocks`]。
|
||||
|
||||
## 循环包装器
|
||||
|
||||
[`~modular_pipelines.LoopSequentialPipelineBlocks`],也被称为 *循环包装器*,因为它定义了循环结构、迭代变量和配置。在循环包装器内,您需要以下变量。
|
||||
|
||||
- `loop_inputs` 是用户提供的值,等同于 [`~modular_pipelines.ModularPipelineBlocks.inputs`]。
|
||||
- `loop_intermediate_inputs` 是来自 [`~modular_pipelines.PipelineState`] 的中间变量,等同于 [`~modular_pipelines.ModularPipelineBlocks.intermediate_inputs`]。
|
||||
- `loop_intermediate_outputs` 是由块创建并添加到 [`~modular_pipelines.PipelineState`] 的新中间变量。它等同于 [`~modular_pipelines.ModularPipelineBlocks.intermediate_outputs`]。
|
||||
- `__call__` 方法定义了循环结构和迭代逻辑。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import LoopSequentialPipelineBlocks, ModularPipelineBlocks, InputParam, OutputParam
|
||||
|
||||
class LoopWrapper(LoopSequentialPipelineBlocks):
|
||||
model_name = "test"
|
||||
@property
|
||||
def description(self):
|
||||
return "I'm a loop!!"
|
||||
@property
|
||||
def loop_inputs(self):
|
||||
return [InputParam(name="num_steps")]
|
||||
@torch.no_grad()
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
# 循环结构 - 可以根据您的需求定制
|
||||
for i in range(block_state.num_steps):
|
||||
# loop_step 按顺序执行所有注册的块
|
||||
components, block_state = self.loop_step(components, block_state, i=i)
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
循环包装器可以传递额外的参数,如当前迭代索引,到循环块。
|
||||
|
||||
## 循环块
|
||||
|
||||
循环块是一个 [`~modular_pipelines.ModularPipelineBlocks`],但 `__call__` 方法的行为不同。
|
||||
|
||||
- 它从循环包装器。
|
||||
- 它直接与[`~modular_pipelines.BlockState`]一起工作,而不是[`~modular_pipelines.PipelineState`]。
|
||||
- 它不需要检索或更新[`~modular_pipelines.BlockState`]。
|
||||
|
||||
循环块共享相同的[`~modular_pipelines.BlockState`],以允许值在循环的每次迭代中累积和变化。
|
||||
|
||||
```py
|
||||
class LoopBlock(ModularPipelineBlocks):
|
||||
model_name = "test"
|
||||
@property
|
||||
def inputs(self):
|
||||
return [InputParam(name="x")]
|
||||
@property
|
||||
def intermediate_outputs(self):
|
||||
# 这个块产生的输出
|
||||
return [OutputParam(name="x")]
|
||||
@property
|
||||
def description(self):
|
||||
return "我是一个在`LoopWrapper`类内部使用的块"
|
||||
def __call__(self, components, block_state, i: int):
|
||||
block_state.x += 1
|
||||
return components, block_state
|
||||
```
|
||||
|
||||
## LoopSequentialPipelineBlocks
|
||||
|
||||
使用[`~modular_pipelines.LoopSequentialPipelineBlocks.from_blocks_dict`]方法将循环块添加到循环包装器中,以创建[`~modular_pipelines.LoopSequentialPipelineBlocks`]。
|
||||
|
||||
```py
|
||||
loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock})
|
||||
```
|
||||
|
||||
添加更多的循环块以在每次迭代中运行,使用[`~modular_pipelines.LoopSequentialPipelineBlocks.from_blocks_dict`]。这允许您在不改变循环逻辑本身的情况下修改块。
|
||||
|
||||
```py
|
||||
loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock(), "block2": LoopBlock})
|
||||
```
|
||||
@@ -1,74 +0,0 @@
|
||||
<!--版权 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版("许可证")授权;除非符合许可证的规定,否则不得使用此文件。
|
||||
您可以在以下网址获取许可证的副本
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件是基于"按原样"分发的,没有任何形式的明示或暗示的担保或条件。有关许可证下特定的语言管理权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 状态
|
||||
|
||||
块依赖于[`~modular_pipelines.PipelineState`]和[`~modular_pipelines.BlockState`]数据结构进行通信和数据共享。
|
||||
|
||||
| 状态 | 描述 |
|
||||
|-------|-------------|
|
||||
| [`~modular_pipelines.PipelineState`] | 维护管道执行所需的整体数据,并允许块读取和更新其数据。 |
|
||||
| [`~modular_pipelines.BlockState`] | 允许每个块使用来自`inputs`的必要数据执行其计算 |
|
||||
|
||||
本指南解释了状态如何工作以及它们如何连接块。
|
||||
|
||||
## PipelineState
|
||||
|
||||
[`~modular_pipelines.PipelineState`]是所有块的全局状态容器。它维护管道的完整运行时状态,并为块提供了一种结构化的方式来读取和写入共享数据。
|
||||
|
||||
[`~modular_pipelines.PipelineState`]中有两个字典用于结构化数据。
|
||||
|
||||
- `values`字典是一个**可变**状态,包含用户提供的输入值的副本和由块生成的中间输出值。如果一个块修改了一个`input`,它将在调用`set_block_state`后反映在`values`字典中。
|
||||
|
||||
```py
|
||||
PipelineState(
|
||||
values={
|
||||
'prompt': 'a cat'
|
||||
'guidance_scale': 7.0
|
||||
'num_inference_steps': 25
|
||||
'prompt_embeds': Tensor(dtype=torch.float32, shape=torch.Size([1, 1, 1, 1]))
|
||||
'negative_prompt_embeds': None
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
## BlockState
|
||||
|
||||
[`~modular_pipelines.BlockState`]是[`~modular_pipelines.PipelineState`]中相关变量的局部视图,单个块需要这些变量来执行其计算。
|
||||
|
||||
直接作为属性访问这些变量,如`block_state.image`。
|
||||
|
||||
```py
|
||||
BlockState(
|
||||
image: <PIL.Image.Image image mode=RGB size=512x512 at 0x7F3ECC494640>
|
||||
)
|
||||
```
|
||||
|
||||
当一个块的`__call__`方法被执行时,它用`self.get_block_state(state)`检索[`BlockState`],执行其操作,并用`self.set_block_state(state, block_state)`更新[`~modular_pipelines.PipelineState`]。
|
||||
|
||||
```py
|
||||
def __call__(self, components, state):
|
||||
# 检索BlockState
|
||||
block_state = self.get_block_state(state)
|
||||
|
||||
# 对输入进行计算的逻辑
|
||||
|
||||
# 更新PipelineState
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
## 状态交互
|
||||
|
||||
[`~modular_pipelines.PipelineState`]和[`~modular_pipelines.BlockState`]的交互由块的`inputs`和`intermediate_outputs`定义。
|
||||
|
||||
- `inputs`,
|
||||
一个块可以修改输入 - 比如 `block_state.image` - 并且这个改变可以通过调用 `set_block_state` 全局传播到 [`~modular_pipelines.PipelineState`]。
|
||||
- `intermediate_outputs`,是一个块创建的新变量。它被添加到 [`~modular_pipelines.PipelineState`] 的 `values` 字典中,并且可以作为后续块的可用变量,或者由用户作为管道的最终输出访问。
|
||||
@@ -1,358 +0,0 @@
|
||||
<!--版权 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版(“许可证”)授权;除非符合许可证的规定,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件是基于“按原样”分发的,没有任何形式的明示或暗示的保证或条件。有关许可证的特定语言,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 模块化管道
|
||||
|
||||
[`ModularPipeline`] 将 [`~modular_pipelines.ModularPipelineBlocks`] 转换为可执行的管道,加载模型并执行块中定义的计算步骤。它是运行管道的主要接口,与 [`DiffusionPipeline`] API 非常相似。
|
||||
|
||||
主要区别在于在管道中包含了一个预期的 `output` 参数。
|
||||
|
||||
<hfoptions id="example">
|
||||
<hfoption id="text-to-image">
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
|
||||
|
||||
blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
|
||||
|
||||
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
|
||||
pipeline = blocks.init_pipeline(modular_repo_id)
|
||||
|
||||
pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
pipeline.to("cuda")
|
||||
|
||||
image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", output="images")[0]
|
||||
image.save("modular_t2i_out.png")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="image-to-image">
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import IMAGE2IMAGE_BLOCKS
|
||||
|
||||
blocks = SequentialPipelineBlocks.from_blocks_dict(IMAGE2IMAGE_BLOCKS)
|
||||
|
||||
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
|
||||
pipeline = blocks.init_pipeline(modular_repo_id)
|
||||
|
||||
pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
pipeline.to("cuda")
|
||||
|
||||
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
|
||||
init_image = load_image(url)
|
||||
prompt = "a dog catching a frisbee in the jungle"
|
||||
image = pipeline(prompt=prompt, image=init_image, strength=0.8, output="images")[0]
|
||||
image.save("modular_i2i_out.png")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="inpainting">
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import INPAINT_BLOCKS
|
||||
from diffusers.utils import load_image
|
||||
|
||||
blocks = SequentialPipelineBlocks.from_blocks_dict(INPAINT_BLOCKS)
|
||||
|
||||
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
|
||||
pipeline = blocks.init_pipeline(modular_repo_id)
|
||||
|
||||
pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
pipeline.to("cuda")
|
||||
|
||||
img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
|
||||
mask_url = "h
|
||||
ttps://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-inpaint-mask.png"
|
||||
|
||||
init_image = load_image(img_url)
|
||||
mask_image = load_image(mask_url)
|
||||
|
||||
prompt = "A deep sea diver floating"
|
||||
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85, output="images")[0]
|
||||
image.save("moduar_inpaint_out.png")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
本指南将向您展示如何创建一个[`ModularPipeline`]并管理其中的组件。
|
||||
|
||||
## 添加块
|
||||
|
||||
块是[`InsertableDict`]对象,可以在特定位置插入,提供了一种灵活的方式来混合和匹配块。
|
||||
|
||||
使用[`~modular_pipelines.modular_pipeline_utils.InsertableDict.insert`]在块类或`sub_blocks`属性上添加一个块。
|
||||
|
||||
```py
|
||||
# BLOCKS是块类的字典,您需要向其中添加类
|
||||
BLOCKS.insert("block_name", BlockClass, index)
|
||||
# sub_blocks属性包含实例,向该属性添加一个块实例
|
||||
t2i_blocks.sub_blocks.insert("block_name", block_instance, index)
|
||||
```
|
||||
|
||||
使用[`~modular_pipelines.modular_pipeline_utils.InsertableDict.pop`]在块类或`sub_blocks`属性上移除一个块。
|
||||
|
||||
```py
|
||||
# 从预设中移除一个块类
|
||||
BLOCKS.pop("text_encoder")
|
||||
# 分离出一个块实例
|
||||
text_encoder_block = t2i_blocks.sub_blocks.pop("text_encoder")
|
||||
```
|
||||
|
||||
通过将现有块设置为新块来交换块。
|
||||
|
||||
```py
|
||||
# 在预设中替换块类
|
||||
BLOCKS["prepare_latents"] = CustomPrepareLatents
|
||||
# 使用块实例在sub_blocks属性中替换
|
||||
t2i_blocks.sub_blocks["prepare_latents"] = CustomPrepareLatents()
|
||||
```
|
||||
|
||||
## 创建管道
|
||||
|
||||
有两种方法可以创建一个[`ModularPipeline`]。从[`ModularPipelineBlocks`]组装并创建管道,或使用[`~ModularPipeline.from_pretrained`]加载现有管道。
|
||||
|
||||
您还应该初始化一个[`ComponentsManager`]来处理设备放置和内存以及组件管理。
|
||||
|
||||
> [!TIP]
|
||||
> 有关它如何帮助管理不同工作流中的组件的更多详细信息,请参阅[ComponentsManager](./components_manager)文档。
|
||||
|
||||
<hfoptions id="create">
|
||||
<hfoption id="ModularPipelineBlocks">
|
||||
|
||||
使用[`~ModularPipelineBlocks.init_pipeline`]方法从组件和配置规范创建一个[`ModularPipeline`]。此方法从`modular_model_index.json`文件加载*规范*,但尚未加载*模型*。
|
||||
|
||||
```py
|
||||
from diffusers import ComponentsManager
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
|
||||
|
||||
t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
|
||||
|
||||
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
|
||||
components = ComponentsManager()
|
||||
t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="from_pretrained">
|
||||
|
||||
[`~ModularPipeline.from_pretrained`]方法创建一个[`ModularPipeline`]从Hub上的模块化仓库加载。
|
||||
|
||||
```py
|
||||
from diffusers import ModularPipeline, ComponentsManager
|
||||
|
||||
components = ComponentsManager()
|
||||
pipeline = ModularPipeline.from_pretrained("YiYiXu/modular-loader-t2i-0704", components_manager=components)
|
||||
```
|
||||
|
||||
添加`trust_remote_code`参数以加载自定义的[`ModularPipeline`]。
|
||||
|
||||
```py
|
||||
from diffusers import ModularPipeline, ComponentsManager
|
||||
|
||||
components = ComponentsManager()
|
||||
modular_repo_id = "YiYiXu/modular-diffdiff-0704"
|
||||
diffdiff_pipeline = ModularPipeline.from_pretrained(modular_repo_id, trust_remote_code=True, components_manager=components)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 加载组件
|
||||
|
||||
一个[`ModularPipeline`]不会自动实例化组件。它只加载配置和组件规范。您可以使用[`~ModularPipeline.load_default_components`]加载所有组件,或仅使用[`~ModularPipeline.load_components`]加载特定组件。
|
||||
|
||||
<hfoptions id="load">
|
||||
<hfoption id="load_default_components">
|
||||
|
||||
```py
|
||||
import torch
|
||||
|
||||
t2i_pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
t2i_pipeline.to("cuda")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="load_components">
|
||||
|
||||
下面的例子仅加载UNet和VAE。
|
||||
|
||||
```py
|
||||
import torch
|
||||
|
||||
t2i_pipeline.load_components(names=["unet", "vae"], torch_dtype=torch.float16)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
打印管道以检查加载的预训练组件。
|
||||
|
||||
```py
|
||||
t2i_pipeline
|
||||
```
|
||||
|
||||
这应该与管道初始化自的模块化仓库中的`modular_model_index.json`文件匹配。如果管道不需要某个组件,即使它在模块化仓库中存在,也不会被包含。
|
||||
|
||||
要修改组件加载的来源,编辑仓库中的`modular_model_index.json`文件,并将其更改为您希望的加载路径。下面的例子从不同的仓库加载UNet。
|
||||
|
||||
```json
|
||||
# 原始
|
||||
"unet": [
|
||||
null, null,
|
||||
{
|
||||
"repo": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"subfolder": "unet",
|
||||
"variant": "fp16"
|
||||
}
|
||||
]
|
||||
|
||||
# 修改后
|
||||
"unet": [
|
||||
null, null,
|
||||
{
|
||||
"repo": "RunDiffusion/Juggernaut-XL-v9",
|
||||
"subfolder": "unet",
|
||||
"variant": "fp16"
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
### 组件加载状态
|
||||
|
||||
下面的管道属性提供了关于哪些组件被加载的更多信息。
|
||||
|
||||
使用`component_names`返回所有预期的组件。
|
||||
|
||||
```py
|
||||
t2i_pipeline.component_names
|
||||
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'guider', 'scheduler', 'unet', 'vae', 'image_processor']
|
||||
```
|
||||
|
||||
使用`null_component_names`返回尚未加载的组件。使用[`~ModularPipeline.from_pretrained`]加载这些组件。
|
||||
|
||||
```py
|
||||
t2i_pipeline.null_component_names
|
||||
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler']
|
||||
```
|
||||
|
||||
使用`pretrained_component_names`返回将从预训练模型加载的组件。
|
||||
|
||||
```py
|
||||
t2i_pipeline.pretrained_component_names
|
||||
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'scheduler', 'unet', 'vae']
|
||||
```
|
||||
|
||||
使用 `config_component_names` 返回那些使用默认配置创建的组件(不是从模块化仓库加载的)。来自配置的组件不包括在内,因为它们已经在管道创建期间初始化。这就是为什么它们没有列在 `null_component_names` 中。
|
||||
|
||||
```py
|
||||
t2i_pipeline.config_component_names
|
||||
['guider', 'image_processor']
|
||||
```
|
||||
|
||||
## 更新组件
|
||||
|
||||
根据组件是*预训练组件*还是*配置组件*,组件可能会被更新。
|
||||
|
||||
> [!WARNING]
|
||||
> 在更新组件时,组件可能会从预训练变为配置。组件类型最初是在块的 `expected_components` 字段中定义的。
|
||||
|
||||
预训练组件通过 [`ComponentSpec`] 更新,而配置组件则通过直接传递对象或使用 [`ComponentSpec`] 更新。
|
||||
|
||||
[`ComponentSpec`] 对于预训练组件显示 `default_creation_method="from_pretrained"`,对于配置组件显示 `default_creation_method="from_config`。
|
||||
|
||||
要更新预训练组件,创建一个 [`ComponentSpec`],指定组件的名称和从哪里加载它。使用 [`~ComponentSpec.load`] 方法来加载组件。
|
||||
|
||||
```py
|
||||
from diffusers import ComponentSpec, UNet2DConditionModel
|
||||
|
||||
unet_spec = ComponentSpec(name="unet",type_hint=UNet2DConditionModel, repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", variant="fp16")
|
||||
unet = unet_spec.load(torch_dtype=torch.float16)
|
||||
```
|
||||
|
||||
[`~ModularPipeline.update_components`] 方法用一个新的组件替换原来的组件。
|
||||
|
||||
```py
|
||||
t2i_pipeline.update_components(unet=unet2)
|
||||
```
|
||||
|
||||
当组件被更新时,加载规范也会在管道配置中更新。
|
||||
|
||||
### 组件提取和修改
|
||||
|
||||
当你使用 [`~ComponentSpec.load`] 时,新组件保持其加载规范。这使得提取规范并重新创建组件成为可能。
|
||||
|
||||
```py
|
||||
spec = ComponentSpec.from_component("unet", unet2)
|
||||
spec
|
||||
ComponentSpec(name='unet', type_hint=<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>, description=None, config=None, repo='stabilityai/stable-diffusion-xl-base-1.0', subfolder='unet', variant='fp16', revision=None, default_creation_method='from_pretrained')
|
||||
unet2_recreated = spec.load(torch_dtype=torch.float16)
|
||||
```
|
||||
|
||||
[`~ModularPipeline.get_component_spec`] 方法获取当前组件规范的副本以进行修改或更新。
|
||||
|
||||
```py
|
||||
unet_spec = t2i_pipeline.get_component_spec("unet")
|
||||
unet_spec
|
||||
ComponentSpec(
|
||||
name='unet',
|
||||
type_hint=<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>,
|
||||
repo='RunDiffusion/Juggernaut-XL-v9',
|
||||
subfolder='unet',
|
||||
variant='fp16',
|
||||
default_creation_method='from_pretrained'
|
||||
)
|
||||
|
||||
# 修改以从不同的仓库加载
|
||||
unet_spec.repo = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
# 使用修改后的规范加载组件
|
||||
unet = unet_spec.load(torch_dtype=torch.float16)
|
||||
```
|
||||
|
||||
## 模块化仓库
|
||||
一个仓库
|
||||
如果管道块使用*预训练组件*,则需要y。该存储库提供了加载规范和元数据。
|
||||
|
||||
[`ModularPipeline`]特别需要*模块化存储库*(参见[示例存储库](https://huggingface.co/YiYiXu/modular-diffdiff)),这比典型的存储库更灵活。它包含一个`modular_model_index.json`文件,包含以下3个元素。
|
||||
|
||||
- `library`和`class`显示组件是从哪个库加载的及其类。如果是`null`,则表示组件尚未加载。
|
||||
- `loading_specs_dict`包含加载组件所需的信息,例如从中加载的存储库和子文件夹。
|
||||
|
||||
与标准存储库不同,模块化存储库可以根据`loading_specs_dict`从不同的存储库获取组件。组件不需要存在于同一个存储库中。
|
||||
|
||||
模块化存储库可能包含用于加载[`ModularPipeline`]的自定义代码。这允许您使用不是Diffusers原生的专用块。
|
||||
|
||||
```
|
||||
modular-diffdiff-0704/
|
||||
├── block.py # 自定义管道块实现
|
||||
├── config.json # 管道配置和auto_map
|
||||
└── modular_model_index.json # 组件加载规范
|
||||
```
|
||||
|
||||
[config.json](https://huggingface.co/YiYiXu/modular-diffdiff-0704/blob/main/config.json)文件包含一个`auto_map`键,指向`block.py`中定义自定义块的位置。
|
||||
|
||||
```json
|
||||
{
|
||||
"_class_name": "DiffDiffBlocks",
|
||||
"auto_map": {
|
||||
"ModularPipelineBlocks": "block.DiffDiffBlocks"
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -1,38 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在以下位置获取许可证的副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证下特定语言的权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 概述
|
||||
|
||||
> [!WARNING]
|
||||
> 模块化Diffusers正在积极开发中,其API可能会发生变化。
|
||||
|
||||
模块化Diffusers是一个统一的管道系统,通过*管道块*简化您的工作流程。
|
||||
|
||||
- 块是可重用的,您只需要为您的管道创建独特的块。
|
||||
- 块可以混合搭配,以适应或为特定工作流程或多个工作流程创建管道。
|
||||
|
||||
模块化Diffusers文档的组织如下所示。
|
||||
|
||||
## 快速开始
|
||||
|
||||
- 一个[快速开始](./quickstart)演示了如何使用模块化Diffusers实现一个示例工作流程。
|
||||
|
||||
## ModularPipelineBlocks
|
||||
|
||||
- [States](./modular_diffusers_states)解释了数据如何在块和[`ModularPipeline`]之间共享和通信。
|
||||
- [ModularPipelineBlocks](./pipeline_block)是[`ModularPipeline`]最基本的单位,本指南向您展示如何创建一个。
|
||||
- [SequentialPipelineBlocks](./sequential_pipeline_blocks)是一种类型的块,它将多个块链接起来,使它们一个接一个地运行,沿着链传递数据。本指南向您展示如何创建[`~modular_pipelines.SequentialPipelineBlocks`]以及它们如何连接和一起工作。
|
||||
- [LoopSequentialPipelineBlocks](./loop_sequential_pipeline_blocks)是一种类型的块,它在循环中运行一系列块。本指南向您展示如何创建[`~modular_pipelines.LoopSequentialPipelineBlocks`]。
|
||||
- [AutoPipelineBlocks](./auto_pipeline_blocks)是一种类型的块,它根据输入自动选择要运行的块。本指南向您展示如何创建[`~modular_pipelines.AutoPipelineBlocks`]。
|
||||
|
||||
## ModularPipeline
|
||||
|
||||
- [ModularPipeline](./modular_pipeline)向您展示如何创建并将管道块转换为可执行的[`ModularPipeline`]。
|
||||
- [ComponentsManager](./components_manager)向您展示如何跨多个管道管理和重用组件。
|
||||
- [Guiders](./guiders)向您展示如何在管道中使用不同的指导方法。
|
||||
@@ -1,114 +0,0 @@
|
||||
<!--版权 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版(“许可证”)授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件是基于“按原样”基础分发的,没有任何明示或暗示的保证或条件。请参阅许可证了解特定语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# ModularPipelineBlocks
|
||||
|
||||
[`~modular_pipelines.ModularPipelineBlocks`] 是构建 [`ModularPipeline`] 的基本块。它定义了管道中特定步骤应执行的组件、输入/输出和计算。一个 [`~modular_pipelines.ModularPipelineBlocks`] 与其他块连接,使用 [状态](./modular_diffusers_states),以实现工作流的模块化构建。
|
||||
|
||||
单独的 [`~modular_pipelines.ModularPipelineBlocks`] 无法执行。它是管道中步骤应执行的操作的蓝图。要实际运行和执行管道,需要将 [`~modular_pipelines.ModularPipelineBlocks`] 转换为 [`ModularPipeline`]。
|
||||
|
||||
本指南将向您展示如何创建 [`~modular_pipelines.ModularPipelineBlocks`]。
|
||||
|
||||
## 输入和输出
|
||||
|
||||
> [!TIP]
|
||||
> 如果您不熟悉Modular Diffusers中状态的工作原理,请参考 [States](./modular_diffusers_states) 指南。
|
||||
|
||||
一个 [`~modular_pipelines.ModularPipelineBlocks`] 需要 `inputs` 和 `intermediate_outputs`。
|
||||
|
||||
- `inputs` 是由用户提供并从 [`~modular_pipelines.PipelineState`] 中检索的值。这很有用,因为某些工作流会调整图像大小,但仍需要原始图像。 [`~modular_pipelines.PipelineState`] 维护原始图像。
|
||||
|
||||
使用 `InputParam` 定义 `inputs`。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import InputParam
|
||||
|
||||
user_inputs = [
|
||||
InputParam(name="image", type_hint="PIL.Image", description="要处理的原始输入图像")
|
||||
]
|
||||
```
|
||||
|
||||
- `intermediate_inputs` 通常由前一个块创建的值,但如果前面的块没有生成它们,也可以直接提供。与 `inputs` 不同,`intermediate_inputs` 可以被修改。
|
||||
|
||||
使用 `InputParam` 定义 `intermediate_inputs`。
|
||||
|
||||
```py
|
||||
user_intermediate_inputs = [
|
||||
InputParam(name="processed_image", type_hint="torch.Tensor", description="image that has been preprocessed and normalized"),
|
||||
]
|
||||
```
|
||||
|
||||
- `intermediate_outputs` 是由块创建并添加到 [`~modular_pipelines.PipelineState`] 的新值。`intermediate_outputs` 可作为后续块的 `intermediate_inputs` 使用,或作为运行管道的最终输出使用。
|
||||
|
||||
使用 `OutputParam` 定义 `intermediate_outputs`。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import OutputParam
|
||||
|
||||
user_intermediate_outputs = [
|
||||
OutputParam(name="image_latents", description="latents representing the image")
|
||||
]
|
||||
```
|
||||
|
||||
中间输入和输出共享数据以连接块。它们可以在任何时候访问,允许你跟踪工作流的进度。
|
||||
|
||||
## 计算逻辑
|
||||
|
||||
一个块执行的计算在`__call__`方法中定义,它遵循特定的结构。
|
||||
|
||||
1. 检索[`~modular_pipelines.BlockState`]以获取`inputs`和`intermediate_inputs`的局部视图。
|
||||
2. 在`inputs`和`intermediate_inputs`上实现计算逻辑。
|
||||
3. 更新[`~modular_pipelines.PipelineState`]以将局部[`~modular_pipelines.BlockState`]的更改推送回全局[`~modular_pipelines.PipelineState`]。
|
||||
4. 返回对下一个块可用的组件和状态。
|
||||
|
||||
```py
|
||||
def __call__(self, components, state):
|
||||
# 获取该块需要的状态变量的局部视图
|
||||
block_state = self.get_block_state(state)
|
||||
|
||||
# 你的计算逻辑在这里
|
||||
# block_state包含你所有的inputs和intermediate_inputs
|
||||
# 像这样访问它们: block_state.image, block_state.processed_image
|
||||
|
||||
# 用你更新的block_states更新管道状态
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
### 组件和配置
|
||||
|
||||
块需要的组件和管道级别的配置在[`ComponentSpec`]和[`~modular_pipelines.ConfigSpec`]中指定。
|
||||
|
||||
- [`ComponentSpec`]包含块使用的预期组件。你需要组件的`name`和理想情况下指定组件确切是什么的`type_hint`。
|
||||
- [`~modular_pipelines.ConfigSpec`]包含控制所有块行为的管道级别设置。
|
||||
|
||||
```py
|
||||
from diffusers import ComponentSpec, ConfigSpec
|
||||
|
||||
expected_components = [
|
||||
ComponentSpec(name="unet", type_hint=UNet2DConditionModel),
|
||||
ComponentSpec(name="scheduler", type_hint=EulerDiscreteScheduler)
|
||||
]
|
||||
|
||||
expected_config = [
|
||||
ConfigSpec("force_zeros_for_empty_prompt", True)
|
||||
]
|
||||
```
|
||||
|
||||
当块被转换为管道时,组件作为`__call__`中的第一个参数对块可用。
|
||||
|
||||
```py
|
||||
def __call__(self, components, state):
|
||||
# 使用点符号访问组件
|
||||
unet = components.unet
|
||||
vae = components.vae
|
||||
scheduler = components.scheduler
|
||||
```
|
||||
@@ -1,346 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证下特定语言的管理权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 快速入门
|
||||
|
||||
模块化Diffusers是一个快速构建灵活和可定制管道的框架。模块化Diffusers的核心是[`ModularPipelineBlocks`],可以与其他块组合以适应新的工作流程。这些块被转换为[`ModularPipeline`],一个开发者可以使用的友好用户界面。
|
||||
|
||||
本文档将向您展示如何使用模块化框架实现[Differential Diffusion](https://differential-diffusion.github.io/)管道。
|
||||
|
||||
## ModularPipelineBlocks
|
||||
|
||||
[`ModularPipelineBlocks`]是*定义*,指定管道中单个步骤的组件、输入、输出和计算逻辑。有四种类型的块。
|
||||
|
||||
- [`ModularPipelineBlocks`]是最基本的单一步骤块。
|
||||
- [`SequentialPipelineBlocks`]是一个多块,线性组合其他块。一个块的输出是下一个块的输入。
|
||||
- [`LoopSequentialPipelineBlocks`]是一个多块,迭代运行,专为迭代工作流程设计。
|
||||
- [`AutoPipelineBlocks`]是一个针对不同工作流程的块集合,它根据输入选择运行哪个块。它旨在方便地将多个工作流程打包到单个管道中。
|
||||
|
||||
[Differential Diffusion](https://differential-diffusion.github.io/)是一个图像到图像的工作流程。从`IMAGE2IMAGE_BLOCKS`预设开始,这是一个用于图像到图像生成的`ModularPipelineBlocks`集合。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl import IMAGE2IMAGE_BLOCKS
|
||||
IMAGE2IMAGE_BLOCKS = InsertableDict([
|
||||
("text_encoder", StableDiffusionXLTextEncoderStep),
|
||||
("image_encoder", StableDiffusionXLVaeEncoderStep),
|
||||
("input", StableDiffusionXLInputStep),
|
||||
("set_timesteps", StableDiffusionXLImg2ImgSetTimestepsStep),
|
||||
("prepare_latents", StableDiffusionXLImg2ImgPrepareLatentsStep),
|
||||
("prepare_add_cond", StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep),
|
||||
("denoise", StableDiffusionXLDenoiseStep),
|
||||
("decode", StableDiffusionXLDecodeStep)
|
||||
])
|
||||
```
|
||||
|
||||
## 管道和块状态
|
||||
|
||||
模块化Diffusers使用*状态*在块之间通信数据。有两种类型的状态。
|
||||
|
||||
- [`PipelineState`]是一个全局状态,可用于跟踪所有块的所有输入和输出。
|
||||
- [`BlockState`]是[`PipelineState`]中相关变量的局部视图,用于单个块。
|
||||
|
||||
## 自定义块
|
||||
|
||||
[Differential Diffusion](https://differential-diffusion.github.io/) 与标准的图像到图像转换在其 `prepare_latents` 和 `denoise` 块上有所不同。所有其他块都可以重用,但你需要修改这两个。
|
||||
|
||||
通过复制和修改现有的块,为 `prepare_latents` 和 `denoise` 创建占位符 `ModularPipelineBlocks`。
|
||||
|
||||
打印 `denoise` 块,可以看到它由 [`LoopSequentialPipelineBlocks`] 组成,包含三个子块,`before_denoiser`、`denoiser` 和 `after_denoiser`。只需要修改 `before_denoiser` 子块,根据变化图为去噪器准备潜在输入。
|
||||
|
||||
```py
|
||||
denoise_blocks = IMAGE2IMAGE_BLOCKS["denoise"]()
|
||||
print(denoise_blocks)
|
||||
```
|
||||
|
||||
用新的 `SDXLDiffDiffLoopBeforeDenoiser` 块替换 `StableDiffusionXLLoopBeforeDenoiser` 子块。
|
||||
|
||||
```py
|
||||
# 复制现有块作为占位符
|
||||
class SDXLDiffDiffPrepareLatentsStep(ModularPipelineBlocks):
|
||||
"""Copied from StableDiffusionXLImg2ImgPrepareLatentsStep - will modify later"""
|
||||
# ... 与 StableDiffusionXLImg2ImgPrepareLatentsStep 相同的实现
|
||||
|
||||
class SDXLDiffDiffDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
|
||||
block_classes = [SDXLDiffDiffLoopBeforeDenoiser, StableDiffusionXLLoopDenoiser, StableDiffusionXLLoopAfterDenoiser]
|
||||
block_names = ["before_denoiser", "denoiser", "after_denoiser"]
|
||||
```
|
||||
|
||||
### prepare_latents
|
||||
|
||||
`prepare_latents` 块需要进行以下更改。
|
||||
|
||||
- 一个处理器来处理变化图
|
||||
- 一个新的 `inputs` 来接受用户提供的变化图,`timestep` 用于预计算所有潜在变量和 `num_inference_steps` 来创建更新图像区域的掩码
|
||||
- 更新 `__call__` 方法中的计算,用于处理变化图和创建掩码,并将其存储在 [`BlockState`] 中
|
||||
|
||||
```diff
|
||||
class SDXLDiffDiffPrepareLatentsStep(ModularPipelineBlocks):
|
||||
@property
|
||||
def expected_components(self) -> List[ComponentSpec]:
|
||||
return [
|
||||
ComponentSpec("vae", AutoencoderKL),
|
||||
ComponentSpec("scheduler", EulerDiscreteScheduler),
|
||||
+ ComponentSpec("mask_processor", VaeImageProcessor, config=FrozenDict({"do_normalize": False, "do_convert_grayscale": True}))
|
||||
]
|
||||
@property
|
||||
def inputs(self) -> List[Tuple[str, Any]]:
|
||||
return [
|
||||
InputParam("generator"),
|
||||
+ InputParam("diffdiff_map", required=True),
|
||||
- InputParam("latent_timestep", required=True, type_hint=torch.Tensor),
|
||||
+ InputParam("timesteps", type_hint=torch.Tensor),
|
||||
+ InputParam("num_inference_steps", type_hint=int),
|
||||
]
|
||||
|
||||
@property
|
||||
def intermediate_outputs(self) -> List[OutputParam]:
|
||||
return [
|
||||
+ OutputParam("original_latents", type_hint=torch.Tensor),
|
||||
+ OutputParam("diffdiff_masks", type_hint=torch.Tensor),
|
||||
]
|
||||
def __call__(self, components, state: PipelineState):
|
||||
# ... existing logic ...
|
||||
+ # Process change map and create masks
|
||||
+ diffdiff_map = components.mask_processor.preprocess(block_state.diffdiff_map, height=latent_height, width=latent_width)
|
||||
+ thresholds = torch.arange(block_state.num_inference_steps, dtype=diffdiff_map.dtype) / block_state.num_inference_steps
|
||||
+ block_state.diffdiff_masks = diffdiff_map > (thresholds + (block_state.denoising_start or 0))
|
||||
+ block_state.original_latents = block_state.latents
|
||||
```
|
||||
|
||||
### 去噪
|
||||
|
||||
`before_denoiser` 子块需要进行以下更改。
|
||||
|
||||
- 新的 `inputs` 以接受 `denoising_start` 参数,`original_latents` 和 `diffdiff_masks` 来自 `prepare_latents` 块
|
||||
- 更新 `__call__` 方法中的计算以应用 Differential Diffusion
|
||||
|
||||
```diff
|
||||
class SDXLDiffDiffLoopBeforeDenoiser(ModularPipelineBlocks):
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return (
|
||||
"Step within the denoising loop for differential diffusion that prepare the latent input for the denoiser"
|
||||
)
|
||||
|
||||
@property
|
||||
def inputs(self) -> List[str]:
|
||||
return [
|
||||
InputParam("latents", required=True, type_hint=torch.Tensor),
|
||||
+ InputParam("denoising_start"),
|
||||
+ InputParam("original_latents", type_hint=torch.Tensor),
|
||||
+ InputParam("diffdiff_masks", type_hint=torch.Tensor),
|
||||
]
|
||||
|
||||
def __call__(self, components, block_state, i, t):
|
||||
+ # Apply differential diffusion logic
|
||||
+ if i == 0 and block_state.denoising_start is None:
|
||||
+ block_state.latents = block_state.original_latents[:1]
|
||||
+ else:
|
||||
+ block_state.mask = block_state.diffdiff_masks[i].unsqueeze(0).unsqueeze(1)
|
||||
+ block_state.latents = block_state.original_latents[i] * block_state.mask + block_state.latents * (1 - block_state.mask)
|
||||
|
||||
# ... rest of existing logic ...
|
||||
```
|
||||
|
||||
## 组装块
|
||||
|
||||
此时,您应该拥有创建 [`ModularPipeline`] 所需的所有块。
|
||||
|
||||
复制现有的 `IMAGE2IMAGE_BLOCKS` 预设,对于 `set_timesteps` 块,使用 `TEXT2IMAGE_BLOCKS` 中的 `set_timesteps`,因为 Differential Diffusion 不需要 `strength` 参数。
|
||||
|
||||
将 `prepare_latents` 和 `denoise` 块设置为您刚刚修改的 `SDXLDiffDiffPrepareLatentsStep` 和 `SDXLDiffDiffDenoiseStep` 块。
|
||||
|
||||
调用 [`SequentialPipelineBlocks.from_blocks_dict`] 在块上创建一个 `SequentialPipelineBlocks`。
|
||||
|
||||
```py
|
||||
DIFFDIFF_BLOCKS = IMAGE2IMAGE_BLOCKS.copy()
|
||||
DIFFDIFF_BLOCKS["set_timesteps"] = TEXT2IMAGE_BLOCKS["set_timesteps"]
|
||||
DIFFDIFF_BLOCKS["prepare_latents"] = SDXLDiffDiffPrepareLatentsStep
|
||||
DIFFDIFF_BLOCKS["denoise"] = SDXLDiffDiffDenoiseStep
|
||||
|
||||
dd_blocks = SequentialPipelineBlocks.from_blocks_dict(DIFFDIFF_BLOCKS)
|
||||
print(dd_blocks)
|
||||
```
|
||||
|
||||
## ModularPipeline
|
||||
|
||||
将 [`SequentialPipelineBlocks`] 转换为 [`ModularPipeline`],使用 [`ModularPipeline.init_pipeline`] 方法。这会初始化从 `modular_model_index.json` 文件加载的预期组件。通过调用 [`ModularPipeline.load_defau
|
||||
lt_components`]。
|
||||
|
||||
初始化[`ComponentManager`]时传入pipeline是一个好主意,以帮助管理不同的组件。一旦调用[`~ModularPipeline.load_default_components`],组件就会被注册到[`ComponentManager`]中,并且可以在工作流之间共享。下面的例子使用`collection`参数为组件分配了一个`"diffdiff"`标签,以便更好地组织。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import ComponentsManager
|
||||
|
||||
components = ComponentManager()
|
||||
|
||||
dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", components_manager=components, collection="diffdiff")
|
||||
dd_pipeline.load_default_componenets(torch_dtype=torch.float16)
|
||||
dd_pipeline.to("cuda")
|
||||
```
|
||||
|
||||
## 添加工作流
|
||||
|
||||
可以向[`ModularPipeline`]添加其他工作流以支持更多功能,而无需从头重写整个pipeline。
|
||||
|
||||
本节演示如何添加IP-Adapter或ControlNet。
|
||||
|
||||
### IP-Adapter
|
||||
|
||||
Stable Diffusion XL已经有一个预设的IP-Adapter块,你可以使用,并且不需要对现有的Differential Diffusion pipeline进行任何更改。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl.encoders import StableDiffusionXLAutoIPAdapterStep
|
||||
|
||||
ip_adapter_block = StableDiffusionXLAutoIPAdapterStep()
|
||||
```
|
||||
|
||||
使用[`sub_blocks.insert`]方法将其插入到[`ModularPipeline`]中。下面的例子在位置`0`插入了`ip_adapter_block`。打印pipeline可以看到`ip_adapter_block`被添加了,并且它需要一个`ip_adapter_image`。这也向pipeline添加了两个组件,`image_encoder`和`feature_extractor`。
|
||||
|
||||
```py
|
||||
dd_blocks.sub_blocks.insert("ip_adapter", ip_adapter_block, 0)
|
||||
```
|
||||
|
||||
调用[`~ModularPipeline.init_pipeline`]来初始化一个[`ModularPipeline`],并使用[`~ModularPipeline.load_default_components`]加载模型组件。加载并设置IP-Adapter以运行pipeline。
|
||||
|
||||
```py
|
||||
dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff")
|
||||
dd_pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
dd_pipeline.loader.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
|
||||
dd_pipeline.loader.set_ip_adapter_scale(0.6)
|
||||
dd_pipeline = dd_pipeline.to(device)
|
||||
|
||||
ip_adapter_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/diffdiff_orange.jpeg")
|
||||
image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/20240329211129_4024911930.png?download=true")
|
||||
mask = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/gradient_mask.png?download=true")
|
||||
|
||||
prompt = "a green pear"
|
||||
negative_prompt = "blurry"
|
||||
generator = torch.Generator(device=device).manual_seed(42)
|
||||
|
||||
image = dd_pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_inference_steps=25,
|
||||
generator=generator,
|
||||
ip_adapter_image=ip_adapter_image,
|
||||
diffdiff_map=mask,
|
||||
image=image,
|
||||
|
||||
output="images"
|
||||
)[0]
|
||||
```
|
||||
|
||||
### ControlNet
|
||||
|
||||
Stable Diffusion XL 已经预设了一个可以立即使用的 ControlNet 块。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines.stable_diffusion_xl.modular_blocks import StableDiffusionXLAutoControlNetInputStep
|
||||
|
||||
control_input_block = StableDiffusionXLAutoControlNetInputStep()
|
||||
```
|
||||
|
||||
然而,它需要修改 `denoise` 块,因为那是 ControlNet 将控制信息注入到 UNet 的地方。
|
||||
|
||||
通过将 `StableDiffusionXLLoopDenoiser` 子块替换为 `StableDiffusionXLControlNetLoopDenoiser` 来修改 `denoise` 块。
|
||||
|
||||
```py
|
||||
class SDXLDiffDiffControlNetDenoiseStep(StableDiffusionXLDenoiseLoopWrapper):
|
||||
block_classes = [SDXLDiffDiffLoopBeforeDenoiser, StableDiffusionXLControlNetLoopDenoiser, StableDiffusionXLDenoiseLoopAfterDenoiser]
|
||||
block_names = ["before_denoiser", "denoiser", "after_denoiser"]
|
||||
|
||||
controlnet_denoise_block = SDXLDiffDiffControlNetDenoiseStep()
|
||||
```
|
||||
|
||||
插入 `controlnet_input` 块并用新的 `controlnet_denoise_block` 替换 `denoise` 块。初始化一个 [`ModularPipeline`] 并将 [`~ModularPipeline.load_default_components`] 加载到其中。
|
||||
|
||||
```py
|
||||
dd_blocks.sub_blocks.insert("controlnet_input", control_input_block, 7)
|
||||
dd_blocks.sub_blocks["denoise"] = controlnet_denoise_block
|
||||
|
||||
dd_pipeline = dd_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff")
|
||||
dd_pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
dd_pipeline = dd_pipeline.to(device)
|
||||
|
||||
control_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/diffdiff_tomato_canny.jpeg")
|
||||
image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/20240329211129_4024911930.png?download=true")
|
||||
mask = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/differential/gradient_mask.png?download=true")
|
||||
|
||||
prompt = "a green pear"
|
||||
negative_prompt = "blurry"
|
||||
generator = torch.Generator(device=device).manual_seed(42)
|
||||
|
||||
image = dd_pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_inference_steps=25,
|
||||
generator=generator,
|
||||
control_image=control_image,
|
||||
controlnet_conditioning_scale=0.5,
|
||||
diffdiff_map=mask,
|
||||
image=image,
|
||||
output="images"
|
||||
)[0]
|
||||
```
|
||||
|
||||
### AutoPipelineBlocks
|
||||
|
||||
差分扩散、IP-Adapter 和 ControlNet 工作流可以通过使用 [`AutoPipelineBlocks`] 捆绑到一个单一的 [`ModularPipeline`] 中。这允许根据输入如 `control_image` 或 `ip_adapter_image` 自动选择要运行的子块。如果没有传递这些输入,则默认为差分扩散。
|
||||
|
||||
使用 `block_trigger_inputs` 仅在提供 `control_image` 输入时运行 `SDXLDiffDiffControlNetDenoiseStep` 块。否则,使用 `SDXLDiffDiffDenoiseStep`。
|
||||
|
||||
```py
|
||||
class SDXLDiffDiffAutoDenoiseStep(AutoPipelineBlocks):
|
||||
block_classes = [SDXLDiffDiffControlNetDenoiseStep, SDXLDiffDiffDenoiseStep]
|
||||
block_names = ["contr
|
||||
olnet_denoise", "denoise"]
|
||||
block_trigger_inputs = ["controlnet_cond", None]
|
||||
```
|
||||
|
||||
添加 `ip_adapter` 和 `controlnet_input` 块。
|
||||
|
||||
```py
|
||||
DIFFDIFF_AUTO_BLOCKS = IMAGE2IMAGE_BLOCKS.copy()
|
||||
DIFFDIFF_AUTO_BLOCKS["prepare_latents"] = SDXLDiffDiffPrepareLatentsStep
|
||||
DIFFDIFF_AUTO_BLOCKS["set_timesteps"] = TEXT2IMAGE_BLOCKS["set_timesteps"]
|
||||
DIFFDIFF_AUTO_BLOCKS["denoise"] = SDXLDiffDiffAutoDenoiseStep
|
||||
DIFFDIFF_AUTO_BLOCKS.insert("ip_adapter", StableDiffusionXLAutoIPAdapterStep, 0)
|
||||
DIFFDIFF_AUTO_BLOCKS.insert("controlnet_input",StableDiffusionXLControlNetAutoInput, 7)
|
||||
```
|
||||
|
||||
调用 [`SequentialPipelineBlocks.from_blocks_dict`] 来创建一个 [`SequentialPipelineBlocks`] 并创建一个 [`ModularPipeline`] 并加载模型组件以运行。
|
||||
|
||||
```py
|
||||
dd_auto_blocks = SequentialPipelineBlocks.from_blocks_dict(DIFFDIFF_AUTO_BLOCKS)
|
||||
dd_pipeline = dd_auto_blocks.init_pipeline("YiYiXu/modular-demo-auto", collection="diffdiff")
|
||||
dd_pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
```
|
||||
|
||||
## 分享
|
||||
|
||||
使用 [`~ModularPipeline.save_pretrained`] 将您的 [`ModularPipeline`] 添加到 Hub,并将 `push_to_hub` 参数设置为 `True`。
|
||||
|
||||
```py
|
||||
dd_pipeline.save_pretrained("YiYiXu/test_modular_doc", push_to_hub=True)
|
||||
```
|
||||
|
||||
其他用户可以使用 [`~ModularPipeline.from_pretrained`] 加载 [`ModularPipeline`]。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import ModularPipeline, ComponentsManager
|
||||
|
||||
components = ComponentsManager()
|
||||
|
||||
diffdiff_pipeline = ModularPipeline.from_pretrained("YiYiXu/modular-diffdiff-0704", trust_remote_code=True, components_manager=components, collection="diffdiff")
|
||||
diffdiff_pipeline.load_default_components(torch_dtype=torch.float16)
|
||||
```
|
||||
@@ -1,112 +0,0 @@
|
||||
<!--版权 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据Apache许可证2.0版("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
获取许可证的副本。
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件是基于"按原样"基础分发的,没有任何形式的明示或暗示的保证或条件。有关许可证下特定语言的管理权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 顺序管道块
|
||||
|
||||
[`~modular_pipelines.SequentialPipelineBlocks`] 是一种多块类型,它将其他 [`~modular_pipelines.ModularPipelineBlocks`] 按顺序组合在一起。数据通过 `intermediate_inputs` 和 `intermediate_outputs` 线性地从一个块流向下一个块。[`~modular_pipelines.SequentialPipelineBlocks`] 中的每个块通常代表管道中的一个步骤,通过组合它们,您逐步构建一个管道。
|
||||
|
||||
本指南向您展示如何将两个块连接成一个 [`~modular_pipelines.SequentialPipelineBlocks`]。
|
||||
|
||||
创建两个 [`~modular_pipelines.ModularPipelineBlocks`]。第一个块 `InputBlock` 输出一个 `batch_size` 值,第二个块 `ImageEncoderBlock` 使用 `batch_size` 作为 `intermediate_inputs`。
|
||||
|
||||
<hfoptions id="sequential">
|
||||
<hfoption id="InputBlock">
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import ModularPipelineBlocks, InputParam, OutputParam
|
||||
|
||||
class InputBlock(ModularPipelineBlocks):
|
||||
|
||||
@property
|
||||
def inputs(self):
|
||||
return [
|
||||
InputParam(name="prompt", type_hint=list, description="list of text prompts"),
|
||||
InputParam(name="num_images_per_prompt", type_hint=int, description="number of images per prompt"),
|
||||
]
|
||||
|
||||
@property
|
||||
def intermediate_outputs(self):
|
||||
return [
|
||||
OutputParam(name="batch_size", description="calculated batch size"),
|
||||
]
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
return "A block that determines batch_size based on the number of prompts and num_images_per_prompt argument."
|
||||
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
batch_size = len(block_state.prompt)
|
||||
block_state.batch_size = batch_size * block_state.num_images_per_prompt
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="ImageEncoderBlock">
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers.modular_pipelines import ModularPipelineBlocks, InputParam, OutputParam
|
||||
|
||||
class ImageEncoderBlock(ModularPipelineBlocks):
|
||||
|
||||
@property
|
||||
def inputs(self):
|
||||
return [
|
||||
InputParam(name="image", type_hint="PIL.Image", description="raw input image to process"),
|
||||
InputParam(name="batch_size", type_hint=int),
|
||||
]
|
||||
|
||||
@property
|
||||
def intermediate_outputs(self):
|
||||
return [
|
||||
OutputParam(name="image_latents", description="latents representing the image"
|
||||
]
|
||||
|
||||
@property
|
||||
def description(self):
|
||||
return "Encode raw image into its latent presentation"
|
||||
|
||||
def __call__(self, components, state):
|
||||
block_state = self.get_block_state(state)
|
||||
# 模拟处理图像
|
||||
# 这将改变所有块的图像状态,从PIL图像变为张量
|
||||
block_state.image = torch.randn(1, 3, 512, 512)
|
||||
block_state.batch_size = block_state.batch_size * 2
|
||||
block_state.image_latents = torch.randn(1, 4, 64, 64)
|
||||
self.set_block_state(state, block_state)
|
||||
return components, state
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
通过定义一个[`InsertableDict`]来连接两个块,将块名称映射到块实例。块按照它们在`blocks_dict`中注册的顺序执行。
|
||||
|
||||
使用[`~modular_pipelines.SequentialPipelineBlocks.from_blocks_dict`]来创建一个[`~modular_pipelines.SequentialPipelineBlocks`]。
|
||||
|
||||
```py
|
||||
from diffusers.modular_pipelines import SequentialPipelineBlocks, InsertableDict
|
||||
|
||||
blocks_dict = InsertableDict()
|
||||
blocks_dict["input"] = input_block
|
||||
blocks_dict["image_encoder"] = image_encoder_block
|
||||
|
||||
blocks = SequentialPipelineBlocks.from_blocks_dict(blocks_dict)
|
||||
```
|
||||
|
||||
通过调用`blocks`来检查[`~modular_pipelines.SequentialPipelineBlocks`]中的子块,要获取更多关于输入和输出的详细信息,可以访问`docs`属性。
|
||||
|
||||
```py
|
||||
print(blocks)
|
||||
print(blocks.doc)
|
||||
```
|
||||
@@ -1,67 +0,0 @@
|
||||
<!-- 版权所有 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本(“许可证”)授权;除非符合许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,否则根据许可证分发的软件按“原样”分发,不附带任何明示或暗示的担保或条件。请参阅许可证以了解具体的语言管理权限和限制。 -->
|
||||
|
||||
# 缓存
|
||||
|
||||
缓存通过存储和重用不同层的中间输出(如注意力层和前馈层)来加速推理,而不是在每个推理步骤执行整个计算。它显著提高了生成速度,但以更多内存为代价,并且不需要额外的训练。
|
||||
|
||||
本指南向您展示如何在 Diffusers 中使用支持的缓存方法。
|
||||
|
||||
## 金字塔注意力广播
|
||||
|
||||
[金字塔注意力广播 (PAB)](https://huggingface.co/papers/2408.12588) 基于这样一种观察:在生成过程的连续时间步之间,注意力输出差异不大。注意力差异在交叉注意力层中最小,并且通常在一个较长的时间步范围内被缓存。其次是时间注意力和空间注意力层。
|
||||
|
||||
> [!TIP]
|
||||
> 并非所有视频模型都有三种类型的注意力(交叉、时间和空间)!
|
||||
|
||||
PAB 可以与其他技术(如序列并行性和无分类器引导并行性(数据并行性))结合,实现近乎实时的视频生成。
|
||||
|
||||
设置并传递一个 [`PyramidAttentionBroadcastConfig`] 到管道的变换器以启用它。`spatial_attention_block_skip_range` 控制跳过空间注意力块中注意力计算的频率,`spatial_attention_timestep_skip_range` 是要跳过的时间步范围。注意选择一个合适的范围,因为较小的间隔可能导致推理速度变慢,而较大的间隔可能导致生成质量降低。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig
|
||||
|
||||
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
|
||||
pipeline.to("cuda")
|
||||
|
||||
config = PyramidAttentionBroadcastConfig(
|
||||
spatial_attention_block_skip_range=2,
|
||||
spatial_attention_timestep_skip_range=(100, 800),
|
||||
current_timestep_callback=lambda: pipe.current_timestep,
|
||||
)
|
||||
pipeline.transformer.enable_cache(config)
|
||||
```
|
||||
|
||||
## FasterCache
|
||||
|
||||
[FasterCache](https://huggingface.co/papers/2410.19355) 缓存并重用注意力特征,类似于 [PAB](#pyramid-attention-broadcast),因为每个连续时间步的输出差异很小。
|
||||
|
||||
此方法在使用无分类器引导进行采样时(在大多数基础模型中常见),也可能选择跳过无条件分支预测,并且
|
||||
如果连续时间步之间的预测潜在输出存在显著冗余,则从条件分支预测中估计它。
|
||||
|
||||
设置并将 [`FasterCacheConfig`] 传递给管道的 transformer 以启用它。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import CogVideoXPipeline, FasterCacheConfig
|
||||
|
||||
pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
|
||||
pipeline.to("cuda")
|
||||
|
||||
config = FasterCacheConfig(
|
||||
spatial_attention_block_skip_range=2,
|
||||
spatial_attention_timestep_skip_range=(-1, 681),
|
||||
current_timestep_callback=lambda: pipe.current_timestep,
|
||||
attention_weight_callback=lambda _: 0.3,
|
||||
unconditional_batch_skip_range=5,
|
||||
unconditional_batch_timestep_skip_range=(-1, 781),
|
||||
tensor_format="BFCHW",
|
||||
)
|
||||
pipeline.transformer.enable_cache(config)
|
||||
```
|
||||
@@ -1,163 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非符合许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证的具体语言,请参阅许可证中的权限和限制。
|
||||
-->
|
||||
|
||||
# 如何使用 Core ML 运行 Stable Diffusion
|
||||
|
||||
[Core ML](https://developer.apple.com/documentation/coreml) 是 Apple 框架支持的模型格式和机器学习库。如果您有兴趣在 macOS 或 iOS/iPadOS 应用中运行 Stable Diffusion 模型,本指南将展示如何将现有的 PyTorch 检查点转换为 Core ML 格式,并使用 Python 或 Swift 进行推理。
|
||||
|
||||
Core ML 模型可以利用 Apple 设备中所有可用的计算引擎:CPU、GPU 和 Apple Neural Engine(或 ANE,一种在 Apple Silicon Mac 和现代 iPhone/iPad 中可用的张量优化加速器)。根据模型及其运行的设备,Core ML 还可以混合和匹配计算引擎,例如,模型的某些部分可能在 CPU 上运行,而其他部分在 GPU 上运行。
|
||||
|
||||
<Tip>
|
||||
|
||||
您还可以使用 PyTorch 内置的 `mps` 加速器在 Apple Silicon Mac 上运行 `diffusers` Python 代码库。这种方法在 [mps 指南](mps) 中有详细解释,但它与原生应用不兼容。
|
||||
|
||||
</Tip>
|
||||
|
||||
## Stable Diffusion Core ML 检查点
|
||||
|
||||
Stable Diffusion 权重(或检查点)以 PyTorch 格式存储,因此在使用它们之前,需要将它们转换为 Core ML 格式。
|
||||
|
||||
幸运的是,Apple 工程师基于 `diffusers` 开发了 [一个转换工具](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml),用于将 PyTorch 检查点转换为 Core ML。
|
||||
|
||||
但在转换模型之前,花点时间探索 Hugging Face Hub——很可能您感兴趣的模型已经以 Core ML 格式提供:
|
||||
|
||||
- [Apple](https://huggingface.co/apple) 组织包括 Stable Diffusion 版本 1.4、1.5、2.0 基础和 2.1 基础
|
||||
- [coreml community](https://huggingface.co/coreml-community) 包括自定义微调模型
|
||||
- 使用此 [过滤器](https://huggingface.co/models?pipeline_tag=text-to-image&library=coreml&p=2&sort=likes) 返回所有可用的 Core ML 检查点
|
||||
|
||||
如果您找不到感兴趣的模型,我们建议您遵循 Apple 的 [Converting Models to Core ML](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml) 说明。
|
||||
|
||||
## 选择要使用的 Core ML 变体
|
||||
|
||||
Stable Diffusion 模型可以转换为不同的 Core ML 变体,用于不同目的:
|
||||
|
||||
- 注意力类型
|
||||
使用了n个块。注意力操作用于“关注”图像表示中不同区域之间的关系,并理解图像和文本表示如何相关。注意力的计算和内存消耗很大,因此存在不同的实现方式,以适应不同设备的硬件特性。对于Core ML Stable Diffusion模型,有两种注意力变体:
|
||||
* `split_einsum`([由Apple引入](https://machinelearning.apple.com/research/neural-engine-transformers))针对ANE设备进行了优化,这些设备在现代iPhone、iPad和M系列计算机中可用。
|
||||
* “原始”注意力(在`diffusers`中使用的基础实现)仅与CPU/GPU兼容,不与ANE兼容。在CPU + GPU上使用`original`注意力运行模型可能比ANE*更快*。请参阅[此性能基准](https://huggingface.co/blog/fast-mac-diffusers#performance-benchmarks)以及社区提供的[一些额外测量](https://github.com/huggingface/swift-coreml-diffusers/issues/31)以获取更多细节。
|
||||
|
||||
- 支持的推理框架。
|
||||
* `packages`适用于Python推理。这可用于在尝试将转换后的Core ML模型集成到原生应用程序之前进行测试,或者如果您想探索Core ML性能但不需要支持原生应用程序。例如,具有Web UI的应用程序完全可以使用Python Core ML后端。
|
||||
* `compiled`模型是Swift代码所必需的。Hub中的`compiled`模型将大型UNet模型权重分成多个文件,以兼容iOS和iPadOS设备。这对应于[`--chunk-unet`转换选项](https://github.com/apple/ml-stable-diffusion#-converting-models-to-core-ml)。如果您想支持原生应用程序,则需要选择`compiled`变体。
|
||||
|
||||
官方的Core ML Stable Diffusion[模型](https://huggingface.co/apple/coreml-stable-diffusion-v1-4/tree/main)包括这些变体,但社区的可能有所不同:
|
||||
|
||||
```
|
||||
coreml-stable-diffusion-v1-4
|
||||
├── README.md
|
||||
├── original
|
||||
│ ├── compiled
|
||||
│ └── packages
|
||||
└── split_einsum
|
||||
├── compiled
|
||||
└── packages
|
||||
```
|
||||
|
||||
您可以下载并使用所需的变体,如下所示。
|
||||
|
||||
## Python中的Core ML推理
|
||||
|
||||
安装以下库以在Python中运行Core ML推理:
|
||||
|
||||
```bash
|
||||
pip install huggingface_hub
|
||||
pip install git+https://github.com/apple/ml-stable-diffusion
|
||||
```
|
||||
|
||||
### 下载模型检查点
|
||||
|
||||
要在Python中运行推理,请使用存储在`packages`文件夹中的版本之一,因为`compiled`版本仅与Swift兼容。您可以选择使用`original`或`split_einsum`注意力。
|
||||
|
||||
这是您如何从Hub下载`original`注意力变体到一个名为`models`的目录:
|
||||
|
||||
```Python
|
||||
from huggingface_hub import snapshot_download
|
||||
from pathlib import Path
|
||||
|
||||
repo_id = "apple/coreml-stable-diffusion-v1-4"
|
||||
variant = "original/packages"
|
||||
|
||||
mo
|
||||
del_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
|
||||
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
|
||||
print(f"Model downloaded at {model_path}")
|
||||
```
|
||||
|
||||
### 推理[[python-inference]]
|
||||
|
||||
下载模型快照后,您可以使用 Apple 的 Python 脚本来测试它。
|
||||
|
||||
```shell
|
||||
python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i ./models/coreml-stable-diffusion-v1-4_original_packages/original/packages -o </path/to/output/image> --compute-unit CPU_AND_GPU --seed 93
|
||||
```
|
||||
|
||||
使用 `-i` 标志将下载的检查点路径传递给脚本。`--compute-unit` 表示您希望允许用于推理的硬件。它必须是以下选项之一:`ALL`、`CPU_AND_GPU`、`CPU_ONLY`、`CPU_AND_NE`。您也可以提供可选的输出路径和用于可重现性的种子。
|
||||
|
||||
推理脚本假设您使用的是 Stable Diffusion 模型的原始版本,`CompVis/stable-diffusion-v1-4`。如果您使用另一个模型,您*必须*在推理命令行中使用 `--model-version` 选项指定其 Hub ID。这适用于已支持的模型以及您自己训练或微调的自定义模型。
|
||||
|
||||
例如,如果您想使用 [`stable-diffusion-v1-5/stable-diffusion-v1-5`](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5):
|
||||
|
||||
```shell
|
||||
python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" --compute-unit ALL -o output --seed 93 -i models/coreml-stable-diffusion-v1-5_original_packages --model-version stable-diffusion-v1-5/stable-diffusion-v1-5
|
||||
```
|
||||
|
||||
## Core ML 在 Swift 中的推理
|
||||
|
||||
在 Swift 中运行推理比在 Python 中稍快,因为模型已经以 `mlmodelc` 格式编译。这在应用启动时加载模型时很明显,但如果在之后运行多次生成,则不应明显。
|
||||
|
||||
### 下载
|
||||
|
||||
要在您的 Mac 上运行 Swift 推理,您需要一个 `compiled` 检查点版本。我们建议您使用类似于先前示例的 Python 代码在本地下载它们,但使用 `compiled` 变体之一:
|
||||
|
||||
```Python
|
||||
from huggingface_hub import snapshot_download
|
||||
from pathlib import Path
|
||||
|
||||
repo_id = "apple/coreml-stable-diffusion-v1-4"
|
||||
variant = "original/compiled"
|
||||
|
||||
model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
|
||||
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
|
||||
print(f"Model downloaded at {model_path}")
|
||||
```
|
||||
|
||||
### 推理[[swift-inference]]
|
||||
|
||||
要运行推理,请克隆 Apple 的仓库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/apple/ml-stable-diffusion
|
||||
cd ml-stable-diffusion
|
||||
```
|
||||
|
||||
然后使用 Apple 的命令行工具,[Swift Package Manager](https://www.swift.org/package-manager/#):
|
||||
|
||||
```bash
|
||||
swift run StableDiffusionSample --resource-path models/coreml-stable-diffusion-v1-4_original_compiled --compute-units all "a photo of an astronaut riding a horse on mars"
|
||||
```
|
||||
|
||||
您必须在 `--resource-path` 中指定上一步下载的检查点之一,请确保它包含扩展名为 `.mlmodelc` 的已编译 Core ML 包。`--compute-units` 必须是以下值之一:`all`、`cpuOnly`、`cpuAndGPU`、`cpuAndNeuralEngine`。
|
||||
|
||||
有关更多详细信息,请参考 [Apple 仓库中的说明](https://github.com/apple/ml-stable-diffusion)。
|
||||
|
||||
## 支持的 Diffusers 功能
|
||||
|
||||
Core ML 模型和推理代码不支持 🧨 Diffusers 的许多功能、选项和灵活性。以下是一些需要注意的限制:
|
||||
|
||||
- Core ML 模型仅适用于推理。它们不能用于训练或微调。
|
||||
- 只有两个调度器已移植到 Swift:Stable Diffusion 使用的默认调度器和我们从 `diffusers` 实现移植到 Swift 的 `DPMSolverMultistepScheduler`。我们推荐您使用 `DPMSolverMultistepScheduler`,因为它在约一半的步骤中产生相同的质量。
|
||||
- 负面提示、无分类器引导尺度和图像到图像任务在推理代码中可用。高级功能如深度引导、ControlNet 和潜在上采样器尚不可用。
|
||||
|
||||
Apple 的 [转换和推理仓库](https://github.com/apple/ml-stable-diffusion) 和我们自己的 [swift-coreml-diffusers](https://github.com/huggingface/swift-coreml-diffusers) 仓库旨在作为技术演示,以帮助其他开发者在此基础上构建。
|
||||
|
||||
如果您对任何缺失功能有强烈需求,请随时提交功能请求或更好的是,贡献一个 PR 🙂。
|
||||
|
||||
## 原生 Diffusers Swift 应用
|
||||
|
||||
一个简单的方法来在您自己的 Apple 硬件上运行 Stable Diffusion 是使用 [我们的开源 Swift 仓库](https://github.com/huggingface/swift-coreml-diffusers),它基于 `diffusers` 和 Apple 的转换和推理仓库。您可以研究代码,使用 [Xcode](https://developer.apple.com/xcode/) 编译它,并根据您的需求进行适配。为了方便,[App Store 中还有一个独立 Mac 应用](https://apps.apple.com/app/diffusers/id1666309574),因此您无需处理代码或 IDE 即可使用它。如果您是开发者,并已确定 Core ML 是构建您的 Stable Diffusion 应用的最佳解决方案,那么您可以使用本指南的其余部分来开始您的项目。我们迫不及待想看看您会构建什么 🙂。
|
||||
@@ -1,59 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,否则根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证的具体语言,请参阅许可证中的权限和限制。
|
||||
-->
|
||||
|
||||
# DeepCache
|
||||
[DeepCache](https://huggingface.co/papers/2312.00858) 通过策略性地缓存和重用高级特征,同时利用 U-Net 架构高效更新低级特征,来加速 [`StableDiffusionPipeline`] 和 [`StableDiffusionXLPipeline`]。
|
||||
|
||||
首先安装 [DeepCache](https://github.com/horseee/DeepCache):
|
||||
```bash
|
||||
pip install DeepCache
|
||||
```
|
||||
|
||||
然后加载并启用 [`DeepCacheSDHelper`](https://github.com/horseee/DeepCache#usage):
|
||||
|
||||
```diff
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
pipe = StableDiffusionPipeline.from_pretrained('stable-diffusion-v1-5/stable-diffusion-v1-5', torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
+ from DeepCache import DeepCacheSDHelper
|
||||
+ helper = DeepCacheSDHelper(pipe=pipe)
|
||||
+ helper.set_params(
|
||||
+ cache_interval=3,
|
||||
+ cache_branch_id=0,
|
||||
+ )
|
||||
+ helper.enable()
|
||||
|
||||
image = pipe("a photo of an astronaut on a moon").images[0]
|
||||
```
|
||||
|
||||
`set_params` 方法接受两个参数:`cache_interval` 和 `cache_branch_id`。`cache_interval` 表示特征缓存的频率,指定为每次缓存操作之间的步数。`cache_branch_id` 标识网络的哪个分支(从最浅层到最深层排序)负责执行缓存过程。
|
||||
选择较低的 `cache_branch_id` 或较大的 `cache_interval` 可以加快推理速度,但会降低图像质量(这些超参数的消融实验可以在[论文](https://huggingface.co/papers/2312.00858)中找到)。一旦设置了这些参数,使用 `enable` 或 `disable` 方法来激活或停用 `DeepCacheSDHelper`。
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://github.com/horseee/Diffusion_DeepCache/raw/master/static/images/example.png">
|
||||
</div>
|
||||
|
||||
您可以在 [WandB 报告](https://wandb.ai/horseee/DeepCache/runs/jwlsqqgt?workspace=user-horseee) 中找到更多生成的样本(原始管道 vs DeepCache)和相应的推理延迟。提示是从 [MS-COCO 2017](https://cocodataset.org/#home) 数据集中随机选择的。
|
||||
|
||||
## 基准测试
|
||||
|
||||
我们在 NVIDIA RTX A5000 上测试了 DeepCache 使用 50 个推理步骤加速 [Stable Diffusion v2.1](https://huggingface.co/stabilityai/stable-diffusion-2-1) 的速度,使用不同的配置,包括分辨率、批处理大小、缓存间隔(I)和缓存分支(B)。
|
||||
|
||||
| **分辨率** | **批次大小** | **原始** | **DeepCache(I=3, B=0)** | **DeepCache(I=5, B=0)** | **DeepCache(I=5, B=1)** |
|
||||
|----------------|----------------|--------------|-------------------------|-------------------------|-------------------------|
|
||||
| 512| 8| 15.96| 6.88(2.32倍)| 5.03(3.18倍)| 7.27(2.20x)|
|
||||
| | 4| 8.39| 3.60(2.33倍)| 2.62(3.21倍)| 3.75(2.24x)|
|
||||
| | 1| 2.61| 1.12(2.33倍)| 0.81(3.24倍)| 1.11(2.35x)|
|
||||
| 768| 8| 43.58| 18.99(2.29倍)| 13.96(3.12倍)| 21.27(2.05x)|
|
||||
| | 4| 22.24| 9.67(2.30倍)| 7.10(3.13倍)| 10.74(2.07x)|
|
||||
| | 1| 6.33| 2.72(2.33倍)| 1.97(3.21倍)| 2.98(2.12x)|
|
||||
| 1024| 8| 101.95| 45.57(2.24倍)| 33.72(3.02倍)| 53.00(1.92x)|
|
||||
| | 4| 49.25| 21.86(2.25倍)| 16.19(3.04倍)| 25.78(1.91x)|
|
||||
| | 1| 13.83| 6.07(2.28倍)| 4.43(3.12倍)| 7.15(1.93x)|
|
||||
@@ -1,28 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。有关许可证管理权限和限制的具体语言,请参阅许可证。
|
||||
-->
|
||||
|
||||
# Intel Gaudi
|
||||
|
||||
Intel Gaudi AI 加速器系列包括 [Intel Gaudi 1](https://habana.ai/products/gaudi/)、[Intel Gaudi 2](https://habana.ai/products/gaudi2/) 和 [Intel Gaudi 3](https://habana.ai/products/gaudi3/)。每台服务器配备 8 个设备,称为 Habana 处理单元 (HPU),在 Gaudi 3 上提供 128GB 内存,在 Gaudi 2 上提供 96GB 内存,在第一代 Gaudi 上提供 32GB 内存。有关底层硬件架构的更多详细信息,请查看 [Gaudi 架构](https://docs.habana.ai/en/latest/Gaudi_Overview/Gaudi_Architecture.html) 概述。
|
||||
|
||||
Diffusers 管道可以利用 HPU 加速,即使管道尚未添加到 [Optimum for Intel Gaudi](https://huggingface.co/docs/optimum/main/en/habana/index),也可以通过 [GPU 迁移工具包](https://docs.habana.ai/en/latest/PyTorch/PyTorch_Model_Porting/GPU_Migration_Toolkit/GPU_Migration_Toolkit.html) 实现。
|
||||
|
||||
在您的管道上调用 `.to("hpu")` 以将其移动到 HPU 设备,如下所示为 Flux 示例:
|
||||
```py
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
|
||||
pipeline.to("hpu")
|
||||
|
||||
image = pipeline("一张松鼠在毕加索风格中的图像").images[0]
|
||||
```
|
||||
|
||||
> [!TIP]
|
||||
> 对于 Gaudi 优化的扩散管道实现,我们推荐使用 [Optimum for Intel Gaudi](https://huggingface.co/docs/optimum/main/en/habana/index)。
|
||||
@@ -1,581 +0,0 @@
|
||||
<!--版权所有 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本(“许可证”)授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按“原样”分发,无任何明示或暗示的担保或条件。有关许可证的特定语言管理权限和限制,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 减少内存使用
|
||||
|
||||
现代diffusion models,如 [Flux](../api/pipelines/flux) 和 [Wan](../api/pipelines/wan),拥有数十亿参数,在您的硬件上进行推理时会占用大量内存。这是一个挑战,因为常见的 GPU 通常没有足够的内存。为了克服内存限制,您可以使用多个 GPU(如果可用)、将一些管道组件卸载到 CPU 等。
|
||||
|
||||
本指南将展示如何减少内存使用。
|
||||
|
||||
> [!TIP]
|
||||
> 请记住,这些技术可能需要根据模型进行调整。例如,基于 transformer 的扩散模型可能不会像基于 UNet 的模型那样从这些内存优化中同等受益。
|
||||
|
||||
## 多个 GPU
|
||||
|
||||
如果您有多个 GPU 的访问权限,有几种选项可以高效地在硬件上加载和分发大型模型。这些功能由 [Accelerate](https://huggingface.co/docs/accelerate/index) 库支持,因此请确保先安装它。
|
||||
|
||||
```bash
|
||||
pip install -U accelerate
|
||||
```
|
||||
|
||||
### 分片检查点
|
||||
|
||||
将大型检查点加载到多个分片中很有用,因为分片会逐个加载。这保持了低内存使用,只需要足够的内存来容纳模型大小和最大分片大小。我们建议当 fp32 检查点大于 5GB 时进行分片。默认分片大小为 5GB。
|
||||
|
||||
在 [`~DiffusionPipeline.save_pretrained`] 中使用 `max_shard_size` 参数对检查点进行分片。
|
||||
|
||||
```py
|
||||
from diffusers import AutoModel
|
||||
|
||||
unet = AutoModel.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
|
||||
)
|
||||
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")
|
||||
```
|
||||
|
||||
现在您可以使用分片检查点,而不是常规检查点,以节省内存。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel, StableDiffusionXLPipeline
|
||||
|
||||
unet = AutoModel.from_pretrained(
|
||||
"username/sdxl-unet-sharded", torch_dtype=torch.float16
|
||||
)
|
||||
pipeline = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
unet=unet,
|
||||
torch_dtype=torch.float16
|
||||
).to("cuda")
|
||||
```
|
||||
|
||||
### 设备放置
|
||||
|
||||
> [!WARNING]
|
||||
> 设备放置是一个实验性功能,API 可能会更改。目前仅支持 `balanced` 策略。我们计划在未来支持额外的映射策略。
|
||||
|
||||
`device_map` 参数控制管道或模型中的组件如何
|
||||
单个模型中的层分布在多个设备上。
|
||||
|
||||
<hfoptions id="device-map">
|
||||
<hfoption id="pipeline level">
|
||||
|
||||
`balanced` 设备放置策略将管道均匀分割到所有可用设备上。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel, StableDiffusionXLPipeline
|
||||
|
||||
pipeline = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
device_map="balanced"
|
||||
)
|
||||
```
|
||||
|
||||
您可以使用 `hf_device_map` 检查管道的设备映射。
|
||||
|
||||
```py
|
||||
print(pipeline.hf_device_map)
|
||||
{'unet': 1, 'vae': 1, 'safety_checker': 0, 'text_encoder': 0}
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="model level">
|
||||
|
||||
`device_map` 对于加载大型模型非常有用,例如具有 125 亿参数的 Flux diffusion transformer。将其设置为 `"auto"` 可以自动将模型首先分布到最快的设备上,然后再移动到较慢的设备。有关更多详细信息,请参阅 [模型分片](../training/distributed_inference#model-sharding) 文档。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel
|
||||
|
||||
transformer = AutoModel.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
subfolder="transformer",
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
```
|
||||
|
||||
您可以使用 `hf_device_map` 检查模型的设备映射。
|
||||
|
||||
```py
|
||||
print(transformer.hf_device_map)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
当设计您自己的 `device_map` 时,它应该是一个字典,包含模型的特定模块名称或层以及设备标识符(整数表示 GPU,`cpu` 表示 CPU,`disk` 表示磁盘)。
|
||||
|
||||
在模型上调用 `hf_device_map` 以查看模型层如何分布,然后设计您自己的映射。
|
||||
|
||||
```py
|
||||
print(transformer.hf_device_map)
|
||||
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}
|
||||
```
|
||||
|
||||
例如,下面的 `device_map` 将 `single_transformer_blocks.10` 到 `single_transformer_blocks.20` 放置在第二个 GPU(`1`)上。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel
|
||||
|
||||
device_map = {
|
||||
'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
|
||||
}
|
||||
|
||||
transformer = AutoModel.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
subfolder="transformer",
|
||||
device_map=device_map,
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
```
|
||||
|
||||
传递一个字典,将最大内存使用量映射到每个设备以强制执行限制。如果设备不在 `max_memory` 中,它将被忽略,管道组件不会分发到该设备。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel, StableDiffusionXLPipeline
|
||||
|
||||
max_memory = {0:"1GB", 1:"1GB"}
|
||||
pipeline = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
device_map="balanced",
|
||||
max_memory=max_memory
|
||||
)
|
||||
```
|
||||
|
||||
Diffusers 默认使用所有设备的最大内存,但如果它们无法适应 GPU,则需要使用单个 GPU 并通过以下方法卸载到 CPU。
|
||||
|
||||
- [`~DiffusionPipeline.enable_model_cpu_offload`] 仅适用于单个 GPU,但非常大的模型可能无法适应它
|
||||
- 使用 [`~DiffusionPipeline.enable_sequential_cpu_offload`] 可能有效,但它极其缓慢,并且仅限于单个 GPU。
|
||||
|
||||
使用 [`~DiffusionPipeline.reset_device_map`] 方法来重置 `device_map`。如果您想在已进行设备映射的管道上使用方法如 `.to()`、[`~DiffusionPipeline.enable_sequential_cpu_offload`] 和 [`~DiffusionPipeline.enable_model_cpu_offload`],这是必要的。
|
||||
|
||||
```py
|
||||
pipeline.reset_device_map()
|
||||
```
|
||||
|
||||
## VAE 切片
|
||||
|
||||
VAE 切片通过将大批次输入拆分为单个数据批次并分别处理它们来节省内存。这种方法在同时生成多个图像时效果最佳。
|
||||
|
||||
例如,如果您同时生成 4 个图像,解码会将峰值激活内存增加 4 倍。VAE 切片通过一次只解码 1 个图像而不是所有 4 个图像来减少这种情况。
|
||||
|
||||
调用 [`~StableDiffusionPipeline.enable_vae_slicing`] 来启用切片 VAE。您可以预期在解码多图像批次时性能会有小幅提升,而在单图像批次时没有性能影响。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoModel, StableDiffusionXLPipeline
|
||||
|
||||
pipeline = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
pipeline.enable_vae_slicing()
|
||||
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
|
||||
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
```
|
||||
|
||||
> [!WARNING]
|
||||
> [`AutoencoderKLWan`] 和 [`AsymmetricAutoencoderKL`] 类不支持切片。
|
||||
|
||||
## VAE 平铺
|
||||
|
||||
VAE 平铺通过将图像划分为较小的重叠图块而不是一次性处理整个图像来节省内存。这也减少了峰值内存使用量,因为 GPU 一次只处理一个图块。
|
||||
|
||||
调用 [`~StableDiffusionPipeline.enable_vae_tiling`] 来启用 VAE 平铺。生成的图像可能因图块到图块的色调变化而有所不同,因为它们被单独解码,但图块之间不应有明显的接缝。对于低于预设(但可配置)限制的分辨率,平铺被禁用。例如,对于 [`StableDiffusionPipeline`] 中的 VAE,此限制为 512x512。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoPipelineForImage2Image
|
||||
from diffusers.utils import load_image
|
||||
|
||||
pipeline = AutoPipelineForImage2Image.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
||||
).to("cuda")
|
||||
pipeline.enable_vae_tiling()
|
||||
|
||||
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
|
||||
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
|
||||
pipeline(prompt, image=init_image, strength=0.5).images[0]
|
||||
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
```
|
||||
|
||||
> [!WARNING]
|
||||
> [`AutoencoderKLWan`] 和 [`AsymmetricAutoencoderKL`] 不支持平铺。
|
||||
|
||||
## 卸载
|
||||
|
||||
卸载策略将非当前活动层移动
|
||||
将模型移动到 CPU 以避免增加 GPU 内存。这些策略可以与量化和 torch.compile 结合使用,以平衡推理速度和内存使用。
|
||||
|
||||
有关更多详细信息,请参考 [编译和卸载量化模型](./speed-memory-optims) 指南。
|
||||
|
||||
### CPU 卸载
|
||||
|
||||
CPU 卸载选择性地将权重从 GPU 移动到 CPU。当需要某个组件时,它被传输到 GPU;当不需要时,它被移动到 CPU。此方法作用于子模块而非整个模型。它通过避免将整个模型存储在 GPU 上来节省内存。
|
||||
|
||||
CPU 卸载显著减少内存使用,但由于子模块在设备之间多次来回传递,它也非常慢。由于速度极慢,它通常不实用。
|
||||
|
||||
> [!WARNING]
|
||||
> 在调用 [`~DiffusionPipeline.enable_sequential_cpu_offload`] 之前,不要将管道移动到 CUDA,否则节省的内存非常有限(更多细节请参考此 [issue](https://github.com/huggingface/diffusers/issues/1934))。这是一个状态操作,会在模型上安装钩子。
|
||||
|
||||
调用 [`~DiffusionPipeline.enable_sequential_cpu_offload`] 以在管道上启用它。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
|
||||
)
|
||||
pipeline.enable_sequential_cpu_offload()
|
||||
|
||||
pipeline(
|
||||
prompt="An astronaut riding a horse on Mars",
|
||||
guidance_scale=0.,
|
||||
height=768,
|
||||
width=1360,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=256,
|
||||
).images[0]
|
||||
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
```
|
||||
|
||||
### 模型卸载
|
||||
|
||||
模型卸载将整个模型移动到 GPU,而不是选择性地移动某些层或模型组件。一个主要管道模型,通常是文本编码器、UNet 和 VAE,被放置在 GPU 上,而其他组件保持在 CPU 上。像 UNet 这样运行多次的组件会一直留在 GPU 上,直到完全完成且不再需要。这消除了 [CPU 卸载](#cpu-offloading) 的通信开销,使模型卸载成为一个更快的替代方案。权衡是内存节省不会那么大。
|
||||
|
||||
> [!WARNING]
|
||||
> 请注意,如果在安装钩子后模型在管道外部被重用(更多细节请参考 [移除钩子](https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.hooks.remove_hook_from_module)),您需要按预期顺序运行整个管道和模型以正确卸载它们。这是一个状态操作,会在模型上安装钩子。
|
||||
|
||||
调用 [`~DiffusionPipeline.enable_model_cpu_offload`] 以在管道上启用它。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
|
||||
)
|
||||
pipeline.enable_model_cpu_offload()
|
||||
|
||||
pipeline(
|
||||
prompt="An astronaut riding a horse on Mars",
|
||||
guidance_scale=0.,
|
||||
height=768,
|
||||
width=1360,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=256,
|
||||
).images[0]
|
||||
print(f"最大内存保留: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
```
|
||||
|
||||
[`~DiffusionPipeline.enable_model_cpu_offload`] 在您单独使用 [`~StableDiffusionXLPipeline.encode_prompt`] 方法生成文本编码器隐藏状态时也有帮助。
|
||||
|
||||
### 组卸载
|
||||
|
||||
组卸载将内部层组([torch.nn.ModuleList](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) 或 [torch.nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html))移动到 CPU。它比[模型卸载](#model-offloading)使用更少的内存,并且比[CPU 卸载](#cpu-offloading)更快,因为它减少了通信开销。
|
||||
|
||||
> [!WARNING]
|
||||
> 如果前向实现包含权重相关的输入设备转换,组卸载可能不适用于所有模型,因为它可能与组卸载的设备转换机制冲突。
|
||||
|
||||
调用 [`~ModelMixin.enable_group_offload`] 为继承自 [`ModelMixin`] 的标准 Diffusers 模型组件启用它。对于不继承自 [`ModelMixin`] 的其他模型组件,例如通用 [torch.nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html),使用 [`~hooks.apply_group_offloading`] 代替。
|
||||
|
||||
`offload_type` 参数可以设置为 `block_level` 或 `leaf_level`。
|
||||
|
||||
- `block_level` 基于 `num_blocks_per_group` 参数卸载层组。例如,如果 `num_blocks_per_group=2` 在一个有 40 层的模型上,每次加载和卸载 2 层(总共 20 次加载/卸载)。这大大减少了内存需求。
|
||||
- `leaf_level` 在最低级别卸载单个层,等同于[CPU 卸载](#cpu-offloading)。但如果您使用流而不放弃推理速度,它可以更快。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import CogVideoXPipeline
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
onload_device = torch.device("cuda")
|
||||
offload_device = torch.device("cpu")
|
||||
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
|
||||
|
||||
# 对 Diffusers 模型实现使用 enable_group_offload 方法
|
||||
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level")
|
||||
pipeline.vae.enable_group_offload(onload_device=onload_device, offload_type="leaf_level")
|
||||
|
||||
# 对其他模型组件使用 apply_group_offloading 方法
|
||||
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
|
||||
|
||||
prompt = (
|
||||
"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
|
||||
"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
|
||||
"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
|
||||
"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
|
||||
"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
|
||||
"atmosphere of this unique musical performance."
|
||||
)
|
||||
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
|
||||
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
export_to_video(video, "output.mp4", fps=8)
|
||||
```
|
||||
|
||||
#### CUDA 流
|
||||
`use_stream` 参数可以激活支持异步数据传输流的 CUDA 设备,以减少整体执行时间,与 [CPU 卸载](#cpu-offloading) 相比。它通过使用层预取重叠数据传输和计算。下一个要执行的层在当前层仍在执行时加载到 GPU 上。这会显著增加 CPU 内存,因此请确保您有模型大小的 2 倍内存。
|
||||
|
||||
设置 `record_stream=True` 以获得更多速度提升,代价是内存使用量略有增加。请参阅 [torch.Tensor.record_stream](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) 文档了解更多信息。
|
||||
|
||||
> [!TIP]
|
||||
> 当 `use_stream=True` 在启用平铺的 VAEs 上时,确保在推理前进行虚拟前向传递(可以使用虚拟输入),以避免设备不匹配错误。这可能不适用于所有实现,因此如果遇到任何问题,请随时提出问题。
|
||||
|
||||
如果您在使用启用 `use_stream` 的 `block_level` 组卸载,`num_blocks_per_group` 参数应设置为 `1`,否则会引发警告。
|
||||
|
||||
```py
|
||||
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)
|
||||
```
|
||||
|
||||
`low_cpu_mem_usage` 参数可以设置为 `True`,以在使用流进行组卸载时减少 CPU 内存使用。它最适合 `leaf_level` 卸载和 CPU 内存瓶颈的情况。通过动态创建固定张量而不是预先固定它们来节省内存。然而,这可能会增加整体执行时间。
|
||||
|
||||
#### 卸载到磁盘
|
||||
组卸载可能会消耗大量系统内存,具体取决于模型大小。在内存有限的系统上,尝试将组卸载到磁盘作为辅助内存。
|
||||
|
||||
在 [`~ModelMixin.enable_group_offload`] 或 [`~hooks.apply_group_offloading`] 中设置 `offload_to_disk_path` 参数,将模型卸载到磁盘。
|
||||
|
||||
```py
|
||||
pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")
|
||||
|
||||
apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")
|
||||
```
|
||||
|
||||
参考这些[两个](https://github.com/huggingface/diffusers/pull/11682#issue-3129365363)[表格](https://github.com/huggingface/diffusers/pull/11682#issuecomment-2955715126)来比较速度和内存的权衡。
|
||||
|
||||
## 分层类型转换
|
||||
|
||||
> [!TIP]
|
||||
> 将分层类型转换与[组卸载](#group-offloading)结合使用,以获得更多内存节省。
|
||||
|
||||
分层类型转换将权重存储在较小的数据格式中(例如 `torch.float8_e4m3fn` 和 `torch.float8_e5m2`),以使用更少的内存,并在计算时将那些权重上转换为更高精度如 `torch.float16` 或 `torch.bfloat16`。某些层(归一化和调制相关权重)被跳过,因为将它们存储在 fp8 中可能会降低生成质量。
|
||||
|
||||
> [!WARNING]
|
||||
> 如果前向实现包含权重的内部类型转换,分层类型转换可能不适用于所有模型。当前的分层类型转换实现假设前向传递独立于权重精度,并且输入数据类型始终在 `compute_dtype` 中指定(请参见[这里](https://github.com/huggingface/transformers/blob/7f5077e53682ca855afc826162b204ebf809f1f9/src/transformers/models/t5/modeling_t5.py#L294-L299)以获取不兼容的实现)。
|
||||
>
|
||||
> 分层类型转换也可能在使用[PEFT](https://huggingface.co/docs/peft/index)层的自定义建模实现上失败。有一些检查可用,但它们没有经过广泛测试或保证在所有情况下都能工作。
|
||||
|
||||
调用 [`~ModelMixin.enable_layerwise_casting`] 来设置存储和计算数据类型。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
transformer = CogVideoXTransformer3DModel.from_pretrained(
|
||||
"THUDM/CogVideoX-5b",
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)
|
||||
|
||||
pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
|
||||
transformer=transformer,
|
||||
torch_dtype=torch.bfloat16
|
||||
).to("cuda")
|
||||
prompt = (
|
||||
"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
|
||||
"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
|
||||
"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
|
||||
"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
|
||||
"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
|
||||
"atmosphere of this unique musical performance."
|
||||
)
|
||||
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
|
||||
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
|
||||
export_to_video(video, "output.mp4", fps=8)
|
||||
```
|
||||
|
||||
[`~hooks.apply_layerwise_casting`] 方法也可以在您需要更多控制和灵活性时使用。它可以通过在特定内部模块上调用它来部分应用于模型层。使用 `skip_modules_pattern` 或 `skip_modules_classes` 参数来指定要避免的模块,例如归一化和调制层。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import CogVideoXTransformer3DModel
|
||||
from diffusers.hooks import apply_layerwise_casting
|
||||
|
||||
transformer = CogVideoXTransformer3DModel.from_pretrained(
|
||||
"THUDM/CogVideoX-5b",
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
# 跳过归一化层
|
||||
apply_layerwise_casting(
|
||||
transformer,
|
||||
storage_dtype=torch.float8_e4m3fn,
|
||||
compute_dtype=torch.bfloat16,
|
||||
skip_modules_classes=["norm"],
|
||||
non_blocking=True,
|
||||
)
|
||||
```
|
||||
|
||||
## torch.channels_last
|
||||
|
||||
[torch.channels_last](https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html) 将张量的存储方式从 `(批次大小, 通道数, 高度, 宽度)` 翻转为 `(批次大小, 高度, 宽度, 通道数)`。这使张量与硬件如何顺序访问存储在内存中的张量对齐,并避免了在内存中跳转以访问像素值。
|
||||
|
||||
并非所有运算符当前都支持通道最后格式,并且可能导致性能更差,但仍然值得尝试。
|
||||
|
||||
```py
|
||||
print(pipeline.unet.conv_out.state_dict()["weight"].stride()) # (2880, 9, 3, 1)
|
||||
pipeline.unet.to(memory_format=torch.channels_last) # 原地操作
|
||||
print(
|
||||
pipeline.unet.conv_out.state_dict()["weight"].stride()
|
||||
) # (2880, 1, 960, 320) 第二个维度的跨度为1证明它有效
|
||||
```
|
||||
|
||||
## torch.jit.trace
|
||||
|
||||
[torch.jit.trace](https://pytorch.org/docs/stable/generated/torch.jit.trace.html) 记录模型在样本输入上执行的操作,并根据记录的执行路径创建一个新的、优化的模型表示。在跟踪过程中,模型被优化以减少来自Python和动态控制流的开销,并且操作被融合在一起以提高效率。返回的可执行文件或 [ScriptFunction](https://pytorch.org/docs/stable/generated/torch.jit.ScriptFunction.html) 可以被编译。
|
||||
|
||||
```py
|
||||
import time
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
import functools
|
||||
|
||||
# torch 禁用梯度
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
# 设置变量
|
||||
n_experiments = 2
|
||||
unet_runs_per_experiment = 50
|
||||
|
||||
# 加载样本输入
|
||||
def generate_inputs():
|
||||
sample = torch.randn((2, 4, 64, 64), device="cuda", dtype=torch.float16)
|
||||
timestep = torch.rand(1, device="cuda", dtype=torch.float16) * 999
|
||||
encoder_hidden_states = torch.randn((2, 77, 768), device="cuda", dtype=torch.float16)
|
||||
return sample, timestep, encoder_hidden_states
|
||||
|
||||
|
||||
pipeline = StableDiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
torch_dtype=torch.float16,
|
||||
use_safetensors=True,
|
||||
).to("cuda")
|
||||
unet = pipeline.unet
|
||||
unet.eval()
|
||||
unet.to(memory
|
||||
_format=torch.channels_last) # 使用 channels_last 内存格式
|
||||
unet.forward = functools.partial(unet.forward, return_dict=False) # 设置 return_dict=False 为默认
|
||||
|
||||
# 预热
|
||||
for _ in range(3):
|
||||
with torch.inference_mode():
|
||||
inputs = generate_inputs()
|
||||
orig_output = unet(*inputs)
|
||||
|
||||
# 追踪
|
||||
print("tracing..")
|
||||
unet_traced = torch.jit.trace(unet, inputs)
|
||||
unet_traced.eval()
|
||||
print("done tracing")
|
||||
|
||||
# 预热和优化图
|
||||
for _ in range(5):
|
||||
with torch.inference_mode():
|
||||
inputs = generate_inputs()
|
||||
orig_output = unet_traced(*inputs)
|
||||
|
||||
# 基准测试
|
||||
with torch.inference_mode():
|
||||
for _ in range(n_experiments):
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(unet_runs_per_experiment):
|
||||
orig_output = unet_traced(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
print(f"unet traced inference took {time.time() - start_time:.2f} seconds")
|
||||
for _ in range(n_experiments):
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(unet_runs_per_experiment):
|
||||
orig_output = unet(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
print(f"unet inference took {time.time() - start_time:.2f} seconds")
|
||||
|
||||
# 保存模型
|
||||
unet_traced.save("unet_traced.pt")
|
||||
```
|
||||
|
||||
替换管道的 UNet 为追踪版本。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class UNet2DConditionOutput:
|
||||
sample: torch.Tensor
|
||||
|
||||
pipeline = StableDiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
torch_dtype=torch.float16,
|
||||
use_safetensors=True,
|
||||
).to("cuda")
|
||||
|
||||
# 使用 jitted unet
|
||||
unet_traced = torch.jit.load("unet_traced.pt")
|
||||
|
||||
# del pipeline.unet
|
||||
class TracedUNet(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.in_channels = pipe.unet.config.in_channels
|
||||
self.device = pipe.unet.device
|
||||
|
||||
def forward(self, latent_model_input, t, encoder_hidden_states):
|
||||
sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
|
||||
return UNet2DConditionOutput(sample=sample)
|
||||
|
||||
pipeline.unet = TracedUNet()
|
||||
|
||||
with torch.inference_mode():
|
||||
image = pipe([prompt] * 1, num_inference_steps=50).images[0]
|
||||
```
|
||||
|
||||
## 内存高效注意力
|
||||
|
||||
> [!TIP]
|
||||
> 内存高效注意力优化内存使用 *和* [推理速度](./fp16#scaled-dot-product-attention)!
|
||||
|
||||
Transformers 注意力机制是内存密集型的,尤其对于长序列,因此您可以尝试使用不同且更内存高效的注意力类型。
|
||||
|
||||
默认情况下,如果安装了 PyTorch >= 2.0,则使用 [scaled dot-product attention (SDPA)](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html)。您无需对代码进行任何额外更改。
|
||||
|
||||
SDPA 还支持 [FlashAttention](https://github.com/Dao-AILab/flash-attention) 和 [xFormers](https://github.com/facebookresearch/xformers),以及 a
|
||||
这是一个原生的 C++ PyTorch 实现。它会根据您的输入自动选择最优的实现。
|
||||
|
||||
您可以使用 [`~ModelMixin.enable_xformers_memory_efficient_attention`] 方法显式地使用 xFormers。
|
||||
|
||||
```py
|
||||
# pip install xformers
|
||||
import torch
|
||||
from diffusers import StableDiffusionXLPipeline
|
||||
|
||||
pipeline = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
pipeline.enable_xformers_memory_efficient_attention()
|
||||
```
|
||||
|
||||
调用 [`~ModelMixin.disable_xformers_memory_efficient_attention`] 来禁用它。
|
||||
|
||||
```py
|
||||
pipeline.disable_xformers_memory_efficient_attention()
|
||||
```
|
||||
@@ -1,82 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。请参阅许可证了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# Metal Performance Shaders (MPS)
|
||||
|
||||
> [!TIP]
|
||||
> 带有 <img alt="MPS" src="https://img.shields.io/badge/MPS-000000?style=flat&logo=apple&logoColor=white%22"> 徽章的管道表示模型可以利用 Apple silicon 设备上的 MPS 后端进行更快的推理。欢迎提交 [Pull Request](https://github.com/huggingface/diffusers/compare) 来为缺少此徽章的管道添加它。
|
||||
|
||||
🤗 Diffusers 与 Apple silicon(M1/M2 芯片)兼容,使用 PyTorch 的 [`mps`](https://pytorch.org/docs/stable/notes/mps.html) 设备,该设备利用 Metal 框架来发挥 MacOS 设备上 GPU 的性能。您需要具备:
|
||||
|
||||
- 配备 Apple silicon(M1/M2)硬件的 macOS 计算机
|
||||
- macOS 12.6 或更高版本(推荐 13.0 或更高)
|
||||
- arm64 版本的 Python
|
||||
- [PyTorch 2.0](https://pytorch.org/get-started/locally/)(推荐)或 1.13(支持 `mps` 的最低版本)
|
||||
|
||||
`mps` 后端使用 PyTorch 的 `.to()` 接口将 Stable Diffusion 管道移动到您的 M1 或 M2 设备上:
|
||||
|
||||
```python
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
|
||||
pipe = pipe.to("mps")
|
||||
|
||||
# 如果您的计算机内存小于 64 GB,推荐使用
|
||||
pipe.enable_attention_slicing()
|
||||
|
||||
prompt = "a photo of an astronaut riding a horse on mars"
|
||||
image = pipe(prompt).images[0]
|
||||
image
|
||||
```
|
||||
|
||||
<Tip warning={true}>
|
||||
|
||||
PyTorch [mps](https://pytorch.org/docs/stable/notes/mps.html) 后端不支持大小超过 `2**32` 的 NDArray。如果您遇到此问题,请提交 [Issue](https://github.com/huggingface/diffusers/issues/new/choose) 以便我们调查。
|
||||
|
||||
</Tip>
|
||||
|
||||
如果您使用 **PyTorch 1.13**,您需要通过管道进行一次额外的"预热"传递。这是一个临时解决方法,用于解决首次推理传递产生的结果与后续传递略有不同的问题。您只需要执行此传递一次,并且在仅进行一次推理步骤后可以丢弃结果。
|
||||
|
||||
```diff
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5").to("mps")
|
||||
pipe.enable_attention_slicing()
|
||||
|
||||
prompt = "a photo of an astronaut riding a horse on mars"
|
||||
# 如果 PyTorch 版本是 1.13,进行首次"预热"传递
|
||||
+ _ = pipe(prompt, num_inference_steps=1)
|
||||
|
||||
# 预热传递后,结果与 CPU 设备上的结果匹配。
|
||||
image = pipe(prompt).images[0]
|
||||
```
|
||||
|
||||
## 故障排除
|
||||
|
||||
本节列出了使用 `mps` 后端时的一些常见问题及其解决方法。
|
||||
|
||||
### 注意力切片
|
||||
|
||||
M1/M2 性能对内存压力非常敏感。当发生这种情况时,系统会自动交换内存,这会显著降低性能。
|
||||
|
||||
为了防止这种情况发生,我们建议使用*注意力切片*来减少推理过程中的内存压力并防止交换。这在您的计算机系统内存少于 64GB 或生成非标准分辨率(大于 512×512 像素)的图像时尤其相关。在您的管道上调用 [`~DiffusionPipeline.enable_attention_slicing`] 函数:
|
||||
|
||||
```py
|
||||
from diffusers import DiffusionPipeline
|
||||
import torch
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("mps")
|
||||
pipeline.enable_attention_slicing()
|
||||
```
|
||||
|
||||
注意力切片将昂贵的注意力操作分多个步骤执行,而不是一次性完成。在没有统一内存的计算机中,它通常能提高约 20% 的性能,但我们观察到在大多数 Apple 芯片计算机中,除非您有 64GB 或更多 RAM,否则性能会*更好*。
|
||||
|
||||
### 批量推理
|
||||
|
||||
批量生成多个提示可能会导致崩溃或无法可靠工作。如果是这种情况,请尝试迭代而不是批量处理。
|
||||
@@ -1,59 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版(“许可证”)授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按“原样”分发,无任何明示或暗示的担保或条件。请参阅许可证了解特定语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# AWS Neuron
|
||||
|
||||
Diffusers 功能可在 [AWS Inf2 实例](https://aws.amazon.com/ec2/instance-types/inf2/)上使用,这些是由 [Neuron 机器学习加速器](https://aws.amazon.com/machine-learning/inferentia/)驱动的 EC2 实例。这些实例旨在提供更好的计算性能(更高的吞吐量、更低的延迟)和良好的成本效益,使其成为 AWS 用户将扩散模型部署到生产环境的良好选择。
|
||||
|
||||
[Optimum Neuron](https://huggingface.co/docs/optimum-neuron/en/index) 是 Hugging Face 库与 AWS 加速器之间的接口,包括 AWS [Trainium](https://aws.amazon.com/machine-learning/trainium/) 和 AWS [Inferentia](https://aws.amazon.com/machine-learning/inferentia/)。它支持 Diffusers 中的许多功能,并具有类似的 API,因此如果您已经熟悉 Diffusers,学习起来更容易。一旦您创建了 AWS Inf2 实例,请安装 Optimum Neuron。
|
||||
|
||||
```bash
|
||||
python -m pip install --upgrade-strategy eager optimum[neuronx]
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
我们提供预构建的 [Hugging Face Neuron 深度学习 AMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2)(DLAMI)和用于 Amazon SageMaker 的 Optimum Neuron 容器。建议正确设置您的环境。
|
||||
|
||||
</Tip>
|
||||
|
||||
下面的示例演示了如何在 inf2.8xlarge 实例上使用 Stable Diffusion XL 模型生成图像(一旦模型编译完成,您可以切换到更便宜的 inf2.xlarge 实例)。要生成一些图像,请使用 [`~optimum.neuron.NeuronStableDiffusionXLPipeline`] 类,该类类似于 Diffusers 中的 [`StableDiffusionXLPipeline`] 类。
|
||||
|
||||
与 Diffusers 不同,您需要将管道中的模型编译为 Neuron 格式,即 `.neuron`。运行以下命令将模型导出为 `.neuron` 格式。
|
||||
|
||||
```bash
|
||||
optimum-cli export neuron --model stabilityai/stable-diffusion-xl-base-1.0 \
|
||||
--batch_size 1 \
|
||||
--height 1024 `# 生成图像的高度(像素),例如 768, 1024` \
|
||||
--width 1024 `# 生成图像的宽度(像素),例如 768, 1024` \
|
||||
--num_images_per_prompt 1 `# 每个提示生成的图像数量,默认为 1` \
|
||||
--auto_cast matmul `# 仅转换矩阵乘法操作` \
|
||||
--auto_cast_type bf16 `# 将操作从 FP32 转换为 BF16` \
|
||||
sd_neuron_xl/
|
||||
```
|
||||
|
||||
现在使用预编译的 SDXL 模型生成一些图像。
|
||||
|
||||
```python
|
||||
>>> from optimum.neuron import Neu
|
||||
ronStableDiffusionXLPipeline
|
||||
|
||||
>>> stable_diffusion_xl = NeuronStableDiffusionXLPipeline.from_pretrained("sd_neuron_xl/")
|
||||
>>> prompt = "a pig with wings flying in floating US dollar banknotes in the air, skyscrapers behind, warm color palette, muted colors, detailed, 8k"
|
||||
>>> image = stable_diffusion_xl(prompt).images[0]
|
||||
```
|
||||
|
||||
<img
|
||||
src="https://huggingface.co/datasets/Jingya/document_images/resolve/main/optimum/neuron/sdxl_pig.png"
|
||||
width="256"
|
||||
height="256"
|
||||
alt="peggy generated by sdxl on inf2"
|
||||
/>
|
||||
|
||||
欢迎查看Optimum Neuron [文档](https://huggingface.co/docs/optimum-neuron/en/inference_tutorials/stable_diffusion#generate-images-with-stable-diffusion-models-on-aws-inferentia)中更多不同用例的指南和示例!
|
||||
@@ -1,77 +0,0 @@
|
||||
<!--版权所有 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按"原样"分发,无任何明示或暗示的担保或条件。请参阅许可证以了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# OpenVINO
|
||||
|
||||
🤗 [Optimum](https://github.com/huggingface/optimum-intel) 提供与 OpenVINO 兼容的 Stable Diffusion 管道,可在各种 Intel 处理器上执行推理(请参阅支持的设备[完整列表](https://docs.openvino.ai/latest/openvino_docs_OV_UG_supported_plugins_Supported_Devices.html))。
|
||||
|
||||
您需要安装 🤗 Optimum Intel,并使用 `--upgrade-strategy eager` 选项以确保 [`optimum-intel`](https://github.com/huggingface/optimum-intel) 使用最新版本:
|
||||
|
||||
```bash
|
||||
pip install --upgrade-strategy eager optimum["openvino"]
|
||||
```
|
||||
|
||||
本指南将展示如何使用 Stable Diffusion 和 Stable Diffusion XL (SDXL) 管道与 OpenVINO。
|
||||
|
||||
## Stable Diffusion
|
||||
|
||||
要加载并运行推理,请使用 [`~optimum.intel.OVStableDiffusionPipeline`]。如果您想加载 PyTorch 模型并即时转换为 OpenVINO 格式,请设置 `export=True`:
|
||||
|
||||
```python
|
||||
from optimum.intel import OVStableDiffusionPipeline
|
||||
|
||||
model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
|
||||
pipeline = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
|
||||
prompt = "sailing ship in storm by Rembrandt"
|
||||
image = pipeline(prompt).images[0]
|
||||
|
||||
# 别忘了保存导出的模型
|
||||
pipeline.save_pretrained("openvino-sd-v1-5")
|
||||
```
|
||||
|
||||
为了进一步加速推理,静态重塑模型。如果您更改任何参数,例如输出高度或宽度,您需要再次静态重塑模型。
|
||||
|
||||
```python
|
||||
# 定义与输入和期望输出相关的形状
|
||||
batch_size, num_images, height, width = 1, 1, 512, 512
|
||||
|
||||
# 静态重塑模型
|
||||
pipeline.reshape(batch_size, height, width, num_images)
|
||||
# 在推理前编译模型
|
||||
pipeline.compile()
|
||||
|
||||
image = pipeline(
|
||||
prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_images_per_prompt=num_images,
|
||||
).images[0]
|
||||
```
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/intel/openvino/stable_diffusion_v1_5_sail_boat_rembrandt.png">
|
||||
</div>
|
||||
|
||||
您可以在 🤗 Optimum [文档](https://huggingface.co/docs/optimum/intel/inference#stable-diffusion) 中找到更多示例,Stable Diffusion 支持文本到图像、图像到图像和修复。
|
||||
|
||||
## Stable Diffusion XL
|
||||
|
||||
要加载并运行 SDXL 推理,请使用 [`~optimum.intel.OVStableDiffusionXLPipeline`]:
|
||||
|
||||
```python
|
||||
from optimum.intel import OVStableDiffusionXLPipeline
|
||||
|
||||
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
pipeline = OVStableDiffusionXLPipeline.from_pretrained(model_id)
|
||||
prompt = "sailing ship in storm by Rembrandt"
|
||||
image = pipeline(prompt).images[0]
|
||||
```
|
||||
|
||||
为了进一步加速推理,可以如Stable Diffusion部分所示[静态重塑](#stable-diffusion)模型。
|
||||
|
||||
您可以在🤗 Optimum[文档](https://huggingface.co/docs/optimum/intel/inference#stable-diffusion-xl)中找到更多示例,并且在OpenVINO中运行SDXL支持文本到图像和图像到图像。
|
||||
@@ -1,497 +0,0 @@
|
||||
# ParaAttention
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-performance.png">
|
||||
</div>
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/hunyuan-video-performance.png">
|
||||
</div>
|
||||
|
||||
大型图像和视频生成模型,如 [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) 和 [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo),由于其规模,可能对实时应用和部署构成推理挑战。
|
||||
|
||||
[ParaAttention](https://github.com/chengzeyi/ParaAttention) 是一个实现了**上下文并行**和**第一块缓存**的库,可以与其他技术(如 torch.compile、fp8 动态量化)结合使用,以加速推理。
|
||||
|
||||
本指南将展示如何在 NVIDIA L20 GPU 上对 FLUX.1-dev 和 HunyuanVideo 应用 ParaAttention。
|
||||
在我们的基线基准测试中,除了 HunyuanVideo 为避免内存不足错误外,未应用任何优化。
|
||||
|
||||
我们的基线基准测试显示,FLUX.1-dev 能够在 28 步中生成 1024x1024 分辨率图像,耗时 26.36 秒;HunyuanVideo 能够在 30 步中生成 129 帧 720p 分辨率视频,耗时 3675.71 秒。
|
||||
|
||||
> [!TIP]
|
||||
> 对于更快的上下文并行推理,请尝试使用支持 NVLink 的 NVIDIA A100 或 H100 GPU(如果可用),尤其是在 GPU 数量较多时。
|
||||
|
||||
## 第一块缓存
|
||||
|
||||
缓存模型中 transformer 块的输出并在后续推理步骤中重用它们,可以降低计算成本并加速推理。
|
||||
|
||||
然而,很难决定何时重用缓存以确保生成图像或视频的质量。ParaAttention 直接使用**第一个 transformer 块输出的残差差异**来近似模型输出之间的差异。当差异足够小时,重用先前推理步骤的残差差异。换句话说,跳过去噪步骤。
|
||||
|
||||
这在 FLUX.1-dev 和 HunyuanVideo 推理上实现了 2 倍加速,且质量非常好。
|
||||
|
||||
<figure>
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/ada-cache.png" alt="Cache in Diffusion Transformer" />
|
||||
<figcaption>AdaCache 的工作原理,第一块缓存是其变体</figcaption>
|
||||
</figure>
|
||||
|
||||
<hfoptions id="first-block-cache">
|
||||
<hfoption id="FLUX-1.dev">
|
||||
|
||||
要在 FLUX.1-dev 上应用第一块缓存,请调用 `apply_cache_on_pipe`,如下所示。0.08 是 FLUX 模型的默认残差差异值。
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(pipe, residual_diff_thre
|
||||
shold=0.08)
|
||||
|
||||
# 启用内存节省
|
||||
# pipe.enable_model_cpu_offload()
|
||||
# pipe.enable_sequential_cpu_offload()
|
||||
|
||||
begin = time.time()
|
||||
image = pipe(
|
||||
"A cat holding a sign that says hello world",
|
||||
num_inference_steps=28,
|
||||
).images[0]
|
||||
end = time.time()
|
||||
print(f"Time: {end - begin:.2f}s")
|
||||
|
||||
print("Saving image to flux.png")
|
||||
image.save("flux.png")
|
||||
```
|
||||
|
||||
| 优化 | 原始 | FBCache rdt=0.06 | FBCache rdt=0.08 | FBCache rdt=0.10 | FBCache rdt=0.12 |
|
||||
| - | - | - | - | - | - |
|
||||
| 预览 |  |  |  |  |  |
|
||||
| 墙时间 (s) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |
|
||||
|
||||
First Block Cache 将推理速度降低到 17.01 秒,与基线相比,或快 1.55 倍,同时保持几乎零质量损失。
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="HunyuanVideo">
|
||||
|
||||
要在 HunyuanVideo 上应用 First Block Cache,请使用 `apply_cache_on_pipe`,如下所示。0.06 是 HunyuanVideo 模型的默认残差差值。
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
model_id = "tencent/HunyuanVideo"
|
||||
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
|
||||
model_id,
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16,
|
||||
revision="refs/pr/18",
|
||||
)
|
||||
pipe = HunyuanVideoPipeline.from_pretrained(
|
||||
model_id,
|
||||
transformer=transformer,
|
||||
torch_dtype=torch.float16,
|
||||
revision="refs/pr/18",
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(pipe, residual_diff_threshold=0.6)
|
||||
|
||||
pipe.vae.enable_tiling()
|
||||
|
||||
begin = time.time()
|
||||
output = pipe(
|
||||
prompt="A cat walks on the grass, realistic",
|
||||
height=720,
|
||||
width=1280,
|
||||
num_frames=129,
|
||||
num_inference_steps=30,
|
||||
).frames[0]
|
||||
end = time.time()
|
||||
print(f"Time: {end - begin:.2f}s")
|
||||
|
||||
print("Saving video to hunyuan_video.mp4")
|
||||
export_to_video(output, "hunyuan_video.mp4", fps=15)
|
||||
```
|
||||
|
||||
<video controls>
|
||||
<source src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/hunyuan-video-original.mp4" type="video/mp4">
|
||||
您的浏览器不支持视频标签。
|
||||
</video>
|
||||
|
||||
<small> HunyuanVideo 无 FBCache </small>
|
||||
|
||||
<video controls>
|
||||
<source src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/hunyuan-video-fbc.mp4" type="video/mp4">
|
||||
Your browser does not support the video tag.
|
||||
</video>
|
||||
|
||||
<small> HunyuanVideo 与 FBCache </small>
|
||||
|
||||
First Block Cache 将推理速度降低至 2271.06 秒,相比基线快了 1.62 倍,同时保持了几乎为零的质量损失。
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## fp8 量化
|
||||
|
||||
fp8 动态量化进一步加速推理并减少内存使用。为了使用 8 位 [NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/),必须对激活和权重进行量化。
|
||||
|
||||
使用 `float8_weight_only` 和 `float8_dynamic_activation_float8_weight` 来量化文本编码器和变换器模型。
|
||||
|
||||
默认量化方法是逐张量量化,但如果您的 GPU 支持逐行量化,您也可以尝试它以获得更好的准确性。
|
||||
|
||||
使用以下命令安装 [torchao](https://github.com/pytorch/ao/tree/main)。
|
||||
|
||||
```bash
|
||||
pip3 install -U torch torchao
|
||||
```
|
||||
|
||||
[torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) 使用 `mode="max-autotune-no-cudagraphs"` 或 `mode="max-autotune"` 选择最佳内核以获得性能。如果是第一次调用模型,编译可能会花费很长时间,但一旦模型编译完成,这是值得的。
|
||||
|
||||
此示例仅量化变换器模型,但您也可以量化文本编码器以进一步减少内存使用。
|
||||
|
||||
> [!TIP]
|
||||
> 动态量化可能会显著改变模型输出的分布,因此您需要将 `residual_diff_threshold` 设置为更大的值以使其生效。
|
||||
|
||||
<hfoptions id="fp8-quantization">
|
||||
<hfoption id="FLUX-1.dev">
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(
|
||||
pipe,
|
||||
residual_diff_threshold=0.12, # 使用更大的值以使缓存生效
|
||||
)
|
||||
|
||||
from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
|
||||
|
||||
quantize_(pipe.text_encoder, float8_weight_only())
|
||||
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
|
||||
pipe.transformer = torch.compile(
|
||||
pipe.transformer, mode="max-autotune-no-cudagraphs",
|
||||
)
|
||||
|
||||
# 启用内存节省
|
||||
# pipe.enable_model_cpu_offload()
|
||||
# pipe.enable_sequential_cpu_offload()
|
||||
|
||||
for i in range(2):
|
||||
begin = time.time()
|
||||
image = pipe(
|
||||
"A cat holding a sign that says hello world",
|
||||
num_inference_steps=28,
|
||||
).images[0]
|
||||
end = time.time()
|
||||
if i == 0:
|
||||
print(f"预热时间: {end - begin:.2f}s")
|
||||
else:
|
||||
print(f"时间: {end - begin:.2f}s")
|
||||
|
||||
print("保存图像到 flux.png")
|
||||
image.save("flux.png")
|
||||
```
|
||||
|
||||
fp8 动态量化和 torch.compile 将推理速度降低至 7.56 秒,相比基线快了 3.48 倍。
|
||||
</hfoption>
|
||||
<hfoption id="HunyuanVideo">
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
model_id = "tencent/HunyuanVideo"
|
||||
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
|
||||
model_id,
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16,
|
||||
revision="refs/pr/18",
|
||||
)
|
||||
pipe = HunyuanVideoPipeline.from_pretrained(
|
||||
model_id,
|
||||
transformer=transformer,
|
||||
torch_dtype=torch.float16,
|
||||
revision="refs/pr/18",
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(pipe)
|
||||
|
||||
from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
|
||||
|
||||
quantize_(pipe.text_encoder, float8_weight_only())
|
||||
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
|
||||
pipe.transformer = torch.compile(
|
||||
pipe.transformer, mode="max-autotune-no-cudagraphs",
|
||||
)
|
||||
|
||||
# Enable memory savings
|
||||
pipe.vae.enable_tiling()
|
||||
# pipe.enable_model_cpu_offload()
|
||||
# pipe.enable_sequential_cpu_offload()
|
||||
|
||||
for i in range(2):
|
||||
begin = time.time()
|
||||
output = pipe(
|
||||
prompt="A cat walks on the grass, realistic",
|
||||
height=720,
|
||||
width=1280,
|
||||
num_frames=129,
|
||||
num_inference_steps=1 if i == 0 else 30,
|
||||
).frames[0]
|
||||
end = time.time()
|
||||
if i == 0:
|
||||
print(f"Warm up time: {end - begin:.2f}s")
|
||||
else:
|
||||
print(f"Time: {end - begin:.2f}s")
|
||||
|
||||
print("Saving video to hunyuan_video.mp4")
|
||||
export_to_video(output, "hunyuan_video.mp4", fps=15)
|
||||
```
|
||||
|
||||
NVIDIA L20 GPU 仅有 48GB 内存,在编译后且如果未调用 `enable_model_cpu_offload` 时,可能会遇到内存不足(OOM)错误,因为 HunyuanVideo 在高分辨率和大量帧数运行时具有非常大的激活张量。对于内存少于 80GB 的 GPU,可以尝试降低分辨率和帧数来避免 OOM 错误。
|
||||
|
||||
大型视频生成模型通常受注意力计算而非全连接层的瓶颈限制。这些模型不会从量化和 torch.compile 中显著受益。
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 上下文并行性
|
||||
|
||||
上下文并行性并行化推理并随多个 GPU 扩展。ParaAttention 组合设计允许您将上下文并行性与第一块缓存和动态量化结合使用。
|
||||
|
||||
> [!TIP]
|
||||
> 请参考 [ParaAttention](https://github.com/chengzeyi/ParaAttention/tree/main) 仓库获取详细说明和如何使用多个 GPU 扩展推理的示例。
|
||||
|
||||
如果推理过程需要持久化和可服务,建议使用 [torch.multiprocessing](https://pytorch.org/docs/stable/multiprocessing.html) 编写您自己的推理处理器。这可以消除启动进程以及加载和重新编译模型的开销。
|
||||
|
||||
<hfoptions id="context-parallelism">
|
||||
<hfoption id="FLUX-1.dev">
|
||||
|
||||
以下代码示例结合了第一块缓存、fp8动态量化、torch.compile和上下文并行,以实现最快的推理速度。
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
dist.init_process_group()
|
||||
|
||||
torch.cuda.set_device(dist.get_rank())
|
||||
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.context_parallel import init_context_parallel_mesh
|
||||
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
|
||||
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
|
||||
|
||||
mesh = init_context_parallel_mesh(
|
||||
pipe.device.type,
|
||||
max_ring_dim_size=2,
|
||||
)
|
||||
parallelize_pipe(
|
||||
pipe,
|
||||
mesh=mesh,
|
||||
)
|
||||
parallelize_vae(pipe.vae, mesh=mesh._flatten())
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(
|
||||
pipe,
|
||||
residual_diff_threshold=0.12, # 使用较大的值以使缓存生效
|
||||
)
|
||||
|
||||
from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
|
||||
|
||||
quantize_(pipe.text_encoder, float8_weight_only())
|
||||
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
|
||||
torch._inductor.config.reorder_for_compute_comm_overlap = True
|
||||
pipe.transformer = torch.compile(
|
||||
pipe.transformer, mode="max-autotune-no-cudagraphs",
|
||||
)
|
||||
|
||||
# 启用内存节省
|
||||
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
|
||||
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())
|
||||
|
||||
for i in range(2):
|
||||
begin = time.time()
|
||||
image = pipe(
|
||||
"A cat holding a sign that says hello world",
|
||||
num_inference_steps=28,
|
||||
output_type="pil" if dist.get_rank() == 0 else "pt",
|
||||
).images[0]
|
||||
end = time.time()
|
||||
if dist.get_rank() == 0:
|
||||
if i == 0:
|
||||
print(f"预热时间: {end - begin:.2f}s")
|
||||
else:
|
||||
print(f"时间: {end - begin:.2f}s")
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
print("将图像保存到flux.png")
|
||||
image.save("flux.png")
|
||||
|
||||
dist.destroy_process_group()
|
||||
```
|
||||
|
||||
保存到`run_flux.py`并使用[torchrun](https://pytorch.org/docs/stable/elastic/run.html)启动。
|
||||
|
||||
```bash
|
||||
# 使用--nproc_per_node指定GPU数量
|
||||
torchrun --nproc_per_node=2 run_flux.py
|
||||
```
|
||||
|
||||
推理速度降至8.20秒,相比基线快了3.21倍,使用2个NVIDIA L20 GPU。在4个L20上,推理速度为3.90秒,快了6.75倍。
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="HunyuanVideo">
|
||||
|
||||
以下代码示例结合了第一块缓存和上下文并行,以实现最快的推理速度。
|
||||
|
||||
```python
|
||||
import time
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
dist.init_process_group()
|
||||
|
||||
torch.cuda.set_device(dist.get_rank())
|
||||
|
||||
model_id = "tencent/HunyuanVideo"
|
||||
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
|
||||
model_id,
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16,
|
||||
revision="refs/pr/18",
|
||||
)
|
||||
pipe = HunyuanVideoPipeline.from_pretrained(
|
||||
model_id,
|
||||
transformer=transformer,
|
||||
torch_dtype=torch.float16,
|
||||
revision="refs/pr/18",
|
||||
).to("cuda")
|
||||
|
||||
from para_attn.context_parallel import init_context_parallel_mesh
|
||||
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
|
||||
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
|
||||
|
||||
mesh = init_context_parallel_mesh(
|
||||
pipe.device.type,
|
||||
)
|
||||
parallelize_pipe(
|
||||
pipe,
|
||||
mesh=mesh,
|
||||
)
|
||||
parallelize_vae(pipe.vae, mesh=mesh._flatten())
|
||||
|
||||
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
|
||||
|
||||
apply_cache_on_pipe(pipe)
|
||||
|
||||
# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
|
||||
#
|
||||
# torch._inductor.config.reorder_for_compute_comm_overlap = True
|
||||
#
|
||||
# quantize_(pipe.text_encoder, float8_weight_only())
|
||||
# quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
|
||||
# pipe.transformer = torch.compile(
|
||||
# pipe.transformer, mode="max-autotune-no-cudagraphs",
|
||||
# )
|
||||
|
||||
# 启用内存节省
|
||||
pipe.vae.enable_tiling()
|
||||
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
|
||||
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())
|
||||
|
||||
for i in range(2):
|
||||
begin = time.time()
|
||||
output = pipe(
|
||||
prompt="A cat walks on the grass, realistic",
|
||||
height=720,
|
||||
width=1280,
|
||||
num_frames=129,
|
||||
num_inference_steps=1 if i == 0 else 30,
|
||||
output_type="pil" if dist.get_rank() == 0 else "pt",
|
||||
).frames[0]
|
||||
end = time.time()
|
||||
if dist.get_rank() == 0:
|
||||
if i == 0:
|
||||
print(f"预热时间: {end - begin:.2f}s")
|
||||
else:
|
||||
print(f"时间: {end - begin:.2f}s")
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
print("保存视频到 hunyuan_video.mp4")
|
||||
export_to_video(output, "hunyuan_video.mp4", fps=15)
|
||||
|
||||
dist.destroy_process_group()
|
||||
```
|
||||
|
||||
保存到 `run_hunyuan_video.py` 并使用 [torchrun](https://pytorch.org/docs/stable/elastic/run.html) 启动。
|
||||
|
||||
```bash
|
||||
# 使用 --nproc_per_node 指定 GPU 数量
|
||||
torchrun --nproc_per_node=8 run_hunyuan_video.py
|
||||
```
|
||||
|
||||
推理速度降低到 649.23 秒,相比基线快 5.66 倍,使用 8 个 NVIDIA L20 GPU。
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 基准测试
|
||||
|
||||
<hfoptions id="conclusion">
|
||||
<hfoption id="FLUX-1.dev">
|
||||
|
||||
| GPU 类型 | GPU 数量 | 优化 | 墙钟时间 (s) | 加速比 |
|
||||
| - | - | - | - | - |
|
||||
| NVIDIA L20 | 1 | 基线 | 26.36 | 1.00x |
|
||||
| NVIDIA L20 | 1 | FBCache (rdt=0.08) | 17.01 | 1.55x |
|
||||
| NVIDIA L20 | 1 | FP8 DQ | 13.40 | 1.96x |
|
||||
| NVIDIA L20 | 1 | FBCache (rdt=0.12) + FP8 DQ | 7.56 | 3.48x |
|
||||
| NVIDIA L20 | 2 | FBCache (rdt=0.12) + FP8 DQ + CP | 4.92 | 5.35x |
|
||||
| NVIDIA L20 | 4 | FBCache (rdt=0.12) + FP8 DQ + CP | 3.90 | 6.75x |
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="HunyuanVideo">
|
||||
|
||||
| GPU 类型 | GPU 数量 | 优化 | 墙钟时间 (s) | 加速比 |
|
||||
| - | - | - | - | - |
|
||||
| NVIDIA L20 | 1 | 基线 | 3675.71 | 1.00x |
|
||||
| NVIDIA
|
||||
L20 | 1 | FBCache | 2271.06 | 1.62x |
|
||||
| NVIDIA L20 | 2 | FBCache + CP | 1132.90 | 3.24x |
|
||||
| NVIDIA L20 | 4 | FBCache + CP | 718.15 | 5.12x |
|
||||
| NVIDIA L20 | 8 | FBCache + CP | 649.23 | 5.66x |
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
@@ -1,184 +0,0 @@
|
||||
# Pruna
|
||||
|
||||
[Pruna](https://github.com/PrunaAI/pruna) 是一个模型优化框架,提供多种优化方法——量化、剪枝、缓存、编译——以加速推理并减少内存使用。以下是优化方法的概览。
|
||||
|
||||
| 技术 | 描述 | 速度 | 内存 | 质量 |
|
||||
|------------|---------------------------------------------------------------------------------------|:----:|:----:|:----:|
|
||||
| `batcher` | 将多个输入分组在一起同时处理,提高计算效率并减少处理时间。 | ✅ | ❌ | ➖ |
|
||||
| `cacher` | 存储计算的中间结果以加速后续操作。 | ✅ | ➖ | ➖ |
|
||||
| `compiler` | 为特定硬件优化模型指令。 | ✅ | ➖ | ➖ |
|
||||
| `distiller`| 训练一个更小、更简单的模型来模仿一个更大、更复杂的模型。 | ✅ | ✅ | ❌ |
|
||||
| `quantizer`| 降低权重和激活的精度,减少内存需求。 | ✅ | ✅ | ❌ |
|
||||
| `pruner` | 移除不重要或冗余的连接和神经元,产生一个更稀疏、更高效的网络。 | ✅ | ✅ | ❌ |
|
||||
| `recoverer`| 在压缩后恢复模型的性能。 | ➖ | ➖ | ✅ |
|
||||
| `factorizer`| 将多个小矩阵乘法批处理为一个大型融合操作。 | ✅ | ➖ | ➖ |
|
||||
| `enhancer` | 通过应用后处理算法(如去噪或上采样)来增强模型输出。 | ❌ | - | ✅ |
|
||||
|
||||
✅ (改进), ➖ (大致相同), ❌ (恶化)
|
||||
|
||||
在 [Pruna 文档](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms) 中探索所有优化方法。
|
||||
|
||||
## 安装
|
||||
|
||||
使用以下命令安装 Pruna。
|
||||
|
||||
```bash
|
||||
pip install pruna
|
||||
```
|
||||
|
||||
## 优化 Diffusers 模型
|
||||
|
||||
Diffusers 模型支持广泛的优化算法,如下所示。
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/PrunaAI/documentation-images/resolve/main/diffusers/diffusers_combinations.png" alt="Diffusers 模型支持的优化算法概览">
|
||||
</div>
|
||||
|
||||
下面的示例使用 factorizer、compiler 和 cacher 算法的组合优化 [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)。这种组合将推理速度加速高达 4.2 倍,并将峰值 GPU 内存使用从 34.7GB 减少到 28.0GB,同时几乎保持相同的输出质量。
|
||||
|
||||
> [!TIP]
|
||||
> 参考 [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html) 文档以了解更多关于该操作的信息。
|
||||
本示例中使用的优化技术。
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/PrunaAI/documentation-images/resolve/main/diffusers/flux_combination.png" alt="用于FLUX.1-dev的优化技术展示,结合了因子分解器、编译器和缓存器算法">
|
||||
</div>
|
||||
|
||||
首先定义一个包含要使用的优化算法的`SmashConfig`。要优化模型,将管道和`SmashConfig`用`smash`包装,然后像往常一样使用管道进行推理。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
from pruna import PrunaModel, SmashConfig, smash
|
||||
|
||||
# 加载模型
|
||||
# 使用小GPU内存尝试segmind/Segmind-Vega或black-forest-labs/FLUX.1-schnell
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
torch_dtype=torch.bfloat16
|
||||
).to("cuda")
|
||||
|
||||
# 定义配置
|
||||
smash_config = SmashConfig()
|
||||
smash_config["factorizer"] = "qkv_diffusers"
|
||||
smash_config["compiler"] = "torch_compile"
|
||||
smash_config["torch_compile_target"] = "module_list"
|
||||
smash_config["cacher"] = "fora"
|
||||
smash_config["fora_interval"] = 2
|
||||
|
||||
# 为了获得最佳速度结果,可以添加这些配置
|
||||
# 但它们会将预热时间从1.5分钟增加到10分钟
|
||||
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
|
||||
# smash_config["quantizer"] = "torchao"
|
||||
# smash_config["torchao_quant_type"] = "fp8dq"
|
||||
# smash_config["torchao_excluded_modules"] = "norm+embedding"
|
||||
|
||||
# 优化模型
|
||||
smashed_pipe = smash(pipe, smash_config)
|
||||
|
||||
# 运行模型
|
||||
smashed_pipe("a knitted purple prune").images[0]
|
||||
```
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/PrunaAI/documentation-images/resolve/main/diffusers/flux_smashed_comparison.png">
|
||||
</div>
|
||||
|
||||
优化后,我们可以使用Hugging Face Hub共享和加载优化后的模型。
|
||||
|
||||
```python
|
||||
# 保存模型
|
||||
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")
|
||||
|
||||
# 加载模型
|
||||
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")
|
||||
```
|
||||
|
||||
## 评估和基准测试Diffusers模型
|
||||
|
||||
Pruna提供了[EvaluationAgent](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)来评估优化后模型的质量。
|
||||
|
||||
我们可以定义我们关心的指标,如总时间和吞吐量,以及要评估的数据集。我们可以定义一个模型并将其传递给`EvaluationAgent`。
|
||||
|
||||
<hfoptions id="eval">
|
||||
<hfoption id="optimized model">
|
||||
|
||||
我们可以通过使用`EvaluationAgent`加载和评估优化后的模型,并将其传递给`Task`。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
from pruna import PrunaModel
|
||||
from pruna.data.pruna_datamodule import PrunaDataModule
|
||||
from pruna.evaluation.evaluation_agent import EvaluationAgent
|
||||
from pruna.evaluation.metrics import (
|
||||
ThroughputMetric,
|
||||
TorchMetricWrapper,
|
||||
TotalTimeMetric,
|
||||
)
|
||||
from pruna.evaluation.task import Task
|
||||
|
||||
# define the device
|
||||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
|
||||
# 加载模型
|
||||
# 使用小GPU内存尝试 PrunaAI/Segmind-Vega-smashed 或 PrunaAI/FLUX.1-dev-smashed
|
||||
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")
|
||||
|
||||
# 定义指标
|
||||
metrics = [
|
||||
TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
|
||||
ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
|
||||
TorchMetricWrapper("clip"),
|
||||
]
|
||||
|
||||
# 定义数据模块
|
||||
datamodule = PrunaDataModule.from_string("LAION256")
|
||||
datamodule.limit_datasets(10)
|
||||
|
||||
# 定义任务和评估代理
|
||||
task = Task(metrics, datamodule=datamodule, device=device)
|
||||
eval_agent = EvaluationAgent(task)
|
||||
|
||||
# 评估优化模型并卸载到CPU
|
||||
smashed_pipe.move_to_device(device)
|
||||
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
|
||||
smashed_pipe.move_to_device("cpu")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="standalone model">
|
||||
|
||||
除了比较优化模型与基础模型,您还可以评估独立的 `diffusers` 模型。这在您想评估模型性能而不考虑优化时非常有用。我们可以通过使用 `PrunaModel` 包装器并运行 `EvaluationAgent` 来实现。
|
||||
|
||||
```python
|
||||
import torch
|
||||
from diffusers import FluxPipeline
|
||||
|
||||
from pruna import PrunaModel
|
||||
|
||||
# 加载模型
|
||||
# 使用小GPU内存尝试 PrunaAI/Segmind-Vega-smashed 或 PrunaAI/FLUX.1-dev-smashed
|
||||
pipe = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
torch_dtype=torch.bfloat16
|
||||
).to("cpu")
|
||||
wrapped_pipe = PrunaModel(model=pipe)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
现在您已经了解了如何优化和评估您的模型,可以开始使用 Pruna 来优化您自己的模型了。幸运的是,我们有许多示例来帮助您入门。
|
||||
|
||||
> [!TIP]
|
||||
> 有关基准测试 Flux 的更多详细信息,请查看 [宣布 FLUX-Juiced:最快的图像生成端点(快 2.6 倍)!](https://huggingface.co/blog/PrunaAI/flux-fastest-image-generation-endpoint) 博客文章和 [InferBench](https://huggingface.co/spaces/PrunaAI/InferBench) 空间。
|
||||
|
||||
## 参考
|
||||
|
||||
- [Pruna](https://github.com/pruna-ai/pruna)
|
||||
- [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)
|
||||
- [Pruna 评估](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)
|
||||
- [Pruna 教程](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html)
|
||||
@@ -1,200 +0,0 @@
|
||||
<!--版权所有 2024 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版(“许可证”)授权;除非符合许可证,否则不得使用此文件。
|
||||
您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按“原样”分发,不附带任何明示或暗示的担保或条件。有关许可证的特定语言,请参阅许可证。
|
||||
-->
|
||||
|
||||
# 编译和卸载量化模型
|
||||
|
||||
优化模型通常涉及[推理速度](./fp16)和[内存使用](./memory)之间的权衡。例如,虽然[缓存](./cache)可以提高推理速度,但它也会增加内存消耗,因为它需要存储中间注意力层的输出。一种更平衡的优化策略结合了量化模型、[torch.compile](./fp16#torchcompile) 和各种[卸载方法](./memory#offloading)。
|
||||
|
||||
> [!TIP]
|
||||
> 查看 [torch.compile](./fp16#torchcompile) 指南以了解更多关于编译以及如何在此处应用的信息。例如,区域编译可以显著减少编译时间,而不会放弃任何加速。
|
||||
|
||||
对于图像生成,结合量化和[模型卸载](./memory#model-offloading)通常可以在质量、速度和内存之间提供最佳权衡。组卸载对于图像生成效果不佳,因为如果计算内核更快完成,通常不可能*完全*重叠数据传输。这会导致 CPU 和 GPU 之间的一些通信开销。
|
||||
|
||||
对于视频生成,结合量化和[组卸载](./memory#group-offloading)往往更好,因为视频模型更受计算限制。
|
||||
|
||||
下表提供了优化策略组合及其对 Flux 延迟和内存使用的影响的比较。
|
||||
|
||||
| 组合 | 延迟 (s) | 内存使用 (GB) |
|
||||
|---|---|---|
|
||||
| 量化 | 32.602 | 14.9453 |
|
||||
| 量化, torch.compile | 25.847 | 14.9448 |
|
||||
| 量化, torch.compile, 模型 CPU 卸载 | 32.312 | 12.2369 |
|
||||
<small>这些结果是在 Flux 上使用 RTX 4090 进行基准测试的。transformer 和 text_encoder 组件已量化。如果您有兴趣评估自己的模型,请参考[基准测试脚本](https://gist.github.com/sayakpaul/0db9d8eeeb3d2a0e5ed7cf0d9ca19b7d)。</small>
|
||||
|
||||
本指南将向您展示如何使用 [bitsandbytes](../quantization/bitsandbytes#torchcompile) 编译和卸载量化模型。确保您正在使用 [PyTorch nightly](https://pytorch.org/get-started/locally/) 和最新版本的 bitsandbytes。
|
||||
|
||||
```bash
|
||||
pip install -U bitsandbytes
|
||||
```
|
||||
|
||||
## 量化和 torch.compile
|
||||
|
||||
首先通过[量化](../quantization/overview)模型来减少存储所需的内存,并[编译](./fp16#torchcompile)它以加速推理。
|
||||
|
||||
配置 [Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) `capture_dynamic_output_shape_ops = True` 以在编译 bitsandbytes 模型时处理动态输出。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
from diffusers.quantizers import PipelineQuantizationConfig
|
||||
|
||||
torch._dynamo.config.capture_dynamic_output_shape_ops = True
|
||||
|
||||
# 量化
|
||||
pipeline_quant_config = PipelineQuantizationConfig(
|
||||
quant_backend="bitsandbytes_4bit",
|
||||
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
|
||||
components_to_quantize=["transformer", "text_encoder_2"],
|
||||
)
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
quantization_config=pipeline_quant_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
# 编译
|
||||
pipeline.transformer.to(memory_format=torch.channels_last)
|
||||
pipeline.transformer.compile(mode="max-autotune", fullgraph=True)
|
||||
pipeline("""
|
||||
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
|
||||
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
|
||||
"""
|
||||
).images[0]
|
||||
```
|
||||
|
||||
## 量化、torch.compile 和卸载
|
||||
|
||||
除了量化和 torch.compile,如果您需要进一步减少内存使用,可以尝试卸载。卸载根据需要将各种层或模型组件从 CPU 移动到 GPU 进行计算。
|
||||
|
||||
在卸载期间配置 [Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) `cache_size_limit` 以避免过多的重新编译,并设置 `capture_dynamic_output_shape_ops = True` 以在编译 bitsandbytes 模型时处理动态输出。
|
||||
|
||||
<hfoptions id="offloading">
|
||||
<hfoption id="model CPU offloading">
|
||||
|
||||
[模型 CPU 卸载](./memory#model-offloading) 将单个管道组件(如 transformer 模型)在需要计算时移动到 GPU。否则,它会被卸载到 CPU。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
from diffusers.quantizers import PipelineQuantizationConfig
|
||||
|
||||
torch._dynamo.config.cache_size_limit = 1000
|
||||
torch._dynamo.config.capture_dynamic_output_shape_ops = True
|
||||
|
||||
# 量化
|
||||
pipeline_quant_config = PipelineQuantizationConfig(
|
||||
quant_backend="bitsandbytes_4bit",
|
||||
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
|
||||
components_to_quantize=["transformer", "text_encoder_2"],
|
||||
)
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
quantization_config=pipeline_quant_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
# 模型 CPU 卸载
|
||||
pipeline.enable_model_cpu_offload()
|
||||
|
||||
# 编译
|
||||
pipeline.transformer.compile()
|
||||
pipeline(
|
||||
"cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain"
|
||||
).images[0]
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="group offloading">
|
||||
|
||||
[组卸载](./memory#group-offloading) 将单个管道组件(如变换器模型)的内部层移动到 GPU 进行计算,并在不需要时将其卸载。同时,它使用 [CUDA 流](./memory#cuda-stream) 功能来预取下一层以执行。
|
||||
|
||||
通过重叠计算和数据传输,它比模型 CPU 卸载更快,同时还能节省内存。
|
||||
|
||||
```py
|
||||
# pip install ftfy
|
||||
import torch
|
||||
from diffusers import AutoModel, DiffusionPipeline
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
from diffusers.utils import export_to_video
|
||||
from diffusers.quantizers import PipelineQuantizationConfig
|
||||
from transformers import UMT5EncoderModel
|
||||
|
||||
torch._dynamo.config.cache_size_limit = 1000
|
||||
torch._dynamo.config.capture_dynamic_output_shape_ops = True
|
||||
|
||||
# 量化
|
||||
pipeline_quant_config = PipelineQuantizationConfig(
|
||||
quant_backend="bitsandbytes_4bit",
|
||||
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
|
||||
components_to_quantize=["transformer", "text_encoder"],
|
||||
)
|
||||
|
||||
text_encoder = UMT5EncoderModel.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16
|
||||
)
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
|
||||
quantization_config=pipeline_quant_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda")
|
||||
|
||||
# 组卸载
|
||||
onload_device = torch.device("cuda")
|
||||
offload_device = torch.device("cpu")
|
||||
|
||||
pipeline.transformer.enable_group_offload(
|
||||
onload_device=onload_device,
|
||||
offload_device=offload_device,
|
||||
offload_type="leaf_level",
|
||||
use_stream=True,
|
||||
non_blocking=True
|
||||
)
|
||||
pipeline.vae.enable_group_offload(
|
||||
onload_device=onload_device,
|
||||
offload_device=offload_device,
|
||||
offload_type="leaf_level",
|
||||
use_stream=True,
|
||||
non_blocking=True
|
||||
)
|
||||
apply_group_offloading(
|
||||
pipeline.text_encoder,
|
||||
onload_device=onload_device,
|
||||
offload_type="leaf_level",
|
||||
use_stream=True,
|
||||
non_blocking=True
|
||||
)
|
||||
|
||||
# 编译
|
||||
pipeline.transformer.compile()
|
||||
|
||||
prompt = """
|
||||
The camera rushes from far to near in a low-angle shot,
|
||||
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
|
||||
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
|
||||
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
|
||||
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
|
||||
"""
|
||||
negative_prompt = """
|
||||
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
|
||||
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
|
||||
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
|
||||
"""
|
||||
|
||||
output = pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_frames=81,
|
||||
guidance_scale=5.0,
|
||||
).frames[0]
|
||||
export_to_video(output, "output.mp4", fps=16)
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
@@ -1,182 +0,0 @@
|
||||
# T-GATE
|
||||
|
||||
[T-GATE](https://github.com/HaozheLiu-ST/T-GATE/tree/main) 通过跳过交叉注意力计算一旦收敛,加速了 [Stable Diffusion](../api/pipelines/stable_diffusion/overview)、[PixArt](../api/pipelines/pixart) 和 [Latency Consistency Model](../api/pipelines/latent_consistency_models.md) 管道的推理。此方法不需要任何额外训练,可以将推理速度提高 10-50%。T-GATE 还与 [DeepCache](./deepcache) 等其他优化方法兼容。
|
||||
|
||||
开始之前,请确保安装 T-GATE。
|
||||
|
||||
```bash
|
||||
pip install tgate
|
||||
pip install -U torch diffusers transformers accelerate DeepCache
|
||||
```
|
||||
|
||||
要使用 T-GATE 与管道,您需要使用其对应的加载器。
|
||||
|
||||
| 管道 | T-GATE 加载器 |
|
||||
|---|---|
|
||||
| PixArt | TgatePixArtLoader |
|
||||
| Stable Diffusion XL | TgateSDXLLoader |
|
||||
| Stable Diffusion XL + DeepCache | TgateSDXLDeepCacheLoader |
|
||||
| Stable Diffusion | TgateSDLoader |
|
||||
| Stable Diffusion + DeepCache | TgateSDDeepCacheLoader |
|
||||
|
||||
接下来,创建一个 `TgateLoader`,包含管道、门限步骤(停止计算交叉注意力的时间步)和推理步骤数。然后在管道上调用 `tgate` 方法,提供提示、门限步骤和推理步骤数。
|
||||
|
||||
让我们看看如何为几个不同的管道启用此功能。
|
||||
|
||||
<hfoptions id="pipelines">
|
||||
<hfoption id="PixArt">
|
||||
|
||||
使用 T-GATE 加速 `PixArtAlphaPipeline`:
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import PixArtAlphaPipeline
|
||||
from tgate import TgatePixArtLoader
|
||||
|
||||
pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16)
|
||||
|
||||
gate_step = 8
|
||||
inference_step = 25
|
||||
pipe = TgatePixArtLoader(
|
||||
pipe,
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step,
|
||||
).to("cuda")
|
||||
|
||||
image = pipe.tgate(
|
||||
"An alpaca made of colorful building blocks, cyberpunk.",
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step,
|
||||
).images[0]
|
||||
```
|
||||
</hfoption>
|
||||
<hfoption id="Stable Diffusion XL">
|
||||
|
||||
使用 T-GATE 加速 `StableDiffusionXLPipeline`:
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import StableDiffusionXLPipeline
|
||||
from diffusers import DPMSolverMultistepScheduler
|
||||
from tgate import TgateSDXLLoader
|
||||
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16",
|
||||
use_safetensors=True,
|
||||
)
|
||||
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
gate_step = 10
|
||||
inference_step = 25
|
||||
pipe = TgateSDXLLoader(
|
||||
pipe,
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step,
|
||||
).to("cuda")
|
||||
|
||||
image = pipe.tgate(
|
||||
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step
|
||||
).images[0]
|
||||
```
|
||||
</hfoption>
|
||||
<hfoption id="StableDiffusionXL with DeepCache">
|
||||
|
||||
使用 [DeepCache](https://github.co 加速 `StableDiffusionXLPipeline`
|
||||
m/horseee/DeepCache) 和 T-GATE:
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import StableDiffusionXLPipeline
|
||||
from diffusers import DPMSolverMultistepScheduler
|
||||
from tgate import TgateSDXLDeepCacheLoader
|
||||
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16",
|
||||
use_safetensors=True,
|
||||
)
|
||||
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
gate_step = 10
|
||||
inference_step = 25
|
||||
pipe = TgateSDXLDeepCacheLoader(
|
||||
pipe,
|
||||
cache_interval=3,
|
||||
cache_branch_id=0,
|
||||
).to("cuda")
|
||||
|
||||
image = pipe.tgate(
|
||||
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step
|
||||
).images[0]
|
||||
```
|
||||
</hfoption>
|
||||
<hfoption id="Latent Consistency Model">
|
||||
|
||||
使用 T-GATE 加速 `latent-consistency/lcm-sdxl`:
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import StableDiffusionXLPipeline
|
||||
from diffusers import UNet2DConditionModel, LCMScheduler
|
||||
from diffusers import DPMSolverMultistepScheduler
|
||||
from tgate import TgateSDXLLoader
|
||||
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
"latent-consistency/lcm-sdxl",
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16",
|
||||
)
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
unet=unet,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16",
|
||||
)
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
|
||||
|
||||
gate_step = 1
|
||||
inference_step = 4
|
||||
pipe = TgateSDXLLoader(
|
||||
pipe,
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step,
|
||||
lcm=True
|
||||
).to("cuda")
|
||||
|
||||
image = pipe.tgate(
|
||||
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k.",
|
||||
gate_step=gate_step,
|
||||
num_inference_steps=inference_step
|
||||
).images[0]
|
||||
```
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
T-GATE 还支持 [`StableDiffusionPipeline`] 和 [PixArt-alpha/PixArt-LCM-XL-2-1024-MS](https://hf.co/PixArt-alpha/PixArt-LCM-XL-2-1024-MS)。
|
||||
|
||||
## 基准测试
|
||||
| 模型 | MACs | 参数 | 延迟 | 零样本 10K-FID on MS-COCO |
|
||||
|-----------------------|----------|-----------|---------|---------------------------|
|
||||
| SD-1.5 | 16.938T | 859.520M | 7.032s | 23.927 |
|
||||
| SD-1.5 w/ T-GATE | 9.875T | 815.557M | 4.313s | 20.789 |
|
||||
| SD-2.1 | 38.041T | 865.785M | 16.121s | 22.609 |
|
||||
| SD-2.1 w/ T-GATE | 22.208T | 815.433 M | 9.878s | 19.940 |
|
||||
| SD-XL | 149.438T | 2.570B | 53.187s | 24.628 |
|
||||
| SD-XL w/ T-GATE | 84.438T | 2.024B | 27.932s | 22.738 |
|
||||
| Pixart-Alpha | 107.031T | 611.350M | 61.502s | 38.669 |
|
||||
| Pixart-Alpha w/ T-GATE | 65.318T | 462.585M | 37.867s | 35.825 |
|
||||
| DeepCache (SD-XL) | 57.888T | - | 19.931s | 23.755 |
|
||||
| DeepCache 配合 T-GATE | 43.868T | - | 14.666秒 | 23.999 |
|
||||
| LCM (SD-XL) | 11.955T | 2.570B | 3.805秒 | 25.044 |
|
||||
| LCM 配合 T-GATE | 11.171T | 2.024B | 3.533秒 | 25.028 |
|
||||
| LCM (Pixart-Alpha) | 8.563T | 611.350M | 4.733秒 | 36.086 |
|
||||
| LCM 配合 T-GATE | 7.623T | 462.585M | 4.543秒 | 37.048 |
|
||||
|
||||
延迟测试基于 NVIDIA 1080TI,MACs 和 Params 使用 [calflops](https://github.com/MrYxJ/calculate-flops.pytorch) 计算,FID 使用 [PytorchFID](https://github.com/mseitzer/pytorch-fid) 计算。
|
||||
@@ -1,90 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版(“许可证”)授权;除非遵守许可证,否则不得使用此文件。
|
||||
您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按“原样”分发,不附带任何明示或暗示的担保或条件。请参阅许可证以了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# 令牌合并
|
||||
|
||||
[令牌合并](https://huggingface.co/papers/2303.17604)(ToMe)在基于 Transformer 的网络的前向传递中逐步合并冗余令牌/补丁,这可以加速 [`StableDiffusionPipeline`] 的推理延迟。
|
||||
|
||||
从 `pip` 安装 ToMe:
|
||||
|
||||
```bash
|
||||
pip install tomesd
|
||||
```
|
||||
|
||||
您可以使用 [`tomesd`](https://github.com/dbolya/tomesd) 库中的 [`apply_patch`](https://github.com/dbolya/tomesd?tab=readme-ov-file#usage) 函数:
|
||||
|
||||
```diff
|
||||
from diffusers import StableDiffusionPipeline
|
||||
import torch
|
||||
import tomesd
|
||||
|
||||
pipeline = StableDiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True,
|
||||
).to("cuda")
|
||||
+ tomesd.apply_patch(pipeline, ratio=0.5)
|
||||
|
||||
image = pipeline("a photo of an astronaut riding a horse on mars").images[0]
|
||||
```
|
||||
|
||||
`apply_patch` 函数公开了多个[参数](https://github.com/dbolya/tomesd#usage),以帮助在管道推理速度和生成令牌的质量之间取得平衡。最重要的参数是 `ratio`,它控制在前向传递期间合并的令牌数量。
|
||||
|
||||
如[论文](https://huggingface.co/papers/2303.17604)中所述,ToMe 可以在显著提升推理速度的同时,很大程度上保留生成图像的质量。通过增加 `ratio`,您可以进一步加速推理,但代价是图像质量有所下降。
|
||||
|
||||
为了测试生成图像的质量,我们从 [Parti Prompts](https://parti.research.google/) 中采样了一些提示,并使用 [`StableDiffusionPipeline`] 进行了推理,设置如下:
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/diffusers/docs-images/resolve/main/tome/tome_samples.png">
|
||||
</div>
|
||||
|
||||
我们没有注意到生成样本的质量有任何显著下降,您可以在此 [WandB 报告](https://wandb.ai/sayakpaul/tomesd-results/runs/23j4bj3i?workspace=)中查看生成的样本。如果您有兴趣重现此实验,请使用此[脚本](https://gist.github.com/sayakpaul/8cac98d7f22399085a060992f411ecbd)。
|
||||
|
||||
## 基准测试
|
||||
|
||||
我们还在启用 [xFormers](https://huggingface.co/docs/diffusers/optimization/xformers) 的情况下,对 [`StableDiffusionPipeline`] 上 `tomesd` 的影响进行了基准测试,涵盖了多个图像分辨率。结果
|
||||
结果是从以下开发环境中的A100和V100 GPU获得的:
|
||||
|
||||
```bash
|
||||
- `diffusers` 版本:0.15.1
|
||||
- Python 版本:3.8.16
|
||||
- PyTorch 版本(GPU?):1.13.1+cu116 (True)
|
||||
- Huggingface_hub 版本:0.13.2
|
||||
- Transformers 版本:4.27.2
|
||||
- Accelerate 版本:0.18.0
|
||||
- xFormers 版本:0.0.16
|
||||
- tomesd 版本:0.1.2
|
||||
```
|
||||
|
||||
要重现此基准测试,请随意使用此[脚本](https://gist.github.com/sayakpaul/27aec6bca7eb7b0e0aa4112205850335)。结果以秒为单位报告,并且在适用的情况下,我们报告了使用ToMe和ToMe + xFormers时相对于原始管道的加速百分比。
|
||||
|
||||
| **GPU** | **分辨率** | **批处理大小** | **原始** | **ToMe** | **ToMe + xFormers** |
|
||||
|----------|----------------|----------------|-------------|----------------|---------------------|
|
||||
| **A100** | 512 | 10 | 6.88 | 5.26 (+23.55%) | 4.69 (+31.83%) |
|
||||
| | 768 | 10 | OOM | 14.71 | 11 |
|
||||
| | | 8 | OOM | 11.56 | 8.84 |
|
||||
| | | 4 | OOM | 5.98 | 4.66 |
|
||||
| | | 2 | 4.99 | 3.24 (+35.07%) | 2.1 (+37.88%) |
|
||||
| | | 1 | 3.29 | 2.24 (+31.91%) | 2.03 (+38.3%) |
|
||||
| | 1024 | 10 | OOM | OOM | OOM |
|
||||
| | | 8 | OOM | OOM | OOM |
|
||||
| | | 4 | OOM | 12.51 | 9.09 |
|
||||
| | | 2 | OOM | 6.52 | 4.96 |
|
||||
| | | 1 | 6.4 | 3.61 (+43.59%) | 2.81 (+56.09%) |
|
||||
| **V100** | 512 | 10 | OOM | 10.03 | 9.29 |
|
||||
| | | 8 | OOM | 8.05 | 7.47 |
|
||||
| | | 4 | 5.7 | 4.3 (+24.56%) | 3.98 (+30.18%) |
|
||||
| | | 2 | 3.14 | 2.43 (+22.61%) | 2.27 (+27.71%) |
|
||||
| | | 1 | 1.88 | 1.57 (+16.49%) | 1.57 (+16.49%) |
|
||||
| | 768 | 10 | OOM | OOM | 23.67 |
|
||||
| | | 8 | OOM | OOM | 18.81 |
|
||||
| | | 4 | OOM | 11.81 | 9.7 |
|
||||
| | | 2 | OOM | 6.27 | 5.2 |
|
||||
| | | 1 | 5.43 | 3.38 (+37.75%) | 2.82 (+48.07%) |
|
||||
| | 1024 | 10 | OOM |
|
||||
如上表所示,`tomesd` 带来的加速效果在更大的图像分辨率下变得更加明显。有趣的是,使用 `tomesd` 可以在更高分辨率如 1024x1024 上运行管道。您可能还可以通过 [`torch.compile`](fp16#torchcompile) 进一步加速推理。
|
||||
@@ -1,119 +0,0 @@
|
||||
# xDiT
|
||||
|
||||
[xDiT](https://github.com/xdit-project/xDiT) 是一个推理引擎,专为大规模并行部署扩散变换器(DiTs)而设计。xDiT 提供了一套用于扩散模型的高效并行方法,以及 GPU 内核加速。
|
||||
|
||||
xDiT 支持四种并行方法,包括[统一序列并行](https://huggingface.co/papers/2405.07719)、[PipeFusion](https://huggingface.co/papers/2405.14430)、CFG 并行和数据并行。xDiT 中的这四种并行方法可以以混合方式配置,优化通信模式以最适合底层网络硬件。
|
||||
|
||||
与并行化正交的优化侧重于加速单个 GPU 的性能。除了利用知名的注意力优化库外,我们还利用编译加速技术,如 torch.compile 和 onediff。
|
||||
|
||||
xDiT 的概述如下所示。
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/methods/xdit_overview.png">
|
||||
</div>
|
||||
您可以使用以下命令安装 xDiT:
|
||||
|
||||
```bash
|
||||
pip install xfuser
|
||||
```
|
||||
|
||||
以下是一个使用 xDiT 加速 Diffusers 模型推理的示例。
|
||||
|
||||
```diff
|
||||
import torch
|
||||
from diffusers import StableDiffusion3Pipeline
|
||||
|
||||
from xfuser import xFuserArgs, xDiTParallel
|
||||
from xfuser.config import FlexibleArgumentParser
|
||||
from xfuser.core.distributed import get_world_group
|
||||
|
||||
def main():
|
||||
+ parser = FlexibleArgumentParser(description="xFuser Arguments")
|
||||
+ args = xFuserArgs.add_cli_args(parser).parse_args()
|
||||
+ engine_args = xFuserArgs.from_cli_args(args)
|
||||
+ engine_config, input_config = engine_args.create_config()
|
||||
|
||||
local_rank = get_world_group().local_rank
|
||||
pipe = StableDiffusion3Pipeline.from_pretrained(
|
||||
pretrained_model_name_or_path=engine_config.model_config.model,
|
||||
torch_dtype=torch.float16,
|
||||
).to(f"cuda:{local_rank}")
|
||||
|
||||
# 在这里对管道进行任何操作
|
||||
|
||||
+ pipe = xDiTParallel(pipe, engine_config, input_config)
|
||||
|
||||
pipe(
|
||||
height=input_config.height,
|
||||
width=input_config.height,
|
||||
prompt=input_config.prompt,
|
||||
num_inference_steps=input_config.num_inference_steps,
|
||||
output_type=input_config.output_type,
|
||||
generator=torch.Generator(device="cuda").manual_seed(input_config.seed),
|
||||
)
|
||||
|
||||
+ if input_config.output_type == "pil":
|
||||
+ pipe.save("results", "stable_diffusion_3")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
如您所见,我们只需要使用 xDiT 中的 xFuserArgs 来获取配置参数,并将这些参数与来自 Diffusers 库的管道对象一起传递给 xDiTParallel,即可完成对 Diffusers 中特定管道的并行化。
|
||||
|
||||
xDiT 运行时参数可以在命令行中使用 `-h` 查看,您可以参考此[使用](https://github.com/xdit-project/xDiT?tab=readme-ov-file#2-usage)示例以获取更多详细信息。
|
||||
ils。
|
||||
|
||||
xDiT 需要使用 torchrun 启动,以支持其多节点、多 GPU 并行能力。例如,以下命令可用于 8-GPU 并行推理:
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node=8 ./inference.py --model models/FLUX.1-dev --data_parallel_degree 2 --ulysses_degree 2 --ring_degree 2 --prompt "A snowy mountain" "A small dog" --num_inference_steps 50
|
||||
```
|
||||
|
||||
## 支持的模型
|
||||
|
||||
在 xDiT 中支持 Diffusers 模型的一个子集,例如 Flux.1、Stable Diffusion 3 等。最新支持的模型可以在[这里](https://github.com/xdit-project/xDiT?tab=readme-ov-file#-supported-dits)找到。
|
||||
|
||||
## 基准测试
|
||||
我们在不同机器上测试了各种模型,以下是一些基准数据。
|
||||
|
||||
### Flux.1-schnell
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/flux/Flux-2k-L40.png">
|
||||
</div>
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/flux/Flux-2K-A100.png">
|
||||
</div>
|
||||
|
||||
### Stable Diffusion 3
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/sd3/L40-SD3.png">
|
||||
</div>
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/sd3/A100-SD3.png">
|
||||
</div>
|
||||
|
||||
### HunyuanDiT
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/hunuyuandit/L40-HunyuanDiT.png">
|
||||
</div>
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/hunuyuandit/V100-HunyuanDiT.png">
|
||||
</div>
|
||||
|
||||
<div class="flex justify-center">
|
||||
<img src="https://huggingface.co/datasets/xDiT/documentation-images/resolve/main/performance/hunuyuandit/T4-HunyuanDiT.png">
|
||||
</div>
|
||||
|
||||
更详细的性能指标可以在我们的 [GitHub 页面](https://github.com/xdit-project/xDiT?tab=readme-ov-file#perf) 上找到。
|
||||
|
||||
## 参考文献
|
||||
|
||||
[xDiT-project](https://github.com/xdit-project/xDiT)
|
||||
|
||||
[USP: A Unified Sequence Parallelism Approach for Long Context Generative AI](https://huggingface.co/papers/2405.07719)
|
||||
|
||||
[PipeFusion: Displaced Patch Pipeline Parallelism for Inference of Diffusion Transformer Models](https://huggingface.co/papers/2405.14430)
|
||||
@@ -1,239 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本(“许可证”)授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,根据许可证分发的软件按“原样”分发,不附带任何明示或暗示的担保或条件。请参阅许可证了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# 分布式推理
|
||||
|
||||
在分布式设置中,您可以使用 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) 或 [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) 在多个 GPU 上运行推理,这对于并行生成多个提示非常有用。
|
||||
|
||||
本指南将向您展示如何使用 🤗 Accelerate 和 PyTorch Distributed 进行分布式推理。
|
||||
|
||||
## 🤗 Accelerate
|
||||
|
||||
🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) 是一个旨在简化在分布式设置中训练或运行推理的库。它简化了设置分布式环境的过程,让您可以专注于您的 PyTorch 代码。
|
||||
|
||||
首先,创建一个 Python 文件并初始化一个 [`accelerate.PartialState`] 来创建分布式环境;您的设置会自动检测,因此您无需明确定义 `rank` 或 `world_size`。将 [`DiffusionPipeline`] 移动到 `distributed_state.device` 以为每个进程分配一个 GPU。
|
||||
|
||||
现在使用 [`~accelerate.PartialState.split_between_processes`] 实用程序作为上下文管理器,自动在进程数之间分发提示。
|
||||
|
||||
```py
|
||||
import torch
|
||||
from accelerate import PartialState
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
|
||||
)
|
||||
distributed_state = PartialState()
|
||||
pipeline.to(distributed_state.device)
|
||||
|
||||
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
|
||||
result = pipeline(prompt).images[0]
|
||||
result.save(f"result_{distributed_state.process_index}.png")
|
||||
```
|
||||
|
||||
使用 `--num_processes` 参数指定要使用的 GPU 数量,并调用 `accelerate launch` 来运行脚本:
|
||||
|
||||
```bash
|
||||
accelerate launch run_distributed.py --num_processes=2
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
参考这个最小示例 [脚本](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba) 以在多个 GPU 上运行推理。要了解更多信息,请查看 [使用 🤗 Accelerate 进行分布式推理](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate) 指南。
|
||||
|
||||
</Tip>
|
||||
|
||||
## PyTorch Distributed
|
||||
|
||||
PyTorch 支持 [`DistributedDataParallel`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html),它启用了数据
|
||||
并行性。
|
||||
|
||||
首先,创建一个 Python 文件并导入 `torch.distributed` 和 `torch.multiprocessing` 来设置分布式进程组,并为每个 GPU 上的推理生成进程。您还应该初始化一个 [`DiffusionPipeline`]:
|
||||
|
||||
```py
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
sd = DiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
|
||||
)
|
||||
```
|
||||
|
||||
您需要创建一个函数来运行推理;[`init_process_group`](https://pytorch.org/docs/stable/distributed.html?highlight=init_process_group#torch.distributed.init_process_group) 处理创建一个分布式环境,指定要使用的后端类型、当前进程的 `rank` 以及参与进程的数量 `world_size`。如果您在 2 个 GPU 上并行运行推理,那么 `world_size` 就是 2。
|
||||
|
||||
将 [`DiffusionPipeline`] 移动到 `rank`,并使用 `get_rank` 为每个进程分配一个 GPU,其中每个进程处理不同的提示:
|
||||
|
||||
```py
|
||||
def run_inference(rank, world_size):
|
||||
dist.init_process_group("nccl", rank=rank, world_size=world_size)
|
||||
|
||||
sd.to(rank)
|
||||
|
||||
if torch.distributed.get_rank() == 0:
|
||||
prompt = "a dog"
|
||||
elif torch.distributed.get_rank() == 1:
|
||||
prompt = "a cat"
|
||||
|
||||
image = sd(prompt).images[0]
|
||||
image.save(f"./{'_'.join(prompt)}.png")
|
||||
```
|
||||
|
||||
要运行分布式推理,调用 [`mp.spawn`](https://pytorch.org/docs/stable/multiprocessing.html#torch.multiprocessing.spawn) 在 `world_size` 定义的 GPU 数量上运行 `run_inference` 函数:
|
||||
|
||||
```py
|
||||
def main():
|
||||
world_size = 2
|
||||
mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
```
|
||||
|
||||
完成推理脚本后,使用 `--nproc_per_node` 参数指定要使用的 GPU 数量,并调用 `torchrun` 来运行脚本:
|
||||
|
||||
```bash
|
||||
torchrun run_distributed.py --nproc_per_node=2
|
||||
```
|
||||
|
||||
> [!TIP]
|
||||
> 您可以在 [`DiffusionPipeline`] 中使用 `device_map` 将其模型级组件分布在多个设备上。请参考 [设备放置](../tutorials/inference_with_big_models#device-placement) 指南了解更多信息。
|
||||
|
||||
## 模型分片
|
||||
|
||||
现代扩散系统,如 [Flux](../api/pipelines/flux),非常大且包含多个模型。例如,[Flux.1-Dev](https://hf.co/black-forest-labs/FLUX.1-dev) 由两个文本编码器 - [T5-XXL](https://hf.co/google/t5-v1_1-xxl) 和 [CLIP-L](https://hf.co/openai/clip-vit-large-patch14) - 一个 [扩散变换器](../api/models/flux_transformer),以及一个 [VAE](../api/models/autoencoderkl) 组成。对于如此大的模型,在消费级 GPU 上运行推理可能具有挑战性。
|
||||
|
||||
模型分片是一种技术,当模型无法容纳在单个 GPU 上时,将模型分布在多个 GPU 上。下面的示例假设有两个 16GB GPU 可用于推理。
|
||||
|
||||
开始使用文本编码器计算文本嵌入。通过设置 `device_map="balanced"` 将文本编码器保持在两个GPU上。`balanced` 策略将模型均匀分布在所有可用GPU上。使用 `max_memory` 参数为每个GPU上的每个文本编码器分配最大内存量。
|
||||
|
||||
> [!TIP]
|
||||
> **仅** 在此步骤加载文本编码器!扩散变换器和VAE在后续步骤中加载以节省内存。
|
||||
|
||||
```py
|
||||
from diffusers import FluxPipeline
|
||||
import torch
|
||||
|
||||
prompt = "a photo of a dog with cat-like look"
|
||||
|
||||
pipeline = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
transformer=None,
|
||||
vae=None,
|
||||
device_map="balanced",
|
||||
max_memory={0: "16GB", 1: "16GB"},
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
with torch.no_grad():
|
||||
print("Encoding prompts.")
|
||||
prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
|
||||
prompt=prompt, prompt_2=None, max_sequence_length=512
|
||||
)
|
||||
```
|
||||
|
||||
一旦文本嵌入计算完成,从GPU中移除它们以为扩散变换器腾出空间。
|
||||
|
||||
```py
|
||||
import gc
|
||||
|
||||
def flush():
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_max_memory_allocated()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
||||
del pipeline.text_encoder
|
||||
del pipeline.text_encoder_2
|
||||
del pipeline.tokenizer
|
||||
del pipeline.tokenizer_2
|
||||
del pipeline
|
||||
|
||||
flush()
|
||||
```
|
||||
|
||||
接下来加载扩散变换器,它有125亿参数。这次,设置 `device_map="auto"` 以自动将模型分布在两个16GB GPU上。`auto` 策略由 [Accelerate](https://hf.co/docs/accelerate/index) 支持,并作为 [大模型推理](https://hf.co/docs/accelerate/concept_guides/big_model_inference) 功能的一部分可用。它首先将模型分布在最快的设备(GPU)上,然后在需要时移动到较慢的设备如CPU和硬盘。将模型参数存储在较慢设备上的权衡是推理延迟较慢。
|
||||
|
||||
```py
|
||||
from diffusers import AutoModel
|
||||
import torch
|
||||
|
||||
transformer = AutoModel.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
subfolder="transformer",
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
```
|
||||
|
||||
> [!TIP]
|
||||
> 在任何时候,您可以尝试 `print(pipeline.hf_device_map)` 来查看各种模型如何在设备上分布。这对于跟踪模型的设备放置很有用。您也可以尝试 `print(transformer.hf_device_map)` 来查看变换器模型如何在设备上分片。
|
||||
|
||||
将变换器模型添加到管道中以进行去噪,但将其他模型级组件如文本编码器和VAE设置为 `None`,因为您还不需要它们。
|
||||
|
||||
```py
|
||||
pipeline = FluxPipeline.from_pretrained(
|
||||
"black-forest-labs/FLUX.1-dev",
|
||||
text_encoder=None,
|
||||
text_encoder_2=None,
|
||||
tokenizer=None,
|
||||
tokenizer_2=None,
|
||||
vae=None,
|
||||
transformer=transformer,
|
||||
torch_dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
print("Running denoising.")
|
||||
height, width = 768, 1360
|
||||
latents = pipeline(
|
||||
|
||||
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
num_inference_steps=50,
|
||||
guidance_scale=3.5,
|
||||
height=height,
|
||||
width=width,
|
||||
output_type="latent",
|
||||
).images
|
||||
```
|
||||
|
||||
从内存中移除管道和变换器,因为它们不再需要。
|
||||
|
||||
```py
|
||||
del pipeline.transformer
|
||||
del pipeline
|
||||
|
||||
flush()
|
||||
```
|
||||
|
||||
最后,使用变分自编码器(VAE)将潜在表示解码为图像。VAE通常足够小,可以在单个GPU上加载。
|
||||
|
||||
```py
|
||||
from diffusers import AutoencoderKL
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
import torch
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16).to("cuda")
|
||||
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
|
||||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
|
||||
|
||||
with torch.no_grad():
|
||||
print("运行解码中。")
|
||||
latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
|
||||
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
|
||||
|
||||
image = vae.decode(latents, return_dict=False)[0]
|
||||
image = image_processor.postprocess(image, output_type="pil")
|
||||
image[0].save("split_transformer.png")
|
||||
```
|
||||
|
||||
通过选择性加载和卸载在特定阶段所需的模型,并将最大模型分片到多个GPU上,可以在消费级GPU上运行大型模型的推理。
|
||||
@@ -1,643 +0,0 @@
|
||||
<!--版权所有 2025 The HuggingFace Team。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版(“许可证”)授权;除非遵守许可证,否则不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,否则根据许可证分发的软件按“原样”分发,不附带任何明示或暗示的担保或条件。请参阅许可证以了解特定的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# DreamBooth
|
||||
|
||||
[DreamBooth](https://huggingface.co/papers/2208.12242) 是一种训练技术,通过仅训练少数主题或风格的图像来更新整个扩散模型。它通过在提示中关联一个特殊词与示例图像来工作。
|
||||
|
||||
如果您在 vRAM 有限的 GPU 上训练,应尝试在训练命令中启用 `gradient_checkpointing` 和 `mixed_precision` 参数。您还可以通过使用 [xFormers](../optimization/xformers) 的内存高效注意力来减少内存占用。JAX/Flax 训练也支持在 TPU 和 GPU 上进行高效训练,但不支持梯度检查点或 xFormers。如果您想使用 Flax 更快地训练,应拥有内存 >30GB 的 GPU。
|
||||
|
||||
本指南将探索 [train_dreambooth.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) 脚本,帮助您更熟悉它,以及如何根据您的用例进行适配。
|
||||
|
||||
在运行脚本之前,请确保从源代码安装库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/huggingface/diffusers
|
||||
cd diffusers
|
||||
pip install .
|
||||
```
|
||||
|
||||
导航到包含训练脚本的示例文件夹,并安装脚本所需的依赖项:
|
||||
|
||||
<hfoptions id="installation">
|
||||
<hfoption id="PyTorch">
|
||||
|
||||
```bash
|
||||
cd examples/dreambooth
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Flax">
|
||||
|
||||
```bash
|
||||
cd examples/dreambooth
|
||||
pip install -r requirements_flax.txt
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
<Tip>
|
||||
|
||||
🤗 Accelerate 是一个库,用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。
|
||||
|
||||
</Tip>
|
||||
|
||||
初始化 🤗 Accelerate 环境:
|
||||
|
||||
```bash
|
||||
accelerate config
|
||||
```
|
||||
|
||||
要设置默认的 🤗 Accelerate 环境而不选择任何配置:
|
||||
|
||||
```bash
|
||||
accelerate config default
|
||||
```
|
||||
|
||||
或者,如果您的环境不支持交互式 shell,例如笔记本,您可以使用:
|
||||
|
||||
```py
|
||||
from accelerate.utils import write_basic_config
|
||||
|
||||
write_basic_config()
|
||||
```
|
||||
|
||||
最后,如果您想在自己的数据集上训练模型,请查看 [创建用于训练的数据集](create_dataset) 指南,了解如何创建与
|
||||
训练脚本。
|
||||
|
||||
<Tip>
|
||||
|
||||
以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分,但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多,请随时阅读[脚本](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py),并告诉我们如果您有任何问题或疑虑。
|
||||
|
||||
</Tip>
|
||||
|
||||
## 脚本参数
|
||||
|
||||
<Tip warning={true}>
|
||||
|
||||
DreamBooth 对训练超参数非常敏感,容易过拟合。阅读 [使用 🧨 Diffusers 训练 Stable Diffusion 与 Dreambooth](https://huggingface.co/blog/dreambooth) 博客文章,了解针对不同主题的推荐设置,以帮助您选择合适的超参数。
|
||||
|
||||
</Tip>
|
||||
|
||||
训练脚本提供了许多参数来自定义您的训练运行。所有参数及其描述都可以在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L228) 函数中找到。参数设置了默认值,这些默认值应该开箱即用效果不错,但如果您愿意,也可以在训练命令中设置自己的值。
|
||||
|
||||
例如,要以 bf16 格式进行训练:
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--mixed_precision="bf16"
|
||||
```
|
||||
|
||||
一些基本且重要的参数需要了解和指定:
|
||||
|
||||
- `--pretrained_model_name_or_path`: Hub 上的模型名称或预训练模型的本地路径
|
||||
- `--instance_data_dir`: 包含训练数据集(示例图像)的文件夹路径
|
||||
- `--instance_prompt`: 包含示例图像特殊单词的文本提示
|
||||
- `--train_text_encoder`: 是否也训练文本编码器
|
||||
- `--output_dir`: 保存训练后模型的位置
|
||||
- `--push_to_hub`: 是否将训练后的模型推送到 Hub
|
||||
- `--checkpointing_steps`: 模型训练时保存检查点的频率;这在训练因某种原因中断时很有用,您可以通过在训练命令中添加 `--resume_from_checkpoint` 来从该检查点继续训练
|
||||
|
||||
### Min-SNR 加权
|
||||
|
||||
[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练,以实现更快的收敛。训练脚本支持预测 `epsilon`(噪声)或 `v_prediction`,但 Min-SNR 与两种预测类型都兼容。此加权策略仅由 PyTorch 支持,在 Flax 训练脚本中不可用。
|
||||
|
||||
添加 `--snr_gamma` 参数并将其设置为推荐值 5.0:
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--snr_gamma=5.0
|
||||
```
|
||||
|
||||
### 先验保持损失
|
||||
|
||||
先验保持损失是一种使用模型自身生成的样本来帮助它学习如何生成更多样化图像的方法。因为这些生成的样本图像属于您提供的图像相同的类别,它们帮助模型 r
|
||||
etain 它已经学到的关于类别的知识,以及它如何利用已经了解的类别信息来创建新的组合。
|
||||
|
||||
- `--with_prior_preservation`: 是否使用先验保留损失
|
||||
- `--prior_loss_weight`: 控制先验保留损失对模型的影响程度
|
||||
- `--class_data_dir`: 包含生成的类别样本图像的文件夹路径
|
||||
- `--class_prompt`: 描述生成的样本图像类别的文本提示
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--with_prior_preservation \
|
||||
--prior_loss_weight=1.0 \
|
||||
--class_data_dir="path/to/class/images" \
|
||||
--class_prompt="text prompt describing class"
|
||||
```
|
||||
|
||||
### 训练文本编码器
|
||||
|
||||
为了提高生成输出的质量,除了 UNet 之外,您还可以训练文本编码器。这需要额外的内存,并且您需要一个至少有 24GB 显存的 GPU。如果您拥有必要的硬件,那么训练文本编码器会产生更好的结果,尤其是在生成面部图像时。通过以下方式启用此选项:
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--train_text_encoder
|
||||
```
|
||||
|
||||
## 训练脚本
|
||||
|
||||
DreamBooth 附带了自己的数据集类:
|
||||
|
||||
- [`DreamBoothDataset`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L604): 预处理图像和类别图像,并对提示进行分词以用于训练
|
||||
- [`PromptDataset`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L738): 生成提示嵌入以生成类别图像
|
||||
|
||||
如果您启用了[先验保留损失](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L842),类别图像在此处生成:
|
||||
|
||||
```py
|
||||
sample_dataset = PromptDataset(args.class_prompt, num_new_images)
|
||||
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)
|
||||
|
||||
sample_dataloader = accelerator.prepare(sample_dataloader)
|
||||
pipeline.to(accelerator.device)
|
||||
|
||||
for example in tqdm(
|
||||
sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
|
||||
):
|
||||
images = pipeline(example["prompt"]).images
|
||||
```
|
||||
|
||||
接下来是 [`main()`](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L799) 函数,它处理设置训练数据集和训练循环本身。脚本加载 [tokenizer](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L898)、[scheduler 和 models](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L912C1-L912C1):
|
||||
|
||||
```py
|
||||
# Load the tokenizer
|
||||
if args.tokenizer_name:
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
|
||||
elif args.pretrained_model_name_or_path:
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path,
|
||||
subfolder="tokenizer",
|
||||
revision=args.revision,
|
||||
use_fast=False,
|
||||
)
|
||||
|
||||
# 加载调度器和模型
|
||||
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
text_encoder = text_encoder_cls.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
|
||||
)
|
||||
|
||||
if model_has_vae(args):
|
||||
vae = AutoencoderKL.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
|
||||
)
|
||||
else:
|
||||
vae = None
|
||||
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
|
||||
)
|
||||
```
|
||||
|
||||
然后,是时候[创建训练数据集](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L1073)和从`DreamBoothDataset`创建DataLoader:
|
||||
|
||||
```py
|
||||
train_dataset = DreamBoothDataset(
|
||||
instance_data_root=args.instance_data_dir,
|
||||
instance_prompt=args.instance_prompt,
|
||||
class_data_root=args.class_data_dir if args.with_prior_preservation else None,
|
||||
class_prompt=args.class_prompt,
|
||||
class_num=args.num_class_images,
|
||||
tokenizer=tokenizer,
|
||||
size=args.resolution,
|
||||
center_crop=args.center_crop,
|
||||
encoder_hidden_states=pre_computed_encoder_hidden_states,
|
||||
class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
|
||||
tokenizer_max_length=args.tokenizer_max_length,
|
||||
)
|
||||
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.train_batch_size,
|
||||
shuffle=True,
|
||||
collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
|
||||
num_workers=args.dataloader_num_workers,
|
||||
)
|
||||
```
|
||||
|
||||
最后,[训练循环](https://github.com/huggingface/diffusers/blob/072e00897a7cf4302c347a63ec917b4b8add16d4/examples/dreambooth/train_dreambooth.py#L1151)处理剩余步骤,例如将图像转换为潜在空间、向输入添加噪声、预测噪声残差和计算损失。
|
||||
|
||||
如果您想了解更多关于训练循环的工作原理,请查看[理解管道、模型和调度器](../using-diffusers/write_own_pipeline)教程,该教程分解了去噪过程的基本模式。
|
||||
|
||||
## 启动脚本
|
||||
|
||||
您现在准备好启动训练脚本了!🚀
|
||||
|
||||
对于本指南,您将下载一些[狗的图片](https://huggingface.co/datasets/diffusers/dog-example)的图像并将它们存储在一个目录中。但请记住,您可以根据需要创建和使用自己的数据集(请参阅[创建用于训练的数据集](create_dataset)指南)。
|
||||
|
||||
```py
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
local_dir = "./dog"
|
||||
snapshot_download(
|
||||
"diffusers/dog-example",
|
||||
local_dir=local_dir,
|
||||
repo_type="dataset",
|
||||
ignore_patterns=".gitattributes",
|
||||
)
|
||||
```
|
||||
|
||||
设置环境变量 `MODEL_NAME` 为 Hub 上的模型 ID 或本地模型路径,`INSTANCE_DIR` 为您刚刚下载狗图像的路径,`OUTPUT_DIR` 为您想保存模型的位置。您将使用 `sks` 作为特殊词来绑定训练。
|
||||
|
||||
如果您有兴趣跟随训练过程,可以定期保存生成的图像作为训练进度。将以下参数添加到训练命令中:
|
||||
|
||||
```bash
|
||||
--validation_prompt="a photo of a sks dog"
|
||||
--num_validation_images=4
|
||||
--validation_steps=100
|
||||
```
|
||||
|
||||
在启动脚本之前,还有一件事!根据您拥有的 GPU,您可能需要启用某些优化来训练 DreamBooth。
|
||||
|
||||
<hfoptions id="gpu-select">
|
||||
<hfoption id="16GB">
|
||||
|
||||
在 16GB GPU 上,您可以使用 bitsandbytes 8 位优化器和梯度检查点来帮助训练 DreamBooth 模型。安装 bitsandbytes:
|
||||
|
||||
```py
|
||||
pip install bitsandbytes
|
||||
```
|
||||
|
||||
然后,将以下参数添加到您的训练命令中:
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--gradient_checkpointing \
|
||||
--use_8bit_adam \
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="12GB">
|
||||
|
||||
在 12GB GPU 上,您需要 bitsandbytes 8 位优化器、梯度检查点、xFormers,并将梯度设置为 `None` 而不是零以减少内存使用。
|
||||
|
||||
```bash
|
||||
accelerate launch train_dreambooth.py \
|
||||
--use_8bit_adam \
|
||||
--gradient_checkpointing \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--set_grads_to_none \
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="8GB">
|
||||
|
||||
在 8GB GPU 上,您需要 [DeepSpeed](https://www.deepspeed.ai/) 将一些张量从 vRAM 卸载到 CPU 或 NVME,以便在更少的 GPU 内存下进行训练。
|
||||
|
||||
运行以下命令来配置您的 🤗 Accelerate 环境:
|
||||
|
||||
```bash
|
||||
accelerate config
|
||||
```
|
||||
|
||||
在配置过程中,确认您想使用 DeepSpeed。现在,通过结合 DeepSpeed 阶段 2、fp16 混合精度以及将模型参数和优化器状态卸载到 CPU,应该可以在低于 8GB vRAM 的情况下进行训练。缺点是这需要更多的系统 RAM(约 25 GB)。有关更多配置选项,请参阅 [DeepSpeed 文档](https://huggingface.co/docs/accelerate/usage_guides/deepspeed)。
|
||||
|
||||
您还应将默认的 Adam 优化器更改为 DeepSpeed 的优化版本 [`deepspeed.ops.adam.DeepSpeedCPUAdam`](https://deepspeed.readthedocs.io/en/latest/optimizers.html#adam-cpu) 以获得显著的速度提升。启用 `DeepSpeedCPUAdam` 要求您的系统 CUDA 工具链版本与 PyTorch 安装的版本相同。
|
||||
|
||||
目前,bitsandbytes 8 位优化器似乎与 DeepSpeed 不兼容。
|
||||
|
||||
就是这样!您不需要向训练命令添加任何额外参数。
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
<hfoptions id="training-inference">
|
||||
<hfoption id="PyTorch">
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
|
||||
export INSTANCE_DIR="./dog"
|
||||
export OUTPUT_DIR="path_to_
|
||||
saved_model"
|
||||
|
||||
accelerate launch train_dreambooth.py \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a photo of sks dog" \
|
||||
--resolution=512 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--learning_rate=5e-6 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--max_train_steps=400 \
|
||||
--push_to_hub
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Flax">
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="duongna/stable-diffusion-v1-4-flax"
|
||||
export INSTANCE_DIR="./dog"
|
||||
export OUTPUT_DIR="path-to-save-model"
|
||||
|
||||
python train_dreambooth_flax.py \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a photo of sks dog" \
|
||||
--resolution=512 \
|
||||
--train_batch_size=1 \
|
||||
--learning_rate=5e-6 \
|
||||
--max_train_steps=400 \
|
||||
--push_to_hub
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
训练完成后,您可以使用新训练的模型进行推理!
|
||||
|
||||
<Tip>
|
||||
|
||||
等不及在训练完成前就尝试您的模型进行推理?🤭 请确保安装了最新版本的 🤗 Accelerate。
|
||||
|
||||
```py
|
||||
from diffusers import DiffusionPipeline, UNet2DConditionModel
|
||||
from transformers import CLIPTextModel
|
||||
import torch
|
||||
|
||||
unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")
|
||||
|
||||
# 如果您使用了 `--args.train_text_encoder` 进行训练,请确保也加载文本编码器
|
||||
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained(
|
||||
"stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
|
||||
).to("cuda")
|
||||
|
||||
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
|
||||
image.save("dog-bucket.png")
|
||||
```
|
||||
|
||||
</Tip>
|
||||
|
||||
<hfoptions id="training-inference">
|
||||
<hfoption id="PyTorch">
|
||||
|
||||
```py
|
||||
from diffusers import DiffusionPipeline
|
||||
import torch
|
||||
|
||||
pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
|
||||
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
|
||||
image.save("dog-bucket.png")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Flax">
|
||||
|
||||
```py
|
||||
import jax
|
||||
import numpy as np
|
||||
from flax.jax_utils import replicate
|
||||
from flax.training.common_utils import shard
|
||||
from diffusers import FlaxStableDiffusionPipeline
|
||||
|
||||
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("path-to-your-trained-model", dtype=jax.numpy.bfloat16)
|
||||
|
||||
prompt = "A photo of sks dog in a bucket"
|
||||
prng_seed = jax.random.PRNGKey(0)
|
||||
num_inference_steps = 50
|
||||
|
||||
num_samples = jax.device_count()
|
||||
prompt = num_samples * [prompt]
|
||||
prompt_ids = pipeline.prepare_inputs(prompt)
|
||||
|
||||
# 分片输入和随机数生成器
|
||||
params = replicate(params)
|
||||
prng_seed = jax.random.split(prng_seed, jax.device_count())
|
||||
prompt_ids = shard(prompt_ids)
|
||||
|
||||
images = pipeline(prompt_ids, params, prng_seed, num_inference_
|
||||
steps, jit=True).images
|
||||
images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
|
||||
image.save("dog-bucket.png")
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## LoRA
|
||||
|
||||
LoRA 是一种训练技术,可显著减少可训练参数的数量。因此,训练速度更快,并且更容易存储生成的权重,因为它们小得多(约 100MB)。使用 [train_dreambooth_lora.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py) 脚本通过 LoRA 进行训练。
|
||||
|
||||
LoRA 训练脚本在 [LoRA 训练](lora) 指南中有更详细的讨论。
|
||||
|
||||
## Stable Diffusion XL
|
||||
|
||||
Stable Diffusion XL (SDXL) 是一个强大的文本到图像模型,可生成高分辨率图像,并在其架构中添加了第二个文本编码器。使用 [train_dreambooth_lora_sdxl.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py) 脚本通过 LoRA 训练 SDXL 模型。
|
||||
|
||||
SDXL 训练脚本在 [SDXL 训练](sdxl) 指南中有更详细的讨论。
|
||||
|
||||
## DeepFloyd IF
|
||||
|
||||
DeepFloyd IF 是一个级联像素扩散模型,包含三个阶段。第一阶段生成基础图像,第二和第三阶段逐步将基础图像放大为高分辨率 1024x1024 图像。使用 [train_dreambooth_lora.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py) 或 [train_dreambooth.py](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) 脚本通过 LoRA 或完整模型训练 DeepFloyd IF 模型。
|
||||
|
||||
DeepFloyd IF 使用预测方差,但 Diffusers 训练脚本使用预测误差,因此训练的 DeepFloyd IF 模型被切换到固定方差调度。训练脚本将为您更新完全训练模型的调度器配置。但是,当您加载保存的 LoRA 权重时,还必须更新管道的调度器配置。
|
||||
|
||||
```py
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)
|
||||
|
||||
pipe.load_lora_weights("<lora weights path>")
|
||||
|
||||
# 更新调度器配置为固定方差调度
|
||||
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")
|
||||
```
|
||||
|
||||
第二阶段模型需要额外的验证图像进行放大。您可以下载并使用训练图像的缩小版本。
|
||||
|
||||
```py
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
local_dir = "./dog_downsized"
|
||||
snapshot_download(
|
||||
"diffusers/dog-example-downsized",
|
||||
local_dir=local_dir,
|
||||
repo_type="dataset",
|
||||
ignore_patterns=".gitattributes",
|
||||
)
|
||||
```
|
||||
|
||||
以下代码示例简要概述了如何结合 DreamBooth 和 LoRA 训练 DeepFloyd IF 模型。一些需要注意的重要参数包括:
|
||||
|
||||
* `--resolution=64`,需要更小的分辨率,因为 DeepFloyd IF 是
|
||||
一个像素扩散模型,用于处理未压缩的像素,输入图像必须更小
|
||||
* `--pre_compute_text_embeddings`,提前计算文本嵌入以节省内存,因为 [`~transformers.T5Model`] 可能占用大量内存
|
||||
* `--tokenizer_max_length=77`,您可以使用更长的默认文本长度与 T5 作为文本编码器,但默认模型编码过程使用较短的文本长度
|
||||
* `--text_encoder_use_attention_mask`,将注意力掩码传递给文本编码器
|
||||
|
||||
<hfoptions id="IF-DreamBooth">
|
||||
<hfoption id="Stage 1 LoRA DreamBooth">
|
||||
|
||||
使用 LoRA 和 DreamBooth 训练 DeepFloyd IF 的第 1 阶段需要约 28GB 内存。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
|
||||
export INSTANCE_DIR="dog"
|
||||
export OUTPUT_DIR="dreambooth_dog_lora"
|
||||
|
||||
accelerate launch train_dreambooth_lora.py \
|
||||
--report_to wandb \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a sks dog" \
|
||||
--resolution=64 \
|
||||
--train_batch_size=4 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--learning_rate=5e-6 \
|
||||
--scale_lr \
|
||||
--max_train_steps=1200 \
|
||||
--validation_prompt="a sks dog" \
|
||||
--validation_epochs=25 \
|
||||
--checkpointing_steps=100 \
|
||||
--pre_compute_text_embeddings \
|
||||
--tokenizer_max_length=77 \
|
||||
--text_encoder_use_attention_mask
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Stage 2 LoRA DreamBooth">
|
||||
|
||||
对于使用 LoRA 和 DreamBooth 的 DeepFloyd IF 第 2 阶段,请注意这些参数:
|
||||
|
||||
* `--validation_images`,验证期间用于上采样的图像
|
||||
* `--class_labels_conditioning=timesteps`,根据需要额外条件化 UNet,如第 2 阶段中所需
|
||||
* `--learning_rate=1e-6`,与第 1 阶段相比使用较低的学习率
|
||||
* `--resolution=256`,上采样器的预期分辨率
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="DeepFloyd/IF-II-L-v1.0"
|
||||
export INSTANCE_DIR="dog"
|
||||
export OUTPUT_DIR="dreambooth_dog_upscale"
|
||||
export VALIDATION_IMAGES="dog_downsized/image_1.png dog_downsized/image_2.png dog_downsized/image_3.png dog_downsized/image_4.png"
|
||||
|
||||
python train_dreambooth_lora.py \
|
||||
--report_to wandb \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a sks dog" \
|
||||
--resolution=256 \
|
||||
--train_batch_size=4 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--learning_rate=1e-6 \
|
||||
--max_train_steps=2000 \
|
||||
--validation_prompt="a sks dog" \
|
||||
--validation_epochs=100 \
|
||||
--checkpointing_steps=500 \
|
||||
--pre_compute_text_embeddings \
|
||||
--tokenizer_max_length=77 \
|
||||
--text_encoder_use_attention_mask \
|
||||
--validation_images $VALIDATION_IMAGES \
|
||||
--class_labels_conditioning=timesteps
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Stage 1 DreamBooth">
|
||||
|
||||
对于使用 DreamBooth 的 DeepFloyd IF 第 1 阶段,请注意这些参数:
|
||||
|
||||
* `--skip_save_text_encoder`,跳过保存完整 T5 文本编码器与微调模型
|
||||
* `--use_8bit_adam`,使用 8 位 Adam 优化器以节省内存,因为
|
||||
|
||||
优化器状态的大小在训练完整模型时
|
||||
* `--learning_rate=1e-7`,对于完整模型训练应使用非常低的学习率,否则模型质量会下降(您可以使用更高的学习率和更大的批次大小)
|
||||
|
||||
使用8位Adam和批次大小为4进行训练,完整模型可以在约48GB内存下训练。
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
|
||||
export INSTANCE_DIR="dog"
|
||||
export OUTPUT_DIR="dreambooth_if"
|
||||
|
||||
accelerate launch train_dreambooth.py \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a photo of sks dog" \
|
||||
--resolution=64 \
|
||||
--train_batch_size=4 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--learning_rate=1e-7 \
|
||||
--max_train_steps=150 \
|
||||
--validation_prompt "a photo of sks dog" \
|
||||
--validation_steps 25 \
|
||||
--text_encoder_use_attention_mask \
|
||||
--tokenizer_max_length 77 \
|
||||
--pre_compute_text_embeddings \
|
||||
--use_8bit_adam \
|
||||
--set_grads_to_none \
|
||||
--skip_save_text_encoder \
|
||||
--push_to_hub
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="Stage 2 DreamBooth">
|
||||
|
||||
对于DeepFloyd IF的第二阶段DreamBooth,请注意这些参数:
|
||||
|
||||
* `--learning_rate=5e-6`,使用较低的学习率和较小的有效批次大小
|
||||
* `--resolution=256`,上采样器的预期分辨率
|
||||
* `--train_batch_size=2` 和 `--gradient_accumulation_steps=6`,为了有效训练包含面部的图像,需要更大的批次大小
|
||||
|
||||
```bash
|
||||
export MODEL_NAME="DeepFloyd/IF-II-L-v1.0"
|
||||
export INSTANCE_DIR="dog"
|
||||
export OUTPUT_DIR="dreambooth_dog_upscale"
|
||||
export VALIDATION_IMAGES="dog_downsized/image_1.png dog_downsized/image_2.png dog_downsized/image_3.png dog_downsized/image_4.png"
|
||||
|
||||
accelerate launch train_dreambooth.py \
|
||||
--report_to wandb \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--instance_data_dir=$INSTANCE_DIR \
|
||||
--output_dir=$OUTPUT_DIR \
|
||||
--instance_prompt="a sks dog" \
|
||||
--resolution=256 \
|
||||
--train_batch_size=2 \
|
||||
--gradient_accumulation_steps=6 \
|
||||
--learning_rate=5e-6 \
|
||||
--max_train_steps=2000 \
|
||||
--validation_prompt="a sks dog" \
|
||||
--validation_steps=150 \
|
||||
--checkpointing_steps=500 \
|
||||
--pre_compute_text_embeddings \
|
||||
--tokenizer_max_length=77 \
|
||||
--text_encoder_use_attention_mask \
|
||||
--validation_images $VALIDATION_IMAGES \
|
||||
--class_labels_conditioning timesteps \
|
||||
--push_to_hub
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
### 训练技巧
|
||||
|
||||
训练DeepFloyd IF模型可能具有挑战性,但以下是我们发现有用的技巧:
|
||||
|
||||
- LoRA对于训练第一阶段模型已足够,因为模型的低分辨率使得表示更精细的细节变得困难,无论如何。
|
||||
- 对于常见或简单的对象,您不一定需要微调上采样器。确保传递给上采样器的提示被调整以移除实例提示中的新令牌。例如,如果您第一阶段提示是"a sks dog",那么您第二阶段的提示应该是"a dog"。
|
||||
- 对于更精细的细节,如面部,完全训练
|
||||
使用阶段2上采样器比使用LoRA训练阶段2模型更好。使用更大的批次大小和较低的学习率也有帮助。
|
||||
- 应使用较低的学习率来训练阶段2模型。
|
||||
- [`DDPMScheduler`] 比训练脚本中使用的DPMSolver效果更好。
|
||||
|
||||
## 下一步
|
||||
|
||||
恭喜您训练了您的DreamBooth模型!要了解更多关于如何使用您的新模型的信息,以下指南可能有所帮助:
|
||||
- 如果您使用LoRA训练了您的模型,请学习如何[加载DreamBooth](../using-diffusers/loading_adapters)模型进行推理。
|
||||
@@ -1,255 +0,0 @@
|
||||
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
-->
|
||||
|
||||
# InstructPix2Pix
|
||||
|
||||
[InstructPix2Pix](https://hf.co/papers/2211.09800) 是一个基于 Stable Diffusion 训练的模型,用于根据人类提供的指令编辑图像。例如,您的提示可以是“将云变成雨天”,模型将相应编辑输入图像。该模型以文本提示(或编辑指令)和输入图像为条件。
|
||||
|
||||
本指南将探索 [train_instruct_pix2pix.py](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py) 训练脚本,帮助您熟悉它,以及如何将其适应您自己的用例。
|
||||
|
||||
在运行脚本之前,请确保从源代码安装库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/huggingface/diffusers
|
||||
cd diffusers
|
||||
pip install .
|
||||
```
|
||||
|
||||
然后导航到包含训练脚本的示例文件夹,并安装脚本所需的依赖项:
|
||||
|
||||
```bash
|
||||
cd examples/instruct_pix2pix
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
🤗 Accelerate 是一个库,用于帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速导览](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。
|
||||
|
||||
</Tip>
|
||||
|
||||
初始化一个 🤗 Accelerate 环境:
|
||||
|
||||
```bash
|
||||
accelerate config
|
||||
```
|
||||
|
||||
要设置一个默认的 🤗 Accelerate 环境,无需选择任何配置:
|
||||
|
||||
```bash
|
||||
accelerate config default
|
||||
```
|
||||
|
||||
或者,如果您的环境不支持交互式 shell,例如笔记本,您可以使用:
|
||||
|
||||
```py
|
||||
from accelerate.utils import write_basic_config
|
||||
|
||||
write_basic_config()
|
||||
```
|
||||
|
||||
最后,如果您想在自己的数据集上训练模型,请查看 [创建用于训练的数据集](create_dataset) 指南,了解如何创建与训练脚本兼容的数据集。
|
||||
|
||||
<Tip>
|
||||
|
||||
以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分,但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多,请随时阅读 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py),并告诉我们如果您有任何问题或疑虑。
|
||||
|
||||
</Tip>
|
||||
|
||||
## 脚本参数
|
||||
|
||||
训练脚本有许多参数可帮助您自定义训练运行。所有
|
||||
参数及其描述可在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L65) 函数中找到。大多数参数都提供了默认值,这些值效果相当不错,但如果您愿意,也可以在训练命令中设置自己的值。
|
||||
|
||||
例如,要增加输入图像的分辨率:
|
||||
|
||||
```bash
|
||||
accelerate launch train_instruct_pix2pix.py \
|
||||
--resolution=512 \
|
||||
```
|
||||
|
||||
许多基本和重要的参数在 [文本到图像](text2image#script-parameters) 训练指南中已有描述,因此本指南仅关注与 InstructPix2Pix 相关的参数:
|
||||
|
||||
- `--original_image_column`:编辑前的原始图像
|
||||
- `--edited_image_column`:编辑后的图像
|
||||
- `--edit_prompt_column`:编辑图像的指令
|
||||
- `--conditioning_dropout_prob`:训练期间编辑图像和编辑提示的 dropout 概率,这为一种或两种条件输入启用了无分类器引导(CFG)
|
||||
|
||||
## 训练脚本
|
||||
|
||||
数据集预处理代码和训练循环可在 [`main()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L374) 函数中找到。这是您将修改训练脚本以适应自己用例的地方。
|
||||
|
||||
与脚本参数类似,[文本到图像](text2image#training-script) 训练指南提供了训练脚本的逐步说明。相反,本指南将查看脚本中与 InstructPix2Pix 相关的部分。
|
||||
|
||||
脚本首先修改 UNet 的第一个卷积层中的 [输入通道数](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L445),以适应 InstructPix2Pix 的额外条件图像:
|
||||
|
||||
```py
|
||||
in_channels = 8
|
||||
out_channels = unet.conv_in.out_channels
|
||||
unet.register_to_config(in_channels=in_channels)
|
||||
|
||||
with torch.no_grad():
|
||||
new_conv_in = nn.Conv2d(
|
||||
in_channels, out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding
|
||||
)
|
||||
new_conv_in.weight.zero_()
|
||||
new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
|
||||
unet.conv_in = new_conv_in
|
||||
```
|
||||
|
||||
这些 UNet 参数由优化器 [更新](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L545C1-L551C6):
|
||||
|
||||
```py
|
||||
optimizer = optimizer_cls(
|
||||
unet.parameters(),
|
||||
lr=args.learning_rate,
|
||||
betas=(args.adam_beta1, args.adam_beta2),
|
||||
weight_decay=args.adam_weight_decay,
|
||||
eps=args.adam_epsilon,
|
||||
)
|
||||
```
|
||||
|
||||
接下来,编辑后的图像和编辑指令被 [预处理](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L624)并被[tokenized](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L610C24-L610C24)。重要的是,对原始图像和编辑后的图像应用相同的图像变换。
|
||||
|
||||
```py
|
||||
def preprocess_train(examples):
|
||||
preprocessed_images = preprocess_images(examples)
|
||||
|
||||
original_images, edited_images = preprocessed_images.chunk(2)
|
||||
original_images = original_images.reshape(-1, 3, args.resolution, args.resolution)
|
||||
edited_images = edited_images.reshape(-1, 3, args.resolution, args.resolution)
|
||||
|
||||
examples["original_pixel_values"] = original_images
|
||||
examples["edited_pixel_values"] = edited_images
|
||||
|
||||
captions = list(examples[edit_prompt_column])
|
||||
examples["input_ids"] = tokenize_captions(captions)
|
||||
return examples
|
||||
```
|
||||
|
||||
最后,在[训练循环](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L730)中,它首先将编辑后的图像编码到潜在空间:
|
||||
|
||||
```py
|
||||
latents = vae.encode(batch["edited_pixel_values"].to(weight_dtype)).latent_dist.sample()
|
||||
latents = latents * vae.config.scaling_factor
|
||||
```
|
||||
|
||||
然后,脚本对原始图像和编辑指令嵌入应用 dropout 以支持 CFG(Classifier-Free Guidance)。这使得模型能够调节编辑指令和原始图像对编辑后图像的影响。
|
||||
|
||||
```py
|
||||
encoder_hidden_states = text_encoder(batch["input_ids"])[0]
|
||||
original_image_embeds = vae.encode(batch["original_pixel_values"].to(weight_dtype)).latent_dist.mode()
|
||||
|
||||
if args.conditioning_dropout_prob is not None:
|
||||
random_p = torch.rand(bsz, device=latents.device, generator=generator)
|
||||
prompt_mask = random_p < 2 * args.conditioning_dropout_prob
|
||||
prompt_mask = prompt_mask.reshape(bsz, 1, 1)
|
||||
null_conditioning = text_encoder(tokenize_captions([""]).to(accelerator.device))[0]
|
||||
encoder_hidden_states = torch.where(prompt_mask, null_conditioning, encoder_hidden_states)
|
||||
|
||||
image_mask_dtype = original_image_embeds.dtype
|
||||
image_mask = 1 - (
|
||||
(random_p >= args.conditioning_dropout_prob).to(image_mask_dtype)
|
||||
* (random_p < 3 * args.conditioning_dropout_prob).to(image_mask_dtype)
|
||||
)
|
||||
image_mask = image_mask.reshape(bsz, 1, 1, 1)
|
||||
original_image_embeds = image_mask * original_image_embeds
|
||||
```
|
||||
|
||||
差不多就是这样了!除了这里描述的不同之处,脚本的其余部分与[文本到图像](text2image#training-script)训练脚本非常相似,所以请随意查看以获取更多细节。如果您想了解更多关于训练循环如何工作的信息,请查看[理解管道、模型和调度器](../using-diffusers/write_own_pipeline)教程,该教程分解了去噪过程的基本模式。
|
||||
|
||||
## 启动脚本
|
||||
|
||||
一旦您对脚本的更改感到满意,或者如果您对默认配置没问题,您
|
||||
准备好启动训练脚本!🚀
|
||||
|
||||
本指南使用 [fusing/instructpix2pix-1000-samples](https://huggingface.co/datasets/fusing/instructpix2pix-1000-samples) 数据集,这是 [原始数据集](https://huggingface.co/datasets/timbrooks/instructpix2pix-clip-filtered) 的一个较小版本。您也可以创建并使用自己的数据集(请参阅 [创建用于训练的数据集](create_dataset) 指南)。
|
||||
|
||||
将 `MODEL_NAME` 环境变量设置为模型名称(可以是 Hub 上的模型 ID 或本地模型的路径),并将 `DATASET_ID` 设置为 Hub 上数据集的名称。脚本会创建并保存所有组件(特征提取器、调度器、文本编码器、UNet 等)到您的仓库中的一个子文件夹。
|
||||
|
||||
<Tip>
|
||||
|
||||
为了获得更好的结果,尝试使用更大的数据集进行更长时间的训练。我们只在较小规模的数据集上测试过此训练脚本。
|
||||
|
||||
<br>
|
||||
|
||||
要使用 Weights and Biases 监控训练进度,请将 `--report_to=wandb` 参数添加到训练命令中,并使用 `--val_image_url` 指定验证图像,使用 `--validation_prompt` 指定验证提示。这对于调试模型非常有用。
|
||||
|
||||
</Tip>
|
||||
|
||||
如果您在多个 GPU 上训练,请将 `--multi_gpu` 参数添加到 `accelerate launch` 命令中。
|
||||
|
||||
```bash
|
||||
accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py \
|
||||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||||
--dataset_name=$DATASET_ID \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resolution=256 \
|
||||
--random_flip \
|
||||
--train_batch_size=4 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--max_train_steps=15000 \
|
||||
--checkpointing_steps=5000 \
|
||||
--checkpoints_total_limit=1 \
|
||||
--learning_rate=5e-05 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_warmup_steps=0 \
|
||||
--conditioning_dropout_prob=0.05 \
|
||||
--mixed_precision=fp16 \
|
||||
--seed=42 \
|
||||
--push_to_hub
|
||||
```
|
||||
|
||||
训练完成后,您可以使用您的新 InstructPix2Pix 进行推理:
|
||||
|
||||
```py
|
||||
import PIL
|
||||
import requests
|
||||
import torch
|
||||
from diffusers import StableDiffusionInstructPix2PixPipeline
|
||||
from diffusers.utils import load_image
|
||||
|
||||
pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained("your_cool_model", torch_dtype=torch.float16).to("cuda")
|
||||
generator = torch.Generator("cuda").manual_seed(0)
|
||||
|
||||
image = load_image("https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/test_pix2pix_4.png")
|
||||
prompt = "add some ducks to the lake"
|
||||
num_inference_steps = 20
|
||||
image_guidance_scale = 1.5
|
||||
guidance_scale = 10
|
||||
|
||||
edited_image = pipeline(
|
||||
prompt,
|
||||
image=image,
|
||||
num_inference_steps=num_inference_steps,
|
||||
image_guidance_scale=image_guidance_scale,
|
||||
guidance_scale=guidance_scale,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
edited_image.save("edited_image.png")
|
||||
```
|
||||
|
||||
您应该尝试不同的 `num_inference_steps`、`image_guidance_scale` 和 `guidance_scale` 值,以查看它们如何影响推理速度和质量。指导比例参数
|
||||
这些参数尤其重要,因为它们控制原始图像和编辑指令对编辑后图像的影响程度。
|
||||
|
||||
## Stable Diffusion XL
|
||||
|
||||
Stable Diffusion XL (SDXL) 是一个强大的文本到图像模型,能够生成高分辨率图像,并在其架构中添加了第二个文本编码器。使用 [`train_instruct_pix2pix_sdxl.py`](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py) 脚本来训练 SDXL 模型以遵循图像编辑指令。
|
||||
|
||||
SDXL 训练脚本在 [SDXL 训练](sdxl) 指南中有更详细的讨论。
|
||||
|
||||
## 后续步骤
|
||||
|
||||
恭喜您训练了自己的 InstructPix2Pix 模型!🥳 要了解更多关于该模型的信息,可能有助于:
|
||||
|
||||
- 阅读 [Instruction-tuning Stable Diffusion with InstructPix2Pix](https://huggingface.co/blog/instruction-tuning-sd) 博客文章,了解更多我们使用 InstructPix2Pix 进行的一些实验、数据集准备以及不同指令的结果。
|
||||
@@ -1,328 +0,0 @@
|
||||
<!--版权所有 2025 HuggingFace 团队。保留所有权利。
|
||||
|
||||
根据 Apache 许可证 2.0 版本("许可证")授权;除非遵守许可证,否则您不得使用此文件。您可以在以下网址获取许可证副本:
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
除非适用法律要求或书面同意,否则根据许可证分发的软件按"原样"分发,不附带任何明示或暗示的担保或条件。请参阅许可证以了解具体的语言管理权限和限制。
|
||||
-->
|
||||
|
||||
# Kandinsky 2.2
|
||||
|
||||
<Tip warning={true}>
|
||||
|
||||
此脚本是实验性的,容易过拟合并遇到灾难性遗忘等问题。尝试探索不同的超参数以在您的数据集上获得最佳结果。
|
||||
|
||||
</Tip>
|
||||
|
||||
Kandinsky 2.2 是一个多语言文本到图像模型,能够生成更逼真的图像。该模型包括一个图像先验模型,用于从文本提示创建图像嵌入,以及一个解码器模型,基于先验模型的嵌入生成图像。这就是为什么在 Diffusers 中您会找到两个独立的脚本用于 Kandinsky 2.2,一个用于训练先验模型,另一个用于训练解码器模型。您可以分别训练这两个模型,但为了获得最佳结果,您应该同时训练先验和解码器模型。
|
||||
|
||||
根据您的 GPU,您可能需要启用 `gradient_checkpointing`(⚠️ 不支持先验模型!)、`mixed_precision` 和 `gradient_accumulation_steps` 来帮助将模型装入内存并加速训练。您可以通过启用 [xFormers](../optimization/xformers) 的内存高效注意力来进一步减少内存使用(版本 [v0.0.16](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212) 在某些 GPU 上训练时失败,因此您可能需要安装开发版本)。
|
||||
|
||||
本指南探讨了 [train_text_to_image_prior.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py) 和 [train_text_to_image_decoder.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py) 脚本,以帮助您更熟悉它,以及如何根据您的用例进行调整。
|
||||
|
||||
在运行脚本之前,请确保从源代码安装库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/huggingface/diffusers
|
||||
cd diffusers
|
||||
pip install .
|
||||
```
|
||||
|
||||
然后导航到包含训练脚本的示例文件夹,并安装脚本所需的依赖项:
|
||||
|
||||
```bash
|
||||
cd examples/kandinsky2_2/text_to_image
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate 的 [快速入门](https://huggingface.co/docs/accelerate/quicktour
|
||||
) 了解更多。
|
||||
|
||||
</Tip>
|
||||
|
||||
初始化一个 🤗 Accelerate 环境:
|
||||
|
||||
```bash
|
||||
accelerate config
|
||||
```
|
||||
|
||||
要设置一个默认的 🤗 Accelerate 环境而不选择任何配置:
|
||||
|
||||
```bash
|
||||
accelerate config default
|
||||
```
|
||||
|
||||
或者,如果您的环境不支持交互式 shell,比如 notebook,您可以使用:
|
||||
|
||||
```py
|
||||
from accelerate.utils import write_basic_config
|
||||
|
||||
write_basic_config()
|
||||
```
|
||||
|
||||
最后,如果您想在自己的数据集上训练模型,请查看 [创建用于训练的数据集](create_dataset) 指南,了解如何创建与训练脚本兼容的数据集。
|
||||
|
||||
<Tip>
|
||||
|
||||
以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分,但并未详细涵盖脚本的每个方面。如果您有兴趣了解更多,请随时阅读脚本,并让我们知道您有任何疑问或顾虑。
|
||||
|
||||
</Tip>
|
||||
|
||||
## 脚本参数
|
||||
|
||||
训练脚本提供了许多参数来帮助您自定义训练运行。所有参数及其描述都可以在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L190) 函数中找到。训练脚本为每个参数提供了默认值,例如训练批次大小和学习率,但如果您愿意,也可以在训练命令中设置自己的值。
|
||||
|
||||
例如,要使用 fp16 格式的混合精度加速训练,请在训练命令中添加 `--mixed_precision` 参数:
|
||||
|
||||
```bash
|
||||
accelerate launch train_text_to_image_prior.py \
|
||||
--mixed_precision="fp16"
|
||||
```
|
||||
|
||||
大多数参数与 [文本到图像](text2image#script-parameters) 训练指南中的参数相同,所以让我们直接进入 Kandinsky 训练脚本的 walkthrough!
|
||||
|
||||
### Min-SNR 加权
|
||||
|
||||
[Min-SNR](https://huggingface.co/papers/2303.09556) 加权策略可以通过重新平衡损失来帮助训练,实现更快的收敛。训练脚本支持预测 `epsilon`(噪声)或 `v_prediction`,但 Min-SNR 与两种预测类型都兼容。此加权策略仅由 PyTorch 支持,在 Flax 训练脚本中不可用。
|
||||
|
||||
添加 `--snr_gamma` 参数并将其设置为推荐值 5.0:
|
||||
|
||||
```bash
|
||||
accelerate launch train_text_to_image_prior.py \
|
||||
--snr_gamma=5.0
|
||||
```
|
||||
|
||||
## 训练脚本
|
||||
|
||||
训练脚本也类似于 [文本到图像](text2image#training-script) 训练指南,但已修改以支持训练 prior 和 decoder 模型。本指南重点介绍 Kandinsky 2.2 训练脚本中独特的代码。
|
||||
|
||||
<hfoptions id="script">
|
||||
<hfoption id="prior model">
|
||||
|
||||
[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L441) 函数包含代码 f
|
||||
或准备数据集和训练模型。
|
||||
|
||||
您会立即注意到的主要区别之一是,训练脚本除了调度器和分词器外,还加载了一个 [`~transformers.CLIPImageProcessor`] 用于预处理图像,以及一个 [`~transformers.CLIPVisionModelWithProjection`] 模型用于编码图像:
|
||||
|
||||
```py
|
||||
noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
|
||||
image_processor = CLIPImageProcessor.from_pretrained(
|
||||
args.pretrained_prior_model_name_or_path, subfolder="image_processor"
|
||||
)
|
||||
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")
|
||||
|
||||
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
args.pretrained_prior_model_name_or_path, subfolder="image_encoder", torch_dtype=weight_dtype
|
||||
).eval()
|
||||
text_encoder = CLIPTextModelWithProjection.from_pretrained(
|
||||
args.pretrained_prior_model_name_or_path, subfolder="text_encoder", torch_dtype=weight_dtype
|
||||
).eval()
|
||||
```
|
||||
|
||||
Kandinsky 使用一个 [`PriorTransformer`] 来生成图像嵌入,因此您需要设置优化器来学习先验模型的参数。
|
||||
|
||||
```py
|
||||
prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
|
||||
prior.train()
|
||||
optimizer = optimizer_cls(
|
||||
prior.parameters(),
|
||||
lr=args.learning_rate,
|
||||
betas=(args.adam_beta1, args.adam_beta2),
|
||||
weight_decay=args.adam_weight_decay,
|
||||
eps=args.adam_epsilon,
|
||||
)
|
||||
```
|
||||
|
||||
接下来,输入标题被分词,图像由 [`~transformers.CLIPImageProcessor`] [预处理](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L632):
|
||||
|
||||
```py
|
||||
def preprocess_train(examples):
|
||||
images = [image.convert("RGB") for image in examples[image_column]]
|
||||
examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
|
||||
examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
|
||||
return examples
|
||||
```
|
||||
|
||||
最后,[训练循环](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L718) 将输入图像转换为潜在表示,向图像嵌入添加噪声,并进行预测:
|
||||
|
||||
```py
|
||||
model_pred = prior(
|
||||
noisy_latents,
|
||||
timestep=timesteps,
|
||||
proj_embedding=prompt_embeds,
|
||||
encoder_hidden_states=text_encoder_hidden_states,
|
||||
attention_mask=text_mask,
|
||||
).predicted_image_embedding
|
||||
```
|
||||
|
||||
如果您想了解更多关于训练循环的工作原理,请查看 [理解管道、模型和调度器](../using-diffusers/write_own_pipeline) 教程,该教程分解了去噪过程的基本模式。
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="decoder model">
|
||||
|
||||
The [`main()`](https://github.com/huggingface/di
|
||||
ffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L440) 函数包含准备数据集和训练模型的代码。
|
||||
|
||||
与之前的模型不同,解码器初始化一个 [`VQModel`] 来将潜在变量解码为图像,并使用一个 [`UNet2DConditionModel`]:
|
||||
|
||||
```py
|
||||
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
|
||||
vae = VQModel.from_pretrained(
|
||||
args.pretrained_decoder_model_name_or_path, subfolder="movq", torch_dtype=weight_dtype
|
||||
).eval()
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
args.pretrained_prior_model_name_or_path, subfolder="image_encoder", torch_dtype=weight_dtype
|
||||
).eval()
|
||||
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")
|
||||
```
|
||||
|
||||
接下来,脚本包括几个图像变换和一个用于对图像应用变换并返回像素值的[预处理](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L622)函数:
|
||||
|
||||
```py
|
||||
def preprocess_train(examples):
|
||||
images = [image.convert("RGB") for image in examples[image_column]]
|
||||
examples["pixel_values"] = [train_transforms(image) for image in images]
|
||||
examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
|
||||
return examples
|
||||
```
|
||||
|
||||
最后,[训练循环](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L706)处理将图像转换为潜在变量、添加噪声和预测噪声残差。
|
||||
|
||||
如果您想了解更多关于训练循环如何工作的信息,请查看[理解管道、模型和调度器](../using-diffusers/write_own_pipeline)教程,该教程分解了去噪过程的基本模式。
|
||||
|
||||
```py
|
||||
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 启动脚本
|
||||
|
||||
一旦您完成了所有更改或接受默认配置,就可以启动训练脚本了!🚀
|
||||
|
||||
您将在[Naruto BLIP 字幕](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions)数据集上进行训练,以生成您自己的Naruto角色,但您也可以通过遵循[创建用于训练的数据集](create_dataset)指南来创建和训练您自己的数据集。将环境变量 `DATASET_NAME` 设置为Hub上数据集的名称,或者如果您在自己的文件上训练,将环境变量 `TRAIN_DIR` 设置为数据集的路径。
|
||||
|
||||
如果您在多个GPU上训练,请在 `accelerate launch` 命令中添加 `--multi_gpu` 参数。
|
||||
|
||||
<Tip>
|
||||
|
||||
要使用Weights & Biases监控训练进度,请在训练命令中添加 `--report_to=wandb` 参数。您还需要
|
||||
建议在训练命令中添加 `--validation_prompt` 以跟踪结果。这对于调试模型和查看中间结果非常有用。
|
||||
|
||||
</Tip>
|
||||
|
||||
<hfoptions id="training-inference">
|
||||
<hfoption id="prior model">
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="lambdalabs/naruto-blip-captions"
|
||||
|
||||
accelerate launch --mixed_precision="fp16" train_text_to_image_prior.py \
|
||||
--dataset_name=$DATASET_NAME \
|
||||
--resolution=768 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--max_train_steps=15000 \
|
||||
--learning_rate=1e-05 \
|
||||
--max_grad_norm=1 \
|
||||
--checkpoints_total_limit=3 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--validation_prompts="A robot naruto, 4k photo" \
|
||||
--report_to="wandb" \
|
||||
--push_to_hub \
|
||||
--output_dir="kandi2-prior-naruto-model"
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="decoder model">
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="lambdalabs/naruto-blip-captions"
|
||||
|
||||
accelerate launch --mixed_precision="fp16" train_text_to_image_decoder.py \
|
||||
--dataset_name=$DATASET_NAME \
|
||||
--resolution=768 \
|
||||
--train_batch_size=1 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--max_train_steps=15000 \
|
||||
--learning_rate=1e-05 \
|
||||
--max_grad_norm=1 \
|
||||
--checkpoints_total_limit=3 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--validation_prompts="A robot naruto, 4k photo" \
|
||||
--report_to="wandb" \
|
||||
--push_to_hub \
|
||||
--output_dir="kandi2-decoder-naruto-model"
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
训练完成后,您可以使用新训练的模型进行推理!
|
||||
|
||||
<hfoptions id="training-inference">
|
||||
<hfoption id="prior model">
|
||||
|
||||
```py
|
||||
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
|
||||
import torch
|
||||
|
||||
prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, torch_dtype=torch.float16)
|
||||
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
|
||||
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, torch_dtype=torch.float16)
|
||||
|
||||
pipe.enable_model_cpu_offload()
|
||||
prompt="A robot naruto, 4k photo"
|
||||
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
可以随意将 `kandinsky-community/kandinsky-2-2-decoder` 替换为您自己训练的 decoder 检查点!
|
||||
|
||||
</Tip>
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="decoder model">
|
||||
|
||||
```py
|
||||
from diffusers import AutoPipelineForText2Image
|
||||
import torch
|
||||
|
||||
pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", torch_dtype=torch.float16)
|
||||
pipeline.enable_model_cpu_offload()
|
||||
|
||||
prompt="A robot naruto, 4k photo"
|
||||
image = pipeline(prompt=prompt).images[0]
|
||||
```
|
||||
|
||||
对于 decoder 模型,您还可以从保存的检查点进行推理,这对于查看中间结果很有用。在这种情况下,将检查点加载到 UNet 中:
|
||||
|
||||
```py
|
||||
from diffusers import AutoPipelineForText2Image, UNet2DConditionModel
|
||||
|
||||
unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")
|
||||
|
||||
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, torch_dtype=torch.float16)
|
||||
pipeline.enable_model_cpu_offload()
|
||||
|
||||
image = pipeline(prompt="A robot naruto, 4k photo").images[0]
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
</hfoptions>
|
||||
|
||||
## 后续步骤
|
||||
|
||||
恭喜您训练了一个 Kandinsky 2.2 模型!要了解更多关于如何使用您的新模型的信息,以下指南可能会有所帮助:
|
||||
|
||||
- 阅读 [Kandinsky](../using-diffusers/kandinsky) 指南,学习如何将其用于各种不同的任务(文本到图像、图像到图像、修复、插值),以及如何与 ControlNet 结合使用。
|
||||
- 查看 [DreamBooth](dreambooth) 和 [LoRA](lora) 训练指南,学习如何使用少量示例图像训练个性化的 Kandinsky 模型。这两种训练技术甚至可以结合使用!
|
||||
@@ -1,191 +0,0 @@
|
||||
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
-->
|
||||
|
||||
# Wuerstchen
|
||||
|
||||
[Wuerstchen](https://hf.co/papers/2306.00637) 模型通过将潜在空间压缩 42 倍,在不影响图像质量的情况下大幅降低计算成本并加速推理。在训练过程中,Wuerstchen 使用两个模型(VQGAN + 自动编码器)来压缩潜在表示,然后第三个模型(文本条件潜在扩散模型)在这个高度压缩的空间上进行条件化以生成图像。
|
||||
|
||||
为了将先验模型放入 GPU 内存并加速训练,尝试分别启用 `gradient_accumulation_steps`、`gradient_checkpointing` 和 `mixed_precision`。
|
||||
|
||||
本指南探讨 [train_text_to_image_prior.py](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) 脚本,帮助您更熟悉它,以及如何根据您的用例进行适配。
|
||||
|
||||
在运行脚本之前,请确保从源代码安装库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/huggingface/diffusers
|
||||
cd diffusers
|
||||
pip install .
|
||||
```
|
||||
|
||||
然后导航到包含训练脚本的示例文件夹,并安装脚本所需的依赖项:
|
||||
|
||||
```bash
|
||||
cd examples/wuerstchen/text_to_image
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
🤗 Accelerate 是一个帮助您在多个 GPU/TPU 上或使用混合精度进行训练的库。它会根据您的硬件和环境自动配置训练设置。查看 🤗 Accelerate [快速入门](https://huggingface.co/docs/accelerate/quicktour) 以了解更多信息。
|
||||
|
||||
</Tip>
|
||||
|
||||
初始化一个 🤗 Accelerate 环境:
|
||||
|
||||
```bash
|
||||
accelerate config
|
||||
```
|
||||
|
||||
要设置一个默认的 🤗 Accelerate 环境而不选择任何配置:
|
||||
|
||||
```bash
|
||||
accelerate config default
|
||||
```
|
||||
|
||||
或者,如果您的环境不支持交互式 shell,例如笔记本,您可以使用:
|
||||
|
||||
```py
|
||||
from accelerate.utils import write_basic_config
|
||||
|
||||
write_basic_config()
|
||||
```
|
||||
|
||||
最后,如果您想在自己的数据集上训练模型,请查看 [创建训练数据集](create_dataset) 指南,了解如何创建与训练脚本兼容的数据集。
|
||||
|
||||
<Tip>
|
||||
|
||||
以下部分重点介绍了训练脚本中对于理解如何修改它很重要的部分,但并未涵盖 [脚本](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py) 的详细信息。如果您有兴趣了解更多,请随时阅读脚本,并告诉我们您是否有任何问题或疑虑。
|
||||
|
||||
</Tip>
|
||||
|
||||
## 脚本参数
|
||||
|
||||
训练脚本提供了许多参数来帮助您自定义训练运行。所有参数及其描述都可以在 [`parse_args()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L192) 函数中找到。它为每个参数提供了默认值,例如训练批次大小和学习率,但如果您愿意,也可以在训练命令中设置自己的值。
|
||||
|
||||
例如,要使用 fp16 格式的混合精度加速训练,请在训练命令中添加 `--mixed_precision` 参数:
|
||||
|
||||
```bash
|
||||
accelerate launch train_text_to_image_prior.py \
|
||||
--mixed_precision="fp16"
|
||||
```
|
||||
|
||||
大多数参数与 [文本到图像](text2image#script-parameters) 训练指南中的参数相同,因此让我们直接深入 Wuerstchen 训练脚本!
|
||||
|
||||
## 训练脚本
|
||||
|
||||
训练脚本也与 [文本到图像](text2image#training-script) 训练指南类似,但已修改以支持 Wuerstchen。本指南重点介绍 Wuerstchen 训练脚本中独特的代码。
|
||||
|
||||
[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L441) 函数首先初始化图像编码器 - 一个 [EfficientNet](https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/modeling_efficient_net_encoder.py) - 以及通常的调度器和分词器。
|
||||
|
||||
```py
|
||||
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
|
||||
pretrained_checkpoint_file = hf_hub_download("dome272/wuerstchen", filename="model_v2_stage_b.pt")
|
||||
state_dict = torch.load(pretrained_checkpoint_file, map_location="cpu")
|
||||
image_encoder = EfficientNetEncoder()
|
||||
image_encoder.load_state_dict(state_dict["effnet_state_dict"])
|
||||
image_encoder.eval()
|
||||
```
|
||||
|
||||
您还将加载 [`WuerstchenPrior`] 模型以进行优化。
|
||||
|
||||
```py
|
||||
prior = WuerstchenPrior.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
|
||||
|
||||
optimizer = optimizer_cls(
|
||||
prior.parameters(),
|
||||
lr=args.learning_rate,
|
||||
betas=(args.adam_beta1, args.adam_beta2),
|
||||
weight_decay=args.adam_weight_decay,
|
||||
eps=args.adam_epsilon,
|
||||
)
|
||||
```
|
||||
|
||||
接下来,您将对图像应用一些 [transforms](https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L656) 并对标题进行 [tokenize](https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L637):
|
||||
|
||||
```py
|
||||
def preprocess_train(examples):
|
||||
images = [image.conver
|
||||
t("RGB") for image in examples[image_column]]
|
||||
examples["effnet_pixel_values"] = [effnet_transforms(image) for image in images]
|
||||
examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
|
||||
return examples
|
||||
```
|
||||
|
||||
最后,[训练循环](https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L656)处理使用`EfficientNetEncoder`将图像压缩到潜在空间,向潜在表示添加噪声,并使用[`WuerstchenPrior`]模型预测噪声残差。
|
||||
|
||||
```py
|
||||
pred_noise = prior(noisy_latents, timesteps, prompt_embeds)
|
||||
```
|
||||
|
||||
如果您想了解更多关于训练循环的工作原理,请查看[理解管道、模型和调度器](../using-diffusers/write_own_pipeline)教程,该教程分解了去噪过程的基本模式。
|
||||
|
||||
## 启动脚本
|
||||
|
||||
一旦您完成了所有更改或对默认配置满意,就可以启动训练脚本了!🚀
|
||||
|
||||
设置`DATASET_NAME`环境变量为Hub中的数据集名称。本指南使用[Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions)数据集,但您也可以创建和训练自己的数据集(参见[创建用于训练的数据集](create_dataset)指南)。
|
||||
|
||||
<Tip>
|
||||
|
||||
要使用Weights & Biases监控训练进度,请在训练命令中添加`--report_to=wandb`参数。您还需要在训练命令中添加`--validation_prompt`以跟踪结果。这对于调试模型和查看中间结果非常有用。
|
||||
|
||||
</Tip>
|
||||
|
||||
```bash
|
||||
export DATASET_NAME="lambdalabs/naruto-blip-captions"
|
||||
|
||||
accelerate launch train_text_to_image_prior.py \
|
||||
--mixed_precision="fp16" \
|
||||
--dataset_name=$DATASET_NAME \
|
||||
--resolution=768 \
|
||||
--train_batch_size=4 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--dataloader_num_workers=4 \
|
||||
--max_train_steps=15000 \
|
||||
--learning_rate=1e-05 \
|
||||
--max_grad_norm=1 \
|
||||
--checkpoints_total_limit=3 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--validation_prompts="A robot naruto, 4k photo" \
|
||||
--report_to="wandb" \
|
||||
--push_to_hub \
|
||||
--output_dir="wuerstchen-prior-naruto-model"
|
||||
```
|
||||
|
||||
训练完成后,您可以使用新训练的模型进行推理!
|
||||
|
||||
```py
|
||||
import torch
|
||||
from diffusers import AutoPipelineForText2Image
|
||||
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
|
||||
|
||||
pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
caption = "A cute bird naruto holding a shield"
|
||||
images = pipeline(
|
||||
caption,
|
||||
width=1024,
|
||||
height=1536,
|
||||
prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
|
||||
prior_guidance_scale=4.0,
|
||||
num_images_per_prompt=2,
|
||||
).images
|
||||
```
|
||||
|
||||
## 下一步
|
||||
|
||||
恭喜您训练了一个Wuerstchen模型!要了解更多关于如何使用您的新模型的信息,请参
|
||||
以下内容可能有所帮助:
|
||||
|
||||
- 查看 [Wuerstchen](../api/pipelines/wuerstchen#text-to-image-generation) API 文档,了解更多关于如何使用该管道进行文本到图像生成及其限制的信息。
|
||||
@@ -90,7 +90,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -88,7 +88,7 @@ from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -95,7 +95,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -60,7 +60,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -43,8 +43,7 @@ from diffusers.utils import BaseOutput, check_min_version
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
|
||||
check_min_version("0.35.0")
|
||||
|
||||
class MarigoldDepthOutput(BaseOutput):
|
||||
"""
|
||||
|
||||
@@ -74,7 +74,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -67,7 +67,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -80,7 +80,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -79,7 +79,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -66,7 +66,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
if is_torch_npu_available():
|
||||
|
||||
@@ -62,7 +62,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -62,7 +62,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
if is_torch_npu_available():
|
||||
|
||||
@@ -64,7 +64,7 @@ from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -35,7 +35,7 @@ from diffusers.utils import check_min_version
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
# Cache compiled models across invocations of this script.
|
||||
cc.initialize_cache(os.path.expanduser("~/.cache/jax/compilation_cache"))
|
||||
|
||||
@@ -80,7 +80,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -75,7 +75,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -87,7 +87,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -75,7 +75,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -75,7 +75,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -87,7 +87,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -80,7 +80,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -55,7 +55,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -58,7 +58,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -58,7 +58,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -60,7 +60,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -53,7 +53,7 @@ if is_wandb_available():
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ from diffusers.utils import check_min_version, is_wandb_available
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ from diffusers.utils import check_min_version, is_wandb_available
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ if is_wandb_available():
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -57,7 +57,7 @@ if is_wandb_available():
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ from diffusers.utils import check_min_version
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -56,7 +56,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -68,7 +68,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
if is_torch_npu_available():
|
||||
|
||||
@@ -55,7 +55,7 @@ from diffusers.utils.torch_utils import is_compiled_module
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
if is_torch_npu_available():
|
||||
|
||||
@@ -82,7 +82,7 @@ else:
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -56,7 +56,7 @@ else:
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -77,7 +77,7 @@ else:
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -50,7 +50,7 @@ if is_wandb_available():
|
||||
import wandb
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.36.0.dev0")
|
||||
check_min_version("0.35.0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
@@ -269,7 +269,7 @@ version_range_max = max(sys.version_info[1], 10) + 1
|
||||
|
||||
setup(
|
||||
name="diffusers",
|
||||
version="0.36.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
|
||||
version="0.35.2", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
|
||||
description="State-of-the-art diffusion in PyTorch and JAX.",
|
||||
long_description=open("README.md", "r", encoding="utf-8").read(),
|
||||
long_description_content_type="text/markdown",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
__version__ = "0.36.0.dev0"
|
||||
__version__ = "0.35.2"
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -13,7 +13,6 @@ from .utils import (
|
||||
is_k_diffusion_available,
|
||||
is_librosa_available,
|
||||
is_note_seq_available,
|
||||
is_nunchaku_available,
|
||||
is_onnx_available,
|
||||
is_opencv_available,
|
||||
is_optimum_quanto_available,
|
||||
@@ -100,18 +99,6 @@ except OptionalDependencyNotAvailable:
|
||||
else:
|
||||
_import_structure["quantizers.quantization_config"].append("TorchAoConfig")
|
||||
|
||||
try:
|
||||
if not is_torch_available() and not is_accelerate_available() and not is_nunchaku_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_nunchaku_objects
|
||||
|
||||
_import_structure["utils.dummy_nunchaku_objects"] = [
|
||||
name for name in dir(dummy_nunchaku_objects) if not name.startswith("_")
|
||||
]
|
||||
else:
|
||||
_import_structure["quantizers.quantization_config"].append("NunchakuConfig")
|
||||
|
||||
try:
|
||||
if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
@@ -194,7 +181,6 @@ else:
|
||||
"AutoencoderOobleck",
|
||||
"AutoencoderTiny",
|
||||
"AutoModel",
|
||||
"BriaTransformer2DModel",
|
||||
"CacheMixin",
|
||||
"ChromaTransformer2DModel",
|
||||
"CogVideoXTransformer3DModel",
|
||||
@@ -411,7 +397,6 @@ else:
|
||||
"AuraFlowPipeline",
|
||||
"BlipDiffusionControlNetPipeline",
|
||||
"BlipDiffusionPipeline",
|
||||
"BriaPipeline",
|
||||
"ChromaImg2ImgPipeline",
|
||||
"ChromaPipeline",
|
||||
"CLIPImageProjection",
|
||||
@@ -804,14 +789,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
else:
|
||||
from .quantizers.quantization_config import QuantoConfig
|
||||
|
||||
try:
|
||||
if not is_nunchaku_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_optimum_quanto_objects import *
|
||||
else:
|
||||
from .quantizers.quantization_config import NunchakuConfig
|
||||
|
||||
try:
|
||||
if not is_onnx_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
@@ -869,7 +846,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
AutoencoderOobleck,
|
||||
AutoencoderTiny,
|
||||
AutoModel,
|
||||
BriaTransformer2DModel,
|
||||
CacheMixin,
|
||||
ChromaTransformer2DModel,
|
||||
CogVideoXTransformer3DModel,
|
||||
@@ -1056,7 +1032,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
AudioLDM2UNet2DConditionModel,
|
||||
AudioLDMPipeline,
|
||||
AuraFlowPipeline,
|
||||
BriaPipeline,
|
||||
ChromaImg2ImgPipeline,
|
||||
ChromaPipeline,
|
||||
CLIPImageProjection,
|
||||
|
||||
@@ -144,7 +144,6 @@ def _register_attention_processors_metadata():
|
||||
def _register_transformer_blocks_metadata():
|
||||
from ..models.attention import BasicTransformerBlock
|
||||
from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
|
||||
from ..models.transformers.transformer_bria import BriaTransformerBlock
|
||||
from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
|
||||
from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
|
||||
from ..models.transformers.transformer_hunyuan_video import (
|
||||
@@ -166,13 +165,6 @@ def _register_transformer_blocks_metadata():
|
||||
return_encoder_hidden_states_index=None,
|
||||
),
|
||||
)
|
||||
TransformerBlockRegistry.register(
|
||||
model_class=BriaTransformerBlock,
|
||||
metadata=TransformerBlockMetadata(
|
||||
return_hidden_states_index=0,
|
||||
return_encoder_hidden_states_index=None,
|
||||
),
|
||||
)
|
||||
|
||||
# CogVideoX
|
||||
TransformerBlockRegistry.register(
|
||||
|
||||
@@ -23,7 +23,6 @@ from typing_extensions import Self
|
||||
|
||||
from .. import __version__
|
||||
from ..quantizers import DiffusersAutoQuantizer
|
||||
from ..quantizers.quantization_config import NunchakuConfig
|
||||
from ..utils import deprecate, is_accelerate_available, is_torch_version, logging
|
||||
from ..utils.torch_utils import empty_device_cache
|
||||
from .single_file_utils import (
|
||||
@@ -43,7 +42,6 @@ from .single_file_utils import (
|
||||
convert_ltx_vae_checkpoint_to_diffusers,
|
||||
convert_lumina2_to_diffusers,
|
||||
convert_mochi_transformer_checkpoint_to_diffusers,
|
||||
convert_nunchaku_flux_to_diffusers,
|
||||
convert_sana_transformer_to_diffusers,
|
||||
convert_sd3_transformer_checkpoint_to_diffusers,
|
||||
convert_stable_cascade_unet_single_file_to_diffusers,
|
||||
@@ -192,23 +190,6 @@ def _get_mapping_function_kwargs(mapping_fn, **kwargs):
|
||||
return mapping_kwargs
|
||||
|
||||
|
||||
def _maybe_determine_modules_to_not_convert(quantization_config, state_dict):
|
||||
if quantization_config is None:
|
||||
return None
|
||||
else:
|
||||
is_nunchaku = quantization_config.quant_method == "nunchaku"
|
||||
if not is_nunchaku:
|
||||
return None
|
||||
else:
|
||||
no_qweight = set()
|
||||
for key in state_dict:
|
||||
if key.endswith(".weight"):
|
||||
# module name is everything except the last piece after "."
|
||||
module_name = ".".join(key.split(".")[:-1])
|
||||
no_qweight.add(module_name)
|
||||
return sorted(no_qweight)
|
||||
|
||||
|
||||
class FromOriginalModelMixin:
|
||||
"""
|
||||
Load pretrained weights saved in the `.ckpt` or `.safetensors` format into a model.
|
||||
@@ -423,14 +404,8 @@ class FromOriginalModelMixin:
|
||||
model = cls.from_config(diffusers_model_config)
|
||||
|
||||
checkpoint_mapping_kwargs = _get_mapping_function_kwargs(checkpoint_mapping_fn, **kwargs)
|
||||
model_state_dict = model.state_dict()
|
||||
# TODO: Only flux nunchaku checkpoint for now. Unify with how checkpoint mappers are done.
|
||||
# For `nunchaku` checkpoints, we might want to determine the `modules_to_not_convert`.
|
||||
if quantization_config is not None and quantization_config.quant_method == "nunchaku":
|
||||
diffusers_format_checkpoint = convert_nunchaku_flux_to_diffusers(
|
||||
checkpoint, model_state_dict=model_state_dict
|
||||
)
|
||||
elif _should_convert_state_dict_to_diffusers(model_state_dict, checkpoint):
|
||||
|
||||
if _should_convert_state_dict_to_diffusers(model.state_dict(), checkpoint):
|
||||
diffusers_format_checkpoint = checkpoint_mapping_fn(
|
||||
config=diffusers_model_config, checkpoint=checkpoint, **checkpoint_mapping_kwargs
|
||||
)
|
||||
@@ -441,27 +416,6 @@ class FromOriginalModelMixin:
|
||||
raise SingleFileComponentError(
|
||||
f"Failed to load {mapping_class_name}. Weights for this component appear to be missing in the checkpoint."
|
||||
)
|
||||
|
||||
# This step is better off here than above because `diffusers_format_checkpoint` holds the keys we expect.
|
||||
# We can move it to a separate function as well.
|
||||
if quantization_config is not None:
|
||||
original_modules_to_not_convert = quantization_config.modules_to_not_convert or []
|
||||
determined_modules_to_not_convert = _maybe_determine_modules_to_not_convert(
|
||||
quantization_config, checkpoint
|
||||
)
|
||||
if determined_modules_to_not_convert:
|
||||
determined_modules_to_not_convert.extend(original_modules_to_not_convert)
|
||||
determined_modules_to_not_convert = list(set(determined_modules_to_not_convert))
|
||||
logger.debug(
|
||||
f"`modules_to_not_convert` in the quantization_config was updated from {quantization_config.modules_to_not_convert} to {determined_modules_to_not_convert}."
|
||||
)
|
||||
modified_quant_config = quantization_config.to_dict()
|
||||
modified_quant_config["modules_to_not_convert"] = determined_modules_to_not_convert
|
||||
# TODO: figure out a better way.
|
||||
modified_quant_config = NunchakuConfig.from_dict(modified_quant_config)
|
||||
setattr(hf_quantizer, "quantization_config", modified_quant_config)
|
||||
logger.debug("TODO")
|
||||
|
||||
# Check if `_keep_in_fp32_modules` is not None
|
||||
use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (
|
||||
(torch_dtype == torch.float16) or hasattr(hf_quantizer, "use_keep_in_fp32_modules")
|
||||
@@ -489,12 +443,6 @@ class FromOriginalModelMixin:
|
||||
unexpected_keys = [
|
||||
param_name for param_name in diffusers_format_checkpoint if param_name not in empty_state_dict
|
||||
]
|
||||
for k in unexpected_keys:
|
||||
if "single_transformer_blocks.0" in k:
|
||||
print(f"Unexpected {k=}")
|
||||
for k in empty_state_dict:
|
||||
if "single_transformer_blocks.0" in k:
|
||||
print(f"model {k=}")
|
||||
device_map = {"": param_device}
|
||||
load_model_dict_into_meta(
|
||||
model,
|
||||
|
||||
@@ -2189,105 +2189,6 @@ def convert_animatediff_checkpoint_to_diffusers(checkpoint, **kwargs):
|
||||
return converted_state_dict
|
||||
|
||||
|
||||
# Adapted from https://github.com/nunchaku-tech/nunchaku/blob/3ec299f439f9986a69ded320798cab4e258c871d/nunchaku/models/transformers/transformer_flux_v2.py#L395
|
||||
def convert_nunchaku_flux_to_diffusers(checkpoint, **kwargs):
|
||||
from .single_file_utils_nunchaku import _unpack_qkv_state_dict
|
||||
|
||||
_SMOOTH_ORIG_RE = re.compile(r"\.smooth_orig(\.|$)")
|
||||
_SMOOTH_RE = re.compile(r"\.smooth(\.|$)")
|
||||
|
||||
new_state_dict = {}
|
||||
model_state_dict = kwargs["model_state_dict"]
|
||||
|
||||
ckpt_keys = list(checkpoint.keys())
|
||||
for k in ckpt_keys:
|
||||
if "qweight" in k:
|
||||
# only the shape information of this tensor is needed
|
||||
v = checkpoint[k]
|
||||
# if the tensor has qweight, but does not have low-rank branch, we need to add some artificial tensors
|
||||
for t in ["lora_up", "lora_down"]:
|
||||
new_k = k.replace(".qweight", f".{t}")
|
||||
if new_k not in ckpt_keys:
|
||||
oc, ic = v.shape
|
||||
ic = ic * 2 # v is packed into INT8, so we need to double the size
|
||||
checkpoint[k.replace(".qweight", f".{t}")] = torch.zeros(
|
||||
(0, ic) if t == "lora_down" else (oc, 0), device=v.device, dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
for k, v in checkpoint.items():
|
||||
new_k = k # start with original, then apply independent replacements
|
||||
|
||||
if k.startswith("single_transformer_blocks."):
|
||||
# attention / qkv / norms
|
||||
new_k = new_k.replace(".qkv_proj.", ".attn.to_qkv.")
|
||||
new_k = new_k.replace(".out_proj.", ".proj_out.")
|
||||
new_k = new_k.replace(".norm_k.", ".attn.norm_k.")
|
||||
new_k = new_k.replace(".norm_q.", ".attn.norm_q.")
|
||||
|
||||
# mlp heads
|
||||
new_k = new_k.replace(".mlp_fc1.", ".proj_mlp.")
|
||||
new_k = new_k.replace(".mlp_fc2.", ".proj_out.")
|
||||
|
||||
# smooth params (use regex to avoid substring collisions)
|
||||
new_k = _SMOOTH_ORIG_RE.sub(r".smooth_factor_orig\1", new_k)
|
||||
new_k = _SMOOTH_RE.sub(r".smooth_factor\1", new_k)
|
||||
|
||||
# lora -> proj
|
||||
new_k = new_k.replace(".lora_down", ".proj_down")
|
||||
new_k = new_k.replace(".lora_up", ".proj_up")
|
||||
|
||||
elif k.startswith("transformer_blocks."):
|
||||
# feed-forward (context & base)
|
||||
new_k = new_k.replace(".mlp_context_fc1.", ".ff_context.net.0.proj.")
|
||||
new_k = new_k.replace(".mlp_context_fc2.", ".ff_context.net.2.")
|
||||
new_k = new_k.replace(".mlp_fc1.", ".ff.net.0.proj.")
|
||||
new_k = new_k.replace(".mlp_fc2.", ".ff.net.2.")
|
||||
|
||||
# attention projections
|
||||
new_k = new_k.replace(".qkv_proj_context.", ".attn.add_qkv_proj.")
|
||||
new_k = new_k.replace(".qkv_proj.", ".attn.to_qkv.")
|
||||
new_k = new_k.replace(".out_proj.", ".attn.to_out.0.")
|
||||
new_k = new_k.replace(".out_proj_context.", ".attn.to_add_out.")
|
||||
|
||||
# norms
|
||||
new_k = new_k.replace(".norm_k.", ".attn.norm_k.")
|
||||
new_k = new_k.replace(".norm_q.", ".attn.norm_q.")
|
||||
new_k = new_k.replace(".norm_added_k.", ".attn.norm_added_k.")
|
||||
new_k = new_k.replace(".norm_added_q.", ".attn.norm_added_q.")
|
||||
|
||||
# smooth params
|
||||
new_k = _SMOOTH_ORIG_RE.sub(r".smooth_factor_orig\1", new_k)
|
||||
new_k = _SMOOTH_RE.sub(r".smooth_factor\1", new_k)
|
||||
|
||||
# lora -> proj
|
||||
new_k = new_k.replace(".lora_down", ".proj_down")
|
||||
new_k = new_k.replace(".lora_up", ".proj_up")
|
||||
|
||||
new_state_dict[new_k] = v
|
||||
|
||||
new_state_dict = _unpack_qkv_state_dict(new_state_dict)
|
||||
|
||||
# some remnant keys need to be patched
|
||||
new_sd_keys = list(new_state_dict.keys())
|
||||
for k in new_sd_keys:
|
||||
if "qweight" in k:
|
||||
no_qweight_k = ".".join(k.split(".qweight")[:-1])
|
||||
for unexpected_k in ["wzeros"]:
|
||||
unexpected_k = no_qweight_k + f".{unexpected_k}"
|
||||
if unexpected_k in new_sd_keys:
|
||||
_ = new_state_dict.pop(unexpected_k)
|
||||
for k in model_state_dict:
|
||||
if k not in new_state_dict:
|
||||
# CPU device for now
|
||||
new_state_dict[k] = torch.ones_like(model_state_dict[k], device="cpu")
|
||||
|
||||
for k in new_state_dict:
|
||||
if "single_transformer_blocks.0" in k and k.endswith(".weight"):
|
||||
print(f"{k=}")
|
||||
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def convert_flux_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
|
||||
converted_state_dict = {}
|
||||
keys = list(checkpoint.keys())
|
||||
|
||||
@@ -1,104 +0,0 @@
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
_QKV_ANCHORS_NUNCHAKU = ("to_qkv", "add_qkv_proj")
|
||||
_ALLOWED_SUFFIXES_NUNCHAKU = {
|
||||
"bias",
|
||||
"proj_down",
|
||||
"proj_up",
|
||||
"qweight",
|
||||
"smooth_factor",
|
||||
"smooth_factor_orig",
|
||||
"wscales",
|
||||
}
|
||||
|
||||
_QKV_NUNCHAKU_REGEX = re.compile(
|
||||
rf"^(?P<prefix>.*)\.(?:{'|'.join(map(re.escape, _QKV_ANCHORS_NUNCHAKU))})\.(?P<suffix>.+)$"
|
||||
)
|
||||
|
||||
|
||||
def _pick_split_dim(t: torch.Tensor, suffix: str) -> int:
|
||||
"""
|
||||
Choose which dimension to split by 3. Heuristics:
|
||||
- 1D -> dim 0
|
||||
- 2D -> prefer dim=1 for 'qweight' (common layout [*, 3*out_features]),
|
||||
otherwise prefer dim=0 (common layout [3*out_features, *]).
|
||||
- If preferred dim isn't divisible by 3, try the other; else error.
|
||||
"""
|
||||
shape = list(t.shape)
|
||||
if len(shape) == 0:
|
||||
raise ValueError("Cannot split a scalar into Q/K/V.")
|
||||
|
||||
if len(shape) == 1:
|
||||
dim = 0
|
||||
if shape[dim] % 3 == 0:
|
||||
return dim
|
||||
raise ValueError(f"1D tensor of length {shape[0]} not divisible by 3.")
|
||||
|
||||
# len(shape) >= 2
|
||||
preferred = 1 if suffix == "qweight" else 0
|
||||
other = 0 if preferred == 1 else 1
|
||||
|
||||
if shape[preferred] % 3 == 0:
|
||||
return preferred
|
||||
if shape[other] % 3 == 0:
|
||||
return other
|
||||
|
||||
# Fall back: any dim divisible by 3
|
||||
for d, s in enumerate(shape):
|
||||
if s % 3 == 0:
|
||||
return d
|
||||
|
||||
raise ValueError(f"None of the dims {shape} are divisible by 3 for suffix '{suffix}'.")
|
||||
|
||||
|
||||
def _split_qkv(t: torch.Tensor, dim: int):
|
||||
return torch.tensor_split(t, 3, dim=dim)
|
||||
|
||||
|
||||
def _unpack_qkv_state_dict(
|
||||
state_dict: dict, anchors=_QKV_ANCHORS_NUNCHAKU, allowed_suffixes=_ALLOWED_SUFFIXES_NUNCHAKU
|
||||
):
|
||||
"""
|
||||
Convert fused QKV entries (e.g., '...to_qkv.bias', '...qkv_proj.wscales') into separate Q/K/V entries:
|
||||
'...to_q.bias', '...to_k.bias', '...to_v.bias' '...to_q.wscales', '...to_k.wscales', '...to_v.wscales'
|
||||
Returns a NEW dict; original is not modified.
|
||||
|
||||
Only keys with suffix in `allowed_suffixes` are processed. Keys with non-divisible-by-3 tensors raise a ValueError.:
|
||||
"""
|
||||
anchors = tuple(anchors)
|
||||
allowed_suffixes = set(allowed_suffixes)
|
||||
|
||||
new_sd: dict = {}
|
||||
sd_keys = list(state_dict.keys())
|
||||
for k in sd_keys:
|
||||
m = _QKV_NUNCHAKU_REGEX.match(k)
|
||||
v = state_dict.pop(k)
|
||||
if m:
|
||||
suffix = m.group("suffix")
|
||||
if suffix not in allowed_suffixes:
|
||||
# keep as-is if it's not one of the targeted suffixes
|
||||
new_sd[k] = v
|
||||
continue
|
||||
|
||||
prefix = m.group("prefix") # everything before .to_qkv/.qkv_proj
|
||||
# Decide split axis
|
||||
split_dim = _pick_split_dim(v, suffix)
|
||||
q, k_, vv = _split_qkv(v, dim=split_dim)
|
||||
|
||||
# Build new keys
|
||||
base_q = f"{prefix}.to_q.{suffix}"
|
||||
base_k = f"{prefix}.to_k.{suffix}"
|
||||
base_v = f"{prefix}.to_v.{suffix}"
|
||||
|
||||
# Write into result dict
|
||||
new_sd[base_q] = q
|
||||
new_sd[base_k] = k_
|
||||
new_sd[base_v] = vv
|
||||
else:
|
||||
# not a fused qkv key
|
||||
new_sd[k] = v
|
||||
|
||||
return new_sd
|
||||
@@ -76,7 +76,6 @@ if is_torch_available():
|
||||
_import_structure["transformers.t5_film_transformer"] = ["T5FilmDecoder"]
|
||||
_import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
|
||||
_import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"]
|
||||
_import_structure["transformers.transformer_bria"] = ["BriaTransformer2DModel"]
|
||||
_import_structure["transformers.transformer_chroma"] = ["ChromaTransformer2DModel"]
|
||||
_import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"]
|
||||
_import_structure["transformers.transformer_cogview4"] = ["CogView4Transformer2DModel"]
|
||||
@@ -159,7 +158,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
from .transformers import (
|
||||
AllegroTransformer3DModel,
|
||||
AuraFlowTransformer2DModel,
|
||||
BriaTransformer2DModel,
|
||||
ChromaTransformer2DModel,
|
||||
CogVideoXTransformer3DModel,
|
||||
CogView3PlusTransformer2DModel,
|
||||
|
||||
@@ -110,6 +110,27 @@ if _CAN_USE_XFORMERS_ATTN:
|
||||
else:
|
||||
xops = None
|
||||
|
||||
# Version guard for PyTorch compatibility - custom_op was added in PyTorch 2.4
|
||||
if torch.__version__ >= "2.4.0":
|
||||
_custom_op = torch.library.custom_op
|
||||
_register_fake = torch.library.register_fake
|
||||
else:
|
||||
|
||||
def custom_op_no_op(name, fn=None, /, *, mutates_args, device_types=None, schema=None):
|
||||
def wrap(func):
|
||||
return func
|
||||
|
||||
return wrap if fn is None else fn
|
||||
|
||||
def register_fake_no_op(op, fn=None, /, *, lib=None, _stacklevel=1):
|
||||
def wrap(func):
|
||||
return func
|
||||
|
||||
return wrap if fn is None else fn
|
||||
|
||||
_custom_op = custom_op_no_op
|
||||
_register_fake = register_fake_no_op
|
||||
|
||||
|
||||
logger = get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
@@ -473,12 +494,11 @@ def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
|
||||
|
||||
# ===== torch op registrations =====
|
||||
# Registrations are required for fullgraph tracing compatibility
|
||||
|
||||
|
||||
# TODO: library.custom_op and register_fake probably need version guards?
|
||||
# TODO: this is only required because the beta release FA3 does not have it. There is a PR adding
|
||||
# this but it was never merged: https://github.com/Dao-AILab/flash-attention/pull/1590
|
||||
@torch.library.custom_op("flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda")
|
||||
|
||||
|
||||
@_custom_op("flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda")
|
||||
def _wrapped_flash_attn_3_original(
|
||||
query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
@@ -487,7 +507,7 @@ def _wrapped_flash_attn_3_original(
|
||||
return out, lse
|
||||
|
||||
|
||||
@torch.library.register_fake("flash_attn_3::_flash_attn_forward")
|
||||
@_register_fake("flash_attn_3::_flash_attn_forward")
|
||||
def _(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
batch_size, seq_len, num_heads, head_dim = query.shape
|
||||
lse_shape = (batch_size, seq_len, num_heads)
|
||||
|
||||
@@ -112,6 +112,9 @@ def _determine_device_map(
|
||||
device_map_kwargs["max_memory"] = max_memory
|
||||
device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs)
|
||||
|
||||
if hf_quantizer is not None:
|
||||
hf_quantizer.validate_environment(device_map=device_map)
|
||||
|
||||
return device_map
|
||||
|
||||
|
||||
@@ -297,13 +300,6 @@ def load_model_dict_into_meta(
|
||||
offload_index = offload_weight(param, param_name, offload_folder, offload_index)
|
||||
elif param_device == "cpu" and state_dict_index is not None:
|
||||
state_dict_index = offload_weight(param, param_name, state_dict_folder, state_dict_index)
|
||||
# This check below might be a bit counter-intuitive in nature. This is because we're checking if the param
|
||||
# or its module is quantized and if so, we're proceeding with creating a quantized param. This is because
|
||||
# of the way pre-trained models are loaded. They're initialized under "meta" device, where
|
||||
# quantization layers are first injected. Hence, for a model that is either pre-quantized or supplemented
|
||||
# with a `quantization_config` during `from_pretrained`, we expect `check_if_quantized_param` to return True.
|
||||
# Then depending on the quantization backend being used, we run the actual quantization step under
|
||||
# `create_quantized_param`.
|
||||
elif is_quantized and (
|
||||
hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)
|
||||
):
|
||||
|
||||
@@ -17,7 +17,6 @@ if is_torch_available():
|
||||
from .t5_film_transformer import T5FilmDecoder
|
||||
from .transformer_2d import Transformer2DModel
|
||||
from .transformer_allegro import AllegroTransformer3DModel
|
||||
from .transformer_bria import BriaTransformer2DModel
|
||||
from .transformer_chroma import ChromaTransformer2DModel
|
||||
from .transformer_cogview3plus import CogView3PlusTransformer2DModel
|
||||
from .transformer_cogview4 import CogView4Transformer2DModel
|
||||
|
||||
@@ -1,719 +0,0 @@
|
||||
import inspect
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ...configuration_utils import ConfigMixin, register_to_config
|
||||
from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
|
||||
from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
|
||||
from ...utils.torch_utils import maybe_allow_in_graph
|
||||
from ..attention import AttentionModuleMixin, FeedForward
|
||||
from ..attention_dispatch import dispatch_attention_fn
|
||||
from ..cache_utils import CacheMixin
|
||||
from ..embeddings import TimestepEmbedding, apply_rotary_emb, get_timestep_embedding
|
||||
from ..modeling_outputs import Transformer2DModelOutput
|
||||
from ..modeling_utils import ModelMixin
|
||||
from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _get_projections(attn: "BriaAttention", hidden_states, encoder_hidden_states=None):
|
||||
query = attn.to_q(hidden_states)
|
||||
key = attn.to_k(hidden_states)
|
||||
value = attn.to_v(hidden_states)
|
||||
|
||||
encoder_query = encoder_key = encoder_value = None
|
||||
if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
|
||||
encoder_query = attn.add_q_proj(encoder_hidden_states)
|
||||
encoder_key = attn.add_k_proj(encoder_hidden_states)
|
||||
encoder_value = attn.add_v_proj(encoder_hidden_states)
|
||||
|
||||
return query, key, value, encoder_query, encoder_key, encoder_value
|
||||
|
||||
|
||||
def _get_fused_projections(attn: "BriaAttention", hidden_states, encoder_hidden_states=None):
|
||||
query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
|
||||
|
||||
encoder_query = encoder_key = encoder_value = (None,)
|
||||
if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
|
||||
encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
|
||||
|
||||
return query, key, value, encoder_query, encoder_key, encoder_value
|
||||
|
||||
|
||||
def _get_qkv_projections(attn: "BriaAttention", hidden_states, encoder_hidden_states=None):
|
||||
if attn.fused_projections:
|
||||
return _get_fused_projections(attn, hidden_states, encoder_hidden_states)
|
||||
return _get_projections(attn, hidden_states, encoder_hidden_states)
|
||||
|
||||
|
||||
def get_1d_rotary_pos_embed(
|
||||
dim: int,
|
||||
pos: Union[np.ndarray, int],
|
||||
theta: float = 10000.0,
|
||||
use_real=False,
|
||||
linear_factor=1.0,
|
||||
ntk_factor=1.0,
|
||||
repeat_interleave_real=True,
|
||||
freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)
|
||||
):
|
||||
"""
|
||||
Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
|
||||
|
||||
This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
|
||||
index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
|
||||
data type.
|
||||
|
||||
Args:
|
||||
dim (`int`): Dimension of the frequency tensor.
|
||||
pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
|
||||
theta (`float`, *optional*, defaults to 10000.0):
|
||||
Scaling factor for frequency computation. Defaults to 10000.0.
|
||||
use_real (`bool`, *optional*):
|
||||
If True, return real part and imaginary part separately. Otherwise, return complex numbers.
|
||||
linear_factor (`float`, *optional*, defaults to 1.0):
|
||||
Scaling factor for the context extrapolation. Defaults to 1.0.
|
||||
ntk_factor (`float`, *optional*, defaults to 1.0):
|
||||
Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
|
||||
repeat_interleave_real (`bool`, *optional*, defaults to `True`):
|
||||
If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
|
||||
Otherwise, they are concateanted with themselves.
|
||||
freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
|
||||
the dtype of the frequency tensor.
|
||||
Returns:
|
||||
`torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
|
||||
"""
|
||||
assert dim % 2 == 0
|
||||
|
||||
if isinstance(pos, int):
|
||||
pos = torch.arange(pos)
|
||||
if isinstance(pos, np.ndarray):
|
||||
pos = torch.from_numpy(pos) # type: ignore # [S]
|
||||
|
||||
theta = theta * ntk_factor
|
||||
freqs = (
|
||||
1.0
|
||||
/ (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))
|
||||
/ linear_factor
|
||||
) # [D/2]
|
||||
freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
|
||||
if use_real and repeat_interleave_real:
|
||||
# bria
|
||||
freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]
|
||||
freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]
|
||||
return freqs_cos, freqs_sin
|
||||
elif use_real:
|
||||
# stable audio, allegro
|
||||
freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]
|
||||
freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]
|
||||
return freqs_cos, freqs_sin
|
||||
else:
|
||||
# lumina
|
||||
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
|
||||
return freqs_cis
|
||||
|
||||
|
||||
class BriaAttnProcessor:
|
||||
_attention_backend = None
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: "BriaAttention",
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
|
||||
attn, hidden_states, encoder_hidden_states
|
||||
)
|
||||
|
||||
query = query.unflatten(-1, (attn.heads, -1))
|
||||
key = key.unflatten(-1, (attn.heads, -1))
|
||||
value = value.unflatten(-1, (attn.heads, -1))
|
||||
|
||||
query = attn.norm_q(query)
|
||||
key = attn.norm_k(key)
|
||||
|
||||
if attn.added_kv_proj_dim is not None:
|
||||
encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
|
||||
encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
|
||||
encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
|
||||
|
||||
encoder_query = attn.norm_added_q(encoder_query)
|
||||
encoder_key = attn.norm_added_k(encoder_key)
|
||||
|
||||
query = torch.cat([encoder_query, query], dim=1)
|
||||
key = torch.cat([encoder_key, key], dim=1)
|
||||
value = torch.cat([encoder_value, value], dim=1)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
|
||||
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
|
||||
|
||||
hidden_states = dispatch_attention_fn(
|
||||
query, key, value, attn_mask=attention_mask, backend=self._attention_backend
|
||||
)
|
||||
hidden_states = hidden_states.flatten(2, 3)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
if encoder_hidden_states is not None:
|
||||
encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
|
||||
[encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
|
||||
)
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
else:
|
||||
return hidden_states
|
||||
|
||||
|
||||
class BriaAttention(torch.nn.Module, AttentionModuleMixin):
|
||||
_default_processor_cls = BriaAttnProcessor
|
||||
_available_processors = [
|
||||
BriaAttnProcessor,
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query_dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = False,
|
||||
added_kv_proj_dim: Optional[int] = None,
|
||||
added_proj_bias: Optional[bool] = True,
|
||||
out_bias: bool = True,
|
||||
eps: float = 1e-5,
|
||||
out_dim: int = None,
|
||||
context_pre_only: Optional[bool] = None,
|
||||
pre_only: bool = False,
|
||||
elementwise_affine: bool = True,
|
||||
processor=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.head_dim = dim_head
|
||||
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
|
||||
self.query_dim = query_dim
|
||||
self.use_bias = bias
|
||||
self.dropout = dropout
|
||||
self.out_dim = out_dim if out_dim is not None else query_dim
|
||||
self.context_pre_only = context_pre_only
|
||||
self.pre_only = pre_only
|
||||
self.heads = out_dim // dim_head if out_dim is not None else heads
|
||||
self.added_kv_proj_dim = added_kv_proj_dim
|
||||
self.added_proj_bias = added_proj_bias
|
||||
|
||||
self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
|
||||
self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
|
||||
self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
|
||||
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
|
||||
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
|
||||
|
||||
if not self.pre_only:
|
||||
self.to_out = torch.nn.ModuleList([])
|
||||
self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
|
||||
self.to_out.append(torch.nn.Dropout(dropout))
|
||||
|
||||
if added_kv_proj_dim is not None:
|
||||
self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
|
||||
self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
|
||||
self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
|
||||
self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
|
||||
self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
|
||||
self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
|
||||
|
||||
if processor is None:
|
||||
processor = self._default_processor_cls()
|
||||
self.set_processor(processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
|
||||
quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
|
||||
unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
|
||||
if len(unused_kwargs) > 0:
|
||||
logger.warning(
|
||||
f"attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
|
||||
)
|
||||
kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
|
||||
return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
|
||||
|
||||
|
||||
class BriaEmbedND(torch.nn.Module):
|
||||
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
|
||||
def __init__(self, theta: int, axes_dim: List[int]):
|
||||
super().__init__()
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
|
||||
def forward(self, ids: torch.Tensor) -> torch.Tensor:
|
||||
n_axes = ids.shape[-1]
|
||||
cos_out = []
|
||||
sin_out = []
|
||||
pos = ids.float()
|
||||
is_mps = ids.device.type == "mps"
|
||||
freqs_dtype = torch.float32 if is_mps else torch.float64
|
||||
for i in range(n_axes):
|
||||
cos, sin = get_1d_rotary_pos_embed(
|
||||
self.axes_dim[i],
|
||||
pos[:, i],
|
||||
theta=self.theta,
|
||||
repeat_interleave_real=True,
|
||||
use_real=True,
|
||||
freqs_dtype=freqs_dtype,
|
||||
)
|
||||
cos_out.append(cos)
|
||||
sin_out.append(sin)
|
||||
freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
|
||||
freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
class BriaTimesteps(nn.Module):
|
||||
def __init__(
|
||||
self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1, time_theta=10000
|
||||
):
|
||||
super().__init__()
|
||||
self.num_channels = num_channels
|
||||
self.flip_sin_to_cos = flip_sin_to_cos
|
||||
self.downscale_freq_shift = downscale_freq_shift
|
||||
self.scale = scale
|
||||
self.time_theta = time_theta
|
||||
|
||||
def forward(self, timesteps):
|
||||
t_emb = get_timestep_embedding(
|
||||
timesteps,
|
||||
self.num_channels,
|
||||
flip_sin_to_cos=self.flip_sin_to_cos,
|
||||
downscale_freq_shift=self.downscale_freq_shift,
|
||||
scale=self.scale,
|
||||
max_period=self.time_theta,
|
||||
)
|
||||
return t_emb
|
||||
|
||||
|
||||
class BriaTimestepProjEmbeddings(nn.Module):
|
||||
def __init__(self, embedding_dim, time_theta):
|
||||
super().__init__()
|
||||
|
||||
self.time_proj = BriaTimesteps(
|
||||
num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, time_theta=time_theta
|
||||
)
|
||||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
|
||||
|
||||
def forward(self, timestep, dtype):
|
||||
timesteps_proj = self.time_proj(timestep)
|
||||
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=dtype)) # (N, D)
|
||||
return timesteps_emb
|
||||
|
||||
|
||||
class BriaPosEmbed(torch.nn.Module):
|
||||
# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
|
||||
def __init__(self, theta: int, axes_dim: List[int]):
|
||||
super().__init__()
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
|
||||
def forward(self, ids: torch.Tensor) -> torch.Tensor:
|
||||
n_axes = ids.shape[-1]
|
||||
cos_out = []
|
||||
sin_out = []
|
||||
pos = ids.float()
|
||||
is_mps = ids.device.type == "mps"
|
||||
freqs_dtype = torch.float32 if is_mps else torch.float64
|
||||
for i in range(n_axes):
|
||||
cos, sin = get_1d_rotary_pos_embed(
|
||||
self.axes_dim[i],
|
||||
pos[:, i],
|
||||
theta=self.theta,
|
||||
repeat_interleave_real=True,
|
||||
use_real=True,
|
||||
freqs_dtype=freqs_dtype,
|
||||
)
|
||||
cos_out.append(cos)
|
||||
sin_out.append(sin)
|
||||
freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
|
||||
freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class BriaTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.norm1 = AdaLayerNormZero(dim)
|
||||
self.norm1_context = AdaLayerNormZero(dim)
|
||||
|
||||
self.attn = BriaAttention(
|
||||
query_dim=dim,
|
||||
added_kv_proj_dim=dim,
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
out_dim=dim,
|
||||
context_pre_only=False,
|
||||
bias=True,
|
||||
processor=BriaAttnProcessor(),
|
||||
eps=eps,
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
|
||||
|
||||
self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
|
||||
|
||||
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
|
||||
encoder_hidden_states, emb=temb
|
||||
)
|
||||
attention_kwargs = attention_kwargs or {}
|
||||
|
||||
# Attention.
|
||||
attention_outputs = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
encoder_hidden_states=norm_encoder_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
**attention_kwargs,
|
||||
)
|
||||
|
||||
if len(attention_outputs) == 2:
|
||||
attn_output, context_attn_output = attention_outputs
|
||||
elif len(attention_outputs) == 3:
|
||||
attn_output, context_attn_output, ip_attn_output = attention_outputs
|
||||
|
||||
# Process attention outputs for the `hidden_states`.
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
hidden_states = hidden_states + attn_output
|
||||
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
|
||||
hidden_states = hidden_states + ff_output
|
||||
if len(attention_outputs) == 3:
|
||||
hidden_states = hidden_states + ip_attn_output
|
||||
|
||||
# Process attention outputs for the `encoder_hidden_states`.
|
||||
context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
|
||||
encoder_hidden_states = encoder_hidden_states + context_attn_output
|
||||
|
||||
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
|
||||
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
|
||||
|
||||
context_ff_output = self.ff_context(norm_encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
|
||||
if encoder_hidden_states.dtype == torch.float16:
|
||||
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
|
||||
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class BriaSingleTransformerBlock(nn.Module):
|
||||
def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0):
|
||||
super().__init__()
|
||||
self.mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
|
||||
self.norm = AdaLayerNormZeroSingle(dim)
|
||||
self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
|
||||
self.act_mlp = nn.GELU(approximate="tanh")
|
||||
self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
|
||||
|
||||
processor = BriaAttnProcessor()
|
||||
|
||||
self.attn = BriaAttention(
|
||||
query_dim=dim,
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
out_dim=dim,
|
||||
bias=True,
|
||||
processor=processor,
|
||||
eps=1e-6,
|
||||
pre_only=True,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> torch.Tensor:
|
||||
text_seq_len = encoder_hidden_states.shape[1]
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
|
||||
residual = hidden_states
|
||||
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
|
||||
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
|
||||
attention_kwargs = attention_kwargs or {}
|
||||
attn_output = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
**attention_kwargs,
|
||||
)
|
||||
|
||||
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
|
||||
gate = gate.unsqueeze(1)
|
||||
hidden_states = gate * self.proj_out(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
if hidden_states.dtype == torch.float16:
|
||||
hidden_states = hidden_states.clip(-65504, 65504)
|
||||
|
||||
encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
|
||||
return encoder_hidden_states, hidden_states
|
||||
|
||||
|
||||
class BriaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin):
|
||||
"""
|
||||
The Transformer model introduced in Flux. Based on FluxPipeline with several changes:
|
||||
- no pooled embeddings
|
||||
- We use zero padding for prompts
|
||||
- No guidance embedding since this is not a distilled version
|
||||
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
|
||||
|
||||
Parameters:
|
||||
patch_size (`int`): Patch size to turn the input data into small patches.
|
||||
in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
|
||||
num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
|
||||
num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
|
||||
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
|
||||
num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
|
||||
joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
|
||||
pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
|
||||
guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int = 1,
|
||||
in_channels: int = 64,
|
||||
num_layers: int = 19,
|
||||
num_single_layers: int = 38,
|
||||
attention_head_dim: int = 128,
|
||||
num_attention_heads: int = 24,
|
||||
joint_attention_dim: int = 4096,
|
||||
pooled_projection_dim: int = None,
|
||||
guidance_embeds: bool = False,
|
||||
axes_dims_rope: List[int] = [16, 56, 56],
|
||||
rope_theta=10000,
|
||||
time_theta=10000,
|
||||
):
|
||||
super().__init__()
|
||||
self.out_channels = in_channels
|
||||
self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
|
||||
|
||||
self.pos_embed = BriaEmbedND(theta=rope_theta, axes_dim=axes_dims_rope)
|
||||
|
||||
self.time_embed = BriaTimestepProjEmbeddings(embedding_dim=self.inner_dim, time_theta=time_theta)
|
||||
if guidance_embeds:
|
||||
self.guidance_embed = BriaTimestepProjEmbeddings(embedding_dim=self.inner_dim)
|
||||
|
||||
self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
|
||||
self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BriaTransformerBlock(
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=self.config.num_attention_heads,
|
||||
attention_head_dim=self.config.attention_head_dim,
|
||||
)
|
||||
for i in range(self.config.num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.single_transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BriaSingleTransformerBlock(
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=self.config.num_attention_heads,
|
||||
attention_head_dim=self.config.attention_head_dim,
|
||||
)
|
||||
for i in range(self.config.num_single_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
|
||||
self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
pooled_projections: torch.Tensor = None,
|
||||
timestep: torch.LongTensor = None,
|
||||
img_ids: torch.Tensor = None,
|
||||
txt_ids: torch.Tensor = None,
|
||||
guidance: torch.Tensor = None,
|
||||
attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
return_dict: bool = True,
|
||||
controlnet_block_samples=None,
|
||||
controlnet_single_block_samples=None,
|
||||
) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
|
||||
"""
|
||||
The [`BriaTransformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
|
||||
Input `hidden_states`.
|
||||
encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
|
||||
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
|
||||
pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
|
||||
from the embeddings of input conditions.
|
||||
timestep ( `torch.LongTensor`):
|
||||
Used to indicate denoising step.
|
||||
block_controlnet_hidden_states: (`list` of `torch.Tensor`):
|
||||
A list of tensors that if specified are added to the residuals of transformer blocks.
|
||||
attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
if attention_kwargs is not None:
|
||||
attention_kwargs = attention_kwargs.copy()
|
||||
lora_scale = attention_kwargs.pop("scale", 1.0)
|
||||
else:
|
||||
lora_scale = 1.0
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# weight the lora layers by setting `lora_scale` for each PEFT layer
|
||||
scale_lora_layers(self, lora_scale)
|
||||
else:
|
||||
if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
|
||||
logger.warning(
|
||||
"Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
|
||||
)
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
|
||||
timestep = timestep.to(hidden_states.dtype)
|
||||
if guidance is not None:
|
||||
guidance = guidance.to(hidden_states.dtype)
|
||||
else:
|
||||
guidance = None
|
||||
|
||||
temb = self.time_embed(timestep, dtype=hidden_states.dtype)
|
||||
|
||||
if guidance:
|
||||
temb += self.guidance_embed(guidance, dtype=hidden_states.dtype)
|
||||
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
|
||||
if len(txt_ids.shape) == 3:
|
||||
txt_ids = txt_ids[0]
|
||||
|
||||
if len(img_ids.shape) == 3:
|
||||
img_ids = img_ids[0]
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=0)
|
||||
image_rotary_emb = self.pos_embed(ids)
|
||||
|
||||
for index_block, block in enumerate(self.transformer_blocks):
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
|
||||
block,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
attention_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_block_samples is not None:
|
||||
interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
|
||||
|
||||
for index_block, block in enumerate(self.single_transformer_blocks):
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
|
||||
block,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
attention_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_single_block_samples is not None:
|
||||
interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
+ controlnet_single_block_samples[index_block // interval_control]
|
||||
)
|
||||
|
||||
hidden_states = self.norm_out(hidden_states, temb)
|
||||
output = self.proj_out(hidden_states)
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# remove `lora_scale` from each PEFT layer
|
||||
unscale_lora_layers(self, lora_scale)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
@@ -350,7 +350,9 @@ class LTXVideoTransformerBlock(nn.Module):
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
|
||||
num_ada_params = self.scale_shift_table.shape[0]
|
||||
ada_values = self.scale_shift_table[None, None] + temb.reshape(batch_size, temb.size(1), num_ada_params, -1)
|
||||
ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape(
|
||||
batch_size, temb.size(1), num_ada_params, -1
|
||||
)
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user