Skip to content

AWS Neuron

Diffusers 功能可在 AWS Inf2 实例 上使用,这些实例是采用 Neuron 机器学习加速器 的 EC2 实例。这些实例旨在提供更好的计算性能(更高的吞吐量,更低的延迟)以及良好的成本效益,使其成为 AWS 用户将扩散模型部署到生产环境的理想选择。

Optimum Neuron 是 Hugging Face 库与 AWS 加速器(包括 AWS Trainium 和 AWS Inferentia)之间的接口。它支持 Diffusers 中的许多功能,并具有类似的 API,因此如果你已经熟悉 Diffusers,那么学习起来会更容易。创建 AWS Inf2 实例后,安装 Optimum Neuron。

bash
python -m pip install --upgrade-strategy eager optimum[neuronx]

我们提供预构建的 Hugging Face Neuron 深度学习 AMI (DLAMI) 和 Optimum Neuron 容器,用于 Amazon SageMaker。建议你正确设置你的环境。

以下示例演示了如何在 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 `# height in pixels of generated image, eg. 768, 1024` \
  --width 1024 `# width in pixels of generated image, eg. 768, 1024` \
  --num_images_per_prompt 1 `# number of images to generate per prompt, defaults to 1` \
  --auto_cast matmul `# cast only matrix multiplication operations` \
  --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
  sd_neuron_xl/

现在使用预编译的 SDXL 模型生成一些图像。

python
>>> from optimum.neuron import NeuronStableDiffusionXLPipeline

>>> 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]
peggy generated by sdxl on inf2

请查看 Optimum Neuron 文档 中更多关于不同用例的指南和示例!