Skip to content

Metal 性能着色器 (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 管道移动到你的 M1 或 M2 设备上:

python
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

如果你使用的是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"
  # 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 像素)的图像,这尤其重要。在您的管道上调用 [~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%,但我们观察到大多数苹果硅计算机的性能更好,除非您拥有 64GB 或更多 RAM。