混元-DiT
来自腾讯混元的《混元-DiT:一种具有细粒度中文理解能力的强大多分辨率扩散Transformer》。
论文的摘要如下:
我们提出了混元-DiT,一种具有细粒度英汉理解能力的文本到图像扩散Transformer。为了构建混元-DiT,我们精心设计了Transformer结构、文本编码器和位置编码。我们还从头开始构建了整个数据流水线,以更新和评估数据,用于迭代模型优化。为了实现细粒度的语言理解,我们训练了一个多模态大语言模型来精炼图像的描述。最终,混元-DiT能够与用户进行多轮多模态对话,根据上下文生成和精炼图像。通过我们全面的50多名专业人类评估者的评估协议,混元-DiT在与其它开源模型相比的中文到图像生成方面达到了新的最先进水平。
你可以在Tencent/HunyuanDiT找到原始代码库,并在Tencent-Hunyuan找到所有可用的检查点。
亮点:HunyuanDiT支持中文/英文到图像的多分辨率生成。
HunyuanDiT包含以下组件:
- 使用扩散Transformer作为主干
- 结合了两种文本编码器,一种双语CLIP和一种多语言T5编码器
优化
你可以使用torch.compile和前馈分块来优化管道的运行时和内存消耗。要了解其他优化方法,请查看加速推理和减少内存使用指南。
推理
使用torch.compile
来减少推理延迟。
首先,加载管道:
from diffusers import HunyuanDiTPipeline
import torch
pipeline = HunyuanDiTPipeline.from_pretrained(
"Tencent-Hunyuan/HunyuanDiT-Diffusers", torch_dtype=torch.float16
).to("cuda")
然后,将管道中的 transformer
和 vae
组件的内存布局更改为 torch.channels-last
:
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
最后,编译这些组件并运行推理:
pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
pipeline.vae.decode = torch.compile(pipeline.vae.decode, mode="max-autotune", fullgraph=True)
image = pipeline(prompt="一个宇航员在骑马").images[0]
在 80GB A100 机器上的 基准测试 结果如下:
With torch.compile(): Average inference time: 12.470 seconds.
Without torch.compile(): Average inference time: 20.570 seconds.
内存优化
通过以 8 位加载 T5 文本编码器,你可以在不到 6 GB 的 GPU VRAM 中运行管道。详情请参阅 此脚本。
此外,你可以使用 [~HunyuanDiT2DModel.enable_forward_chunking
] 方法来减少内存使用。前馈分块会在循环中运行 transformer 块中的前馈层,而不是一次性全部运行。这为你提供了一个在内存消耗和推理运行时间之间的权衡。
+ pipeline.transformer.enable_forward_chunking(chunk_size=1, dim=1)
HunyuanDiTPipeline
[[autodoc]] HunyuanDiTPipeline - all - call