DiffusionPipeline

Les modèles de diffusion sont composés de plusieurs éléments tels que des UNets ou des transformeurs de diffusion (DiTs), des encodeurs de texte, des auto-encodeurs variationnels (VAE) et des planificateurs. Le DiffusionPipeline encapsule l'ensemble de ces composants dans une API unique et facile à utiliser, sans renoncer à la flexibilité de modifier ses composants.

Ce guide vous montre comment charger un DiffusionPipeline.

Chargement d'un pipeline

DiffusionPipeline est une classe de pipeline de base qui sélectionne et renvoie automatiquement une instance d’une sous-classe de pipeline du modèle, comme QwenImagePipeline, en parcourant le fichier model_index.json à la recherche du nom de la classe.

Passez un identifiant de modèle à from_pretrained() pour charger un pipeline.

import torch
from diffusers import DiffusionPipeline

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

Chaque modèle possède un sous-classe de pipeline spécifique qui hérite de DiffusionPipeline. Un sous-classe a généralement un champ d’application restreint et est spécifique à une tâche. Consultez le tableau ci-dessous pour un exemple.

Vous pouvez utiliser la sous-classe directement en passant un identifiant de modèle à 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"
)

Consultez la documentation Format de fichier unique pour apprendre à charger des modèles en fichier unique.

Pipelines locaux

Les pipelines peuvent également être exécutés localement. Utilisez snapshot_download pour télécharger un dépôt de modèle.

from huggingface_hub import snapshot_download

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

Passez ce chemin à from_pretrained() pour le charger.

import torch
from diffusers import QwenImagePipeline

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

La méthode from_pretrained() ne télécharge pas les fichiers depuis le Hub lorsqu’elle détecte un chemin local. Cela signifie également qu’elle ne téléchargera ni ne mettra en cache les mises à jour apportées au modèle.

Types de données du pipeline

Utilisez l'argument dtype dans from_pretrained() pour charger un modèle avec un type de données spécifique. Cela permet de charger différents modèles dans différentes précisions. Par exemple, charger un grand modèle transformer en demi-précision réduit la mémoire requise.

Transmettez le type de données de chaque modèle sous forme de dictionnaire à dtype. Utilisez la clé default pour définir le type de données par défaut. Si un modèle n’est pas présent dans le dictionnaire et que default n’est pas fourni, il est chargé en précision complète (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)

Vous n’avez pas besoin d’utiliser un dictionnaire si vous chargez tous les modèles dans le même type de données.

import torch
from diffusers import QwenImagePipeline

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

Placement sur l'appareil

L'argument device_map place un pipeline sur des appareils. À l'exécution, Diffusers accepte "balanced", "cpu" et la chaîne de caractères de l'accélérateur de la machine utilisée (généralement "cuda", "mps" ou "xpu").

Utilisez l'argument max_memory dans from_pretrained() pour allouer une quantité maximale de mémoire à utiliser sur chaque appareil. Par défaut, Diffusers utilise la quantité maximale disponible.

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,
)

L'attribut hf_device_map vous permet d'accéder et de consulter le device_map.

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

Réinitialisez le device_map d’un pipeline avec la méthode reset_device_map(). C’est nécessaire si vous souhaitez utiliser des méthodes telles que .to(), enable_sequential_cpu_offload() et enable_model_cpu_offload().

pipeline.reset_device_map()

Chargement parallèle

Les grands modèles sont souvent fragmentés en fichiers plus petits afin de faciliter leur chargement. Diffusers prend en charge le chargement parallèle des fragments pour accélérer le processus de chargement.

Définissez HF_ENABLE_PARALLEL_LOADING sur "YES" pour activer le chargement parallèle des shards.

L'argument device_map doit être défini sur "cuda" afin de pré-allouer un grand bloc de mémoire en fonction de la taille du modèle. Cela réduit considérablement le temps de chargement du modèle, car le réchauffement de l'allocationnaire de mémoire évite ainsi de nombreux petits appels à l'allocationnaire par la suite.

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"
)

Remplacer des modèles dans un pipeline

DiffusionPipeline est flexible et permet de charger d’autres modèles ou planificateurs. Vous pouvez expérimenter avec différents planificateurs pour optimiser la vitesse ou la qualité de génération, et remplacer les modèles par des versions plus performantes.

L'exemple ci-dessous utilise une version plus stable du 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"
)

Réutilisation de modèles dans plusieurs pipelines

Lorsque vous travaillez avec plusieurs pipelines utilisant le même modèle, la méthode from_pipe() permet de réutiliser le modèle au lieu de le recharger à chaque fois. Cela vous permet d'utiliser plusieurs pipelines sans augmenter l'utilisation de la mémoire.

La consommation mémoire est déterminée par le pipeline ayant la plus forte exigence mémoire, indépendamment du nombre de pipelines.

L’exemple ci-dessous charge un pipeline, puis charge un second pipeline avec from_pipe() afin d’utiliser perturbed-attention guidance (PAG) pour améliorer la qualité de la génération.

Utilisez AutoPipelineForText2Image car DiffusionPipeline ne prend pas en charge PAG. Consultez la documentation de AutoPipeline pour en savoir plus.

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

Définissez enable_pag=True dans le second pipeline pour activer PAG. Le second pipeline utilise la même quantité de mémoire car il partage les poids du modèle avec le premier.

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

Les pipelines créés par from_pipe() partagent les mêmes modèles et state. Modifier l'état d'un modèle dans un pipeline affecte tous les autres pipelines qui partagent le même modèle.

Certaines méthodes peuvent ne pas fonctionner correctement sur les pipelines créés avec from_pipe(). Par exemple, enable_model_cpu_offload() repose sur un ordre d'exécution unique du modèle, qui peut différer dans le nouveau pipeline. Pour garantir un fonctionnement approprié, réappliquez ces méthodes sur le nouveau pipeline.

Étapes suivantes

Une fois qu'un pipeline est chargé, vous ajustez généralement le calendrier de débruitage, changez les formats de poids ou attachez des adaptateurs.

• Schedulers couvre l'échange et la configuration de l'algorithme de débruitage.
• Formats de modèles couvre GGUF, les points de contrôle en fichier unique et d'autres dispositions de poids.
• Réduire l'utilisation de la mémoire couvre le déchargement et d'autres outils de gestion de la mémoire.
• LoRA couvre le chargement d’adaptateurs sur un pipeline.
• Points de contrôle legacy couvre le vérificateur de sécurité pour les points de contrôle Stable Diffusion plus anciens.