DiffusionPipeline

Модели диффузии состоят из нескольких компонентов, таких как UNet или диффузионные трансформеры (DiT), текстовые энкодеры, вариационные автоэнкодеры (VAE) и планировщики. DiffusionPipeline объединяет все эти компоненты в единый простой в использовании API, не теряя при этом гибкости для модификации отдельных компонентов.

В этом руководстве показано, как загрузить DiffusionPipeline.

Загрузка конвейера

DiffusionPipeline — это базовый класс конвейера, который автоматически выбирает и возвращает экземпляр подкласса конвейера модели, например QwenImagePipeline, сканируя файл model_index.json в поисках имени класса.

Передайте идентификатор модели в from_pretrained(), чтобы загрузить конвейер.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

Каждая модель имеет собственный подкласс пайплайна, наследующий от DiffusionPipeline. Как правило, подкласс имеет узкую специализацию и предназначен для конкретной задачи. Пример приведён в таблице ниже.

Вы можете использовать подкласс напрямую, передав идентификатор модели в from_pretrained().

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

См. документацию Формат одного файла, чтобы узнать, как загружать модели в формате одного файла.

Локальные конвейеры

Пайплайны также можно запускать локально. Используйте snapshot_download для загрузки репозитория модели.

from huggingface_hub import snapshot_download

local_dir = snapshot_download(repo_id="Qwen/Qwen-Image")

Передайте этот путь в from_pretrained(), чтобы загрузить его.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  local_dir, dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

Метод from_pretrained() не скачивает файлы с Hub, если обнаруживает локальный путь. Однако это также означает, что он не скачает и не закэширует какие-либо обновления, внесённые в модель.

Типы данных Pipeline

Используйте аргумент dtype в from_pretrained(), чтобы загрузить модель с конкретным типом данных. Это позволяет загружать разные модели с разной точностью. Например, загрузка большой трансформерной модели в полупrecisии уменьшает требуемый объём памяти.

Передайте типы данных для каждой модели в виде словаря в dtype. Используйте ключ default, чтобы задать тип данных по умолчанию. Если модель отсутствует в словаре и default не указан, она загружается с полной точностью (torch.float32).

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image",
  dtype={"transformer": torch.bfloat16, "default": torch.float16},
)
print(pipeline.transformer.dtype, pipeline.vae.dtype)

Вам не нужно использовать словарь, если вы загружаете все модели в одном и том же типе данных.

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image", dtype=torch.bfloat16
)
print(pipeline.transformer.dtype, pipeline.vae.dtype)

Размещение на устройствах

Аргумент device_map размещает пайплайн на устройствах. Во время выполнения Diffusers принимает значения "balanced", "cpu" и строку ускорителя для вашей машины (обычно "cuda", "mps" или "xpu").

Используйте аргумент max_memory в from_pretrained(), чтобы задать максимальный объём памяти, выделяемый на каждом устройстве. По умолчанию Diffusers использует весь доступный объём памяти.

import torch
from diffusers import DiffusionPipeline

max_memory = {0: "16GB", 1: "16GB"}
pipeline = DiffusionPipeline.from_pretrained(
  "Qwen/Qwen-Image",
  dtype=torch.bfloat16,
  device_map="balanced",
  max_memory=max_memory,
)

Атрибут hf_device_map позволяет получить доступ к device_map и просмотреть его.

print(pipeline.hf_device_map)
# {'transformer': 1, 'vae': 1, 'text_encoder': 0}

Сбросьте device_map конвейера с помощью метода reset_device_map(). Это необходимо, если вы хотите использовать методы, такие как .to(), enable_sequential_cpu_offload() и enable_model_cpu_offload().

pipeline.reset_device_map()

Параллельная загрузка

Крупные модели часто разбиваются на более мелкие файлы, чтобы их было проще загружать. Diffusers поддерживает параллельную загрузку фрагментов для ускорения процесса загрузки.

Установите HF_ENABLE_PARALLEL_LOADING в значение "YES", чтобы включить параллельную загрузку шардов.

Аргумент device_map следует установить в значение "cuda", чтобы предварительно выделить большой блок памяти в соответствии с размером модели. Это значительно сокращает время загрузки модели, поскольку прогрев аллокатора памяти позволяет избежать множества мелких обращений к аллокатору на более поздних этапах.

import os
import torch
from diffusers import DiffusionPipeline

os.environ["HF_ENABLE_PARALLEL_LOADING"] = "YES"

pipeline = DiffusionPipeline.from_pretrained(
  "Wan-AI/Wan2.2-I2V-A14B-Diffusers", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

Замена моделей в конвейере

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

Приведённый ниже пример использует более стабильную версию VAE.

import torch
from diffusers import DiffusionPipeline, AutoModel

vae = AutoModel.from_pretrained(
  "madebyollin/sdxl-vae-fp16-fix", dtype=torch.float16
)

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

Повторное использование моделей в нескольких конвейерах

При работе с несколькими конвейерами, использующими одну и ту же модель, метод from_pipe() позволяет переиспользовать модель, а не загружать её заново каждый раз. Это позволяет использовать несколько конвейеров без увеличения потребления памяти.

Объём используемой памяти определяется пайплайном с наибольшими требованиями к памяти, независимо от количества пайплайнов.

The example below loads a pipeline and then loads a second pipeline with from_pipe() to use perturbed-attention guidance (PAG) to improve generation quality.

Используйте AutoPipelineForText2Image, так как DiffusionPipeline не поддерживает PAG. Подробнее см. в документации AutoPipeline.

import torch
from diffusers import AutoPipelineForText2Image

pipeline_sdxl = AutoPipelineForText2Image.from_pretrained(
  "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
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
"""
image = pipeline_sdxl(prompt).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Max memory reserved: 10.47 GB

Установите enable_pag=True во втором пайплайне, чтобы включить PAG. Второй пайплайн использует тот же объём памяти, так как разделяет веса модели с первым.

pipeline = AutoPipelineForText2Image.from_pipe(
  pipeline_sdxl, enable_pag=True
)
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
"""
image = pipeline(prompt).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
# Max memory reserved: 10.47 GB

Пайплайны, созданные с помощью from_pipe(), используют одни и те же модели и state. Изменение состояния модели в одном пайплайне влияет на все остальные пайплайны, которые используют эту же модель.

Some methods may not work correctly on pipelines created with from_pipe(). For example, enable_model_cpu_offload() relies on a unique model execution order, which may differ in the new pipeline. To ensure proper functionality, reapply these methods on the new pipeline.

Дальнейшие шаги

После загрузки конвейера обычно настраивают расписание удаления шума, меняют формат весов или подключают адаптеры.

• Schedulers описывает замену и настройку алгоритма удаления шума.
• Форматы моделей описывает GGUF, контрольные точки в одном файле и другие схемы расположения весов.
• Сокращение использования памяти описывает выгрузку данных и другие инструменты для работы с памятью.
• LoRA описывает загрузку адаптеров в конвейер.
• Устаревшие контрольные точки описывает модуль проверки безопасности для старых контрольных точек Stable Diffusion.