Scalable Diffusion Models with Transformers(DiT)由 William Peebles 和 Saining Xie 提出。
论文摘要如下:
我们探索了一类基于 transformer 架构的新型扩散模型。我们训练图像的潜在扩散模型,将常用的 U-Net 骨干网络替换为在潜在补丁上运行的 transformer。我们通过以 Gflops 衡量的前向传播复杂度来分析我们的扩散 Transformer(DiTs)的可扩展性。我们发现,具有更高 Gflops 的 DiT——通过增加 transformer 的深度/宽度或增加输入标记的数量——始终具有更低的 FID。除了具有良好的可扩展性外,我们最大的 DiT-XL/2 模型在类别条件 ImageNet 512x512 和 256x256 基准测试中优于所有先前的扩散模型,在后者上实现了 2.27 的最先进 FID。
原始代码库可在 facebookresearch/dit 找到。
from mindone.diffusers import DiTPipeline, DPMSolverMultistepScheduler
import mindspore as ms
import numpy as np
pipe = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256", mindspore_dtype=ms.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
# pick words from Imagenet class labels
pipe.labels # to print all available words
# pick words that exist in ImageNet
words = ["white shark", "umbrella"]
class_ids = pipe.get_label_ids(words)
generator = np.random.default_rng(33)
output = pipe(class_labels=class_ids, num_inference_steps=25, generator=generator)
image = output[0][0] # label 'white shark'.jpeg)