Diffusers 文档

Metal Performance Shaders (MPS)

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Metal Performance Shaders (MPS)

🤗 Diffusers 兼容使用 PyTorch mps 设备的 Apple 芯片 (M1/M2 芯片),它使用 Metal 框架来利用 MacOS 设备上的 GPU。您需要具备

  • 配备 Apple 芯片 (M1/M2) 硬件的 macOS 电脑
  • macOS 12.6 或更高版本(推荐 13.0 或更高版本)
  • arm64 版本的 Python
  • PyTorch 2.0(推荐)或 1.13(mps 支持的最低版本)

mps 后端使用 PyTorch 的 .to() 接口将 Stable Diffusion pipeline 移动到您的 M1 或 M2 设备上

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
pipe = pipe.to("mps")

# Recommended if your computer has < 64 GB of RAM
pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image

批量生成多个 prompt 可能会 崩溃 或无法可靠工作。我们认为这与 PyTorch 中的 mps 后端有关。在对此进行调查期间,您应该迭代而不是批量处理。

如果您使用的是 PyTorch 1.13,您需要通过额外的单次传递来“预热” pipeline。这是一个临时解决方法,用于解决首次推理传递产生的结果与后续传递略有不同的问题。您只需要执行此传递一次,并且在仅一步推理后,您可以丢弃结果。

  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"
  # First-time "warmup" pass if PyTorch version is 1.13
+ _ = pipe(prompt, num_inference_steps=1)

  # Results match those from the CPU device after the warmup pass.
  image = pipe(prompt).images[0]

故障排除

M1/M2 性能对内存压力非常敏感。当发生这种情况时,系统会在需要时自动进行交换,这会显着降低性能。

为了防止这种情况发生,我们建议使用注意力切片来减少推理期间的内存压力并防止交换。如果您的计算机系统 RAM 小于 64GB,或者您生成的分辨率大于 512×512 像素的非标准图像,则这一点尤其重要。在您的 pipeline 上调用 enable_attention_slicing() 函数

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%,但我们观察到,除非您拥有 64GB 或更多的 RAM,否则在大多数 Apple 芯片计算机中性能更好

< > 在 GitHub 上更新