Обзор

Вывод в диффузионных моделях требует значительных вычислительных ресурсов, и часто приходится запускать DiffusionPipeline несколько раз, прежде чем результат будет удовлетворительным. На этой странице приведён обзор основных методов оптимизации в Diffusers, их назначение и случаи применения.

Начальный путь

Если модель помещается на одном GPU, начните с этой базовой загрузки. Установите dtype и разместите конвейер на ускорителе. Обращайтесь к выгрузке модели на CPU только в том случае, если памяти не хватает. Также вы можете ускорить вывод, уменьшив количество шагов или выбрав более быстрый планировщик.

Если вы не указываете dtype, Diffusers загружает компоненты в формате float32. Передайте dtype=torch.bfloat16 (или torch.float16, если bfloat16 не поддерживается), а затем переместите конвейер на ускоритель с помощью pipeline.to("cuda").

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu"

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
pipeline(prompt).images[0]

Если пайплайн не помещается в память или ресурсов GPU недостаточно, вместо удержания всех компонентов на GPU вызовите enable_model_cpu_offload(). Он размещает активную модель на GPU, а остальные компоненты оставляет на CPU.

Пропускайте этот шаг, если модель помещается в память. Разгрузка (offloading) замедляет процесс, когда в ней нет необходимости.

Для получения дополнительных вариантов выгрузки данных см. Снижение потребления памяти.

pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.bfloat16,
)
pipeline.enable_model_cpu_offload()

Меньшая задержка при использовании меньшего количества шагов num_inference_steps или более быстрого планировщика, например DPMSolverMultistepScheduler. Обычно это ускоряет генерацию, но может снизить качество изображения по сравнению с более медленным планировщиком, обеспечивающим высокое качество. Дополнительные способы ускорения описаны в разделе Техники оптимизации ниже.

import time
from diffusers import DPMSolverMultistepScheduler

pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)

start_time = time.perf_counter()
image = pipeline(prompt, num_inference_steps=25).images[0]
end_time = time.perf_counter()

print(f"Image generation took {end_time - start_time:.3f} seconds")

Техники оптимизации

Если стартового пути недостаточно, используйте эти техники. Если инференс слишком медленный, начните с torch.compile, кэширования или бэкендов внимания. Если память исчерпана, начните с offloading или квантизации.

Более быстрый вывод:

• torch.compile — компиляция UNet, трансформера или VAE в оптимизированные ядра.
• Региональная компиляция — компиляция повторяющихся блоков для сокращения задержки холодного запуска torch.compile и повторного использования скомпилированных артефактов.
• Кэширование — повторное использование промежуточных результатов между шагами денойзинга, когда важна скорость и можно пожертвовать памятью.
• Бэкенды внимания — Заменяйте реализации внимания в Diffusers через единый API, если именно внимание является узким местом.
• Kernels — Загружайте оптимизированные вычислительные ядра Hub (внимание и пользовательские CUDA-операции, такие как RMSNorm или RoPE), когда вам нужны аппаратно-специфичные ускорения, выходящие за рамки стандартного PyTorch.

Меньше памяти:

• Выгрузка (Offloading) — перемещение неактивных моделей или слоёв на CPU с помощью выгрузки на CPU, выгрузки модели или групповой выгрузки.
• Квантование — загрузка весов меньшего размера для экономии памяти (некоторые бэкенды также ускоряют инференс). GGUF — популярная отправная точка.
• Разбиение VAE и тайлинг VAE — декодируйте большие батчи или изображения высокого разрешения по частям, чтобы снизить пиковое потребление памяти.

Оба:

• Квантование, компиляция и выгрузка — комбинируйте квантование, torch.compile и выгрузку, когда одного метода недостаточно.