Diffusione stabile con 🧨 diffusori

'Diffusione stabile con diffusori'

…usando 🧨 Diffusers

Stable Diffusion è un modello di diffusione latente testo-immagine creato dai ricercatori e dagli ingegneri di CompVis, Stability AI e LAION. È allenato su immagini di 512×512 pixel di un sottoinsieme del database LAION-5B. LAION-5B è attualmente il dataset multimodale più grande e liberamente accessibile.

In questo post vogliamo mostrare come utilizzare Stable Diffusion con la libreria 🧨 Diffusers, spiegare come funziona il modello e infine approfondire come diffusers consente di personalizzare il processo di generazione delle immagini.

Nota: È altamente consigliato avere una comprensione di base di come funzionano i modelli di diffusione. Se i modelli di diffusione ti sono completamente nuovi, ti consigliamo di leggere uno dei seguenti articoli:

Il modello di diffusione annotato
Primi passi con 🧨 Diffusers

Ora, cominciamo generando alcune immagini 🎨.

Esecuzione di Stable Diffusion

Licenza

Prima di utilizzare il modello, è necessario accettare la licenza del modello per poter scaricare e utilizzare i pesi. Nota: la licenza non deve più essere accettata esplicitamente tramite l’interfaccia utente.

La licenza è progettata per mitigare i potenziali effetti dannosi di un sistema di apprendimento automatico così potente. Richiediamo agli utenti di leggere attentamente l’intera licenza. Qui offriamo un riassunto:

Non puoi utilizzare il modello per produrre o condividere intenzionalmente output o contenuti illegali o dannosi,
Non rivendichiamo alcun diritto sugli output che generi, sei libero di usarli e sei responsabile del loro utilizzo che non deve violare le disposizioni stabilite nella licenza, e
Puoi ridistribuire i pesi e utilizzare il modello commercialmente e/o come servizio. Se lo fai, ti preghiamo di essere consapevole che devi includere le stesse restrizioni d’uso presenti nella licenza e condividere una copia del CreativeML OpenRAIL-M con tutti i tuoi utenti.

Utilizzo

Prima di tutto, devi installare diffusers==0.10.2 per eseguire i seguenti snippet di codice:

pip install diffusers==0.10.2 transformers scipy ftfy accelerate

In questo post utilizzeremo la versione del modello v1-4, ma è possibile utilizzare anche altre versioni del modello come 1.5, 2 e 2.1 con modifiche minime al codice.

Il modello Stable Diffusion può essere eseguito in inferenza con solo un paio di righe utilizzando la pipeline StableDiffusionPipeline. La pipeline configura tutto ciò che serve per generare immagini da testo con una semplice chiamata alla funzione from_pretrained.

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

Se è disponibile una GPU, spostiamolo su di essa!

pipe.to("cuda")

Nota: Se sei limitato dalla memoria della GPU e hai meno di 10 GB di RAM GPU disponibile, assicurati di caricare la StableDiffusionPipeline in precisione float16 anziché nella precisione float32 predefinita come fatto sopra.

Puoi farlo caricando i pesi dal ramo fp16 e dicendo a diffusers di aspettarsi che i pesi siano in precisione float16:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)

Per eseguire la pipeline, basta definire il prompt e chiamare pipe.

prompt = "una fotografia di un astronauta che cavalca un cavallo"

immagine = pipe(prompt).images[0]

# puoi salvare l'immagine con
# image.save(f"astronauta_cavalca_cavallo.png")

Il risultato sarebbe il seguente

Il codice precedente ti darà un’immagine diversa ogni volta che lo esegui.

Se a un certo punto ottieni un’immagine nera, potrebbe essere perché il filtro dei contenuti incorporato nel modello potrebbe aver rilevato un risultato NSFW. Se ritieni che questo non dovrebbe essere il caso, prova a modificare il tuo prompt o utilizza un seed diverso. Infatti, le previsioni del modello includono informazioni su se è stato rilevato NSFW per un determinato risultato. Vediamo come sono:

result = pipe(prompt)
print(result)

{
    'images': [<PIL.Image.Image image mode=RGB size=512x512>],
    'nsfw_content_detected': [False]
}

Se desideri un output deterministico, puoi impostare un seed casuale e passare un generatore alla pipeline. Ogni volta che usi un generatore con lo stesso seed, otterrai lo stesso output dell’immagine.

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]

# puoi salvare l'immagine con
# image.save(f"astronaut_rides_horse.png")

Il risultato sarebbe il seguente:

Puoi modificare il numero di passaggi di inferenza utilizzando l’argomento num_inference_steps.

In generale, i risultati sono migliori quanto più passaggi si utilizzano, tuttavia più passaggi richiedono più tempo per la generazione. Stable Diffusion funziona molto bene con un numero relativamente piccolo di passaggi, quindi consigliamo di utilizzare il numero predefinito di passaggi di inferenza di 50. Se desideri risultati più veloci, puoi utilizzare un numero inferiore. Se desideri risultati potenzialmente di migliore qualità, puoi utilizzare numeri più grandi.

Proviamo a eseguire la pipeline con meno passaggi di denoising.

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, guidance_scale=7.5, num_inference_steps=15, generator=generator).images[0]

# puoi salvare l'immagine con
# image.save(f"astronaut_rides_horse.png")

Nota come la struttura sia la stessa, ma ci sono problemi nella tuta dell’astronauta e nella forma generale del cavallo. Questo mostra che utilizzare solo 15 passaggi di denoising ha significativamente degradato la qualità del risultato della generazione. Come già detto, 50 passaggi di denoising sono di solito sufficienti per generare immagini di alta qualità.

Oltre a num_inference_steps, abbiamo utilizzato un altro argomento di funzione chiamato guidance_scale in tutti gli esempi precedenti. guidance_scale è un modo per aumentare l’aderenza al segnale condizionale che guida la generazione (testo, in questo caso) così come la qualità complessiva del campione. È noto anche come orientamento libero dal classificatore, che in termini semplici costringe la generazione a corrispondere meglio al prompt potenzialmente a scapito della qualità o della diversità dell’immagine. Valori tra 7 e 8.5 sono di solito delle buone scelte per Stable Diffusion. Per impostazione predefinita, la pipeline utilizza un guidance_scale di 7.5.

Se si utilizza un valore molto grande, le immagini potrebbero sembrare buone, ma saranno meno diverse. È possibile approfondire i dettagli tecnici di questo parametro in questa sezione del post.

Successivamente, vediamo come generare più immagini dello stesso prompt contemporaneamente. Per prima cosa, creeremo una funzione image_grid per aiutarci a visualizzarle in modo ordinato in una griglia.

from PIL import Image

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

Possiamo generare più immagini per lo stesso prompt semplicemente utilizzando una lista con lo stesso prompt ripetuto più volte. Inviaremo la lista alla pipeline invece della stringa che abbiamo usato in precedenza.

num_images = 3
prompt = ["una fotografia di un astronauta che cavalca un cavallo"] * num_images

images = pipe(prompt).images

grid = image_grid(images, rows=1, cols=3)

# puoi salvare la griglia con
# grid.save(f"astronaut_rides_horse.png")

Per impostazione predefinita, la diffusione stabile produce immagini di dimensioni 512 × 512 pixel. È molto facile sovrascrivere il valore predefinito utilizzando gli argomenti height e width per creare immagini rettangolari con rapporto di aspetto verticale o orizzontale.

Quando si scelgono le dimensioni dell’immagine, si consiglia quanto segue:

Assicurarsi che height e width siano entrambi multipli di 8.
Andare al di sotto di 512 potrebbe comportare una riduzione della qualità delle immagini.
Superare 512 in entrambe le direzioni comporterà la ripetizione di aree dell’immagine (si perde la coerenza globale).
Il modo migliore per creare immagini non quadrate è utilizzare 512 in una dimensione e un valore superiore a quello nell’altra dimensione.

Eseguiamo un esempio:

prompt = "una foto di un astronauta che cavalca un cavallo"
image = pipe(prompt, height=512, width=768).images[0]

# puoi salvare l'immagine con
# image.save(f"astronaut_rides_horse.png")

Come funziona la Diffusione Stabile?

Dopo aver visto le immagini di alta qualità che la diffusione stabile può produrre, cerchiamo di capire meglio come funziona il modello.

La Diffusione Stabile si basa su un particolare tipo di modello di diffusione chiamato Diffusione Latente, proposto in “High-Resolution Image Synthesis with Latent Diffusion Models”.

In generale, i modelli di diffusione sono sistemi di apprendimento automatico che vengono addestrati a rimuovere gradualmente il rumore gaussiano casuale per ottenere un campione di interesse, come ad esempio un’immagine. Per una panoramica più dettagliata su come funzionano, consulta questo colab.

I modelli di diffusione hanno dimostrato di ottenere risultati all’avanguardia nella generazione di dati di immagine. Tuttavia, uno svantaggio dei modelli di diffusione è che il processo di denoising inverso è lento a causa della sua natura ripetitiva e sequenziale. Inoltre, questi modelli consumano molta memoria perché operano nello spazio dei pixel, che diventa enorme quando si generano immagini ad alta risoluzione. Pertanto, è difficile addestrare questi modelli e utilizzarli per l’elaborazione.

La diffusione latente può ridurre la complessità di memoria e calcolo applicando il processo di diffusione su uno spazio latente di dimensioni inferiori, anziché utilizzare lo spazio dei pixel effettivo. Questa è la differenza chiave tra i modelli di diffusione standard e quelli di diffusione latente: nel caso della diffusione latente, il modello è addestrato a generare rappresentazioni latenti (compresse) delle immagini.

Ci sono tre componenti principali nella diffusione latente.

Un autoencoder (VAE).
Una U-Net.
Un text-encoder, ad esempio il Text Encoder di CLIP.

1. L’autoencoder (VAE)

Il modello VAE ha due parti, un encoder e un decoder. L’encoder viene utilizzato per convertire l’immagine in una rappresentazione latente a bassa dimensionalità, che servirà come input per il modello U-Net. Il decoder, al contrario, trasforma la rappresentazione latente in un’immagine.

Durante l’addestramento della diffusione latente, l’encoder viene utilizzato per ottenere le rappresentazioni latenti (latents) delle immagini per il processo di diffusione in avanti, che applica sempre più rumore ad ogni passo. Durante l’inferenza, le rappresentazioni latenti denoised generate dal processo di diffusione inversa vengono convertite nuovamente in immagini utilizzando il decoder VAE. Come vedremo durante l’inferenza, abbiamo bisogno solo del decoder VAE.

2. La U-Net

La U-Net è composta da una parte di encoder e una parte di decoder, entrambe composte da blocchi ResNet. L’encoder comprime una rappresentazione dell’immagine in una rappresentazione dell’immagine a risoluzione inferiore e il decoder decodifica la rappresentazione dell’immagine a risoluzione inferiore nella rappresentazione dell’immagine originale a risoluzione superiore, che si presume sia meno rumorosa. In modo più specifico, l’output della U-Net predice il residuo del rumore che può essere utilizzato per calcolare la rappresentazione dell’immagine denoised prevista.

Per evitare che la U-Net perda informazioni importanti durante il campionamento a risoluzione inferiore, di solito vengono aggiunte connessioni shortcut tra i ResNet di downsampling dell’encoder e i ResNet di upsampling del decoder. Inoltre, la U-Net di diffusione stabile è in grado di condizionare il suo output su text-embeddings tramite strati di cross-attention. Gli strati di cross-attention sono aggiunti sia alla parte di encoder che alla parte di decoder della U-Net, di solito tra i blocchi ResNet.

3. L’encoder di testo

L’encoder di testo è responsabile della trasformazione del prompt di input, ad esempio “Un astronauta che cavalca un cavallo”, in uno spazio di embedding comprensibile per la U-Net. Di solito si tratta di un encoder basato su transformer semplice che mappa una sequenza di token di input in una sequenza di embedding di testo latenti.

Ispirato da Imagen, Stable Diffusion non addestra l’encoder di testo durante l’addestramento e semplicemente utilizza un encoder di testo già addestrato di CLIP, CLIPTextModel.

Perché la diffusione latente è veloce ed efficiente?

Poiché la diffusione latente opera su uno spazio a bassa dimensione, riduce notevolmente i requisiti di memoria e calcolo rispetto ai modelli di diffusione nello spazio dei pixel. Ad esempio, l’autoencoder utilizzato in Stable Diffusion ha un fattore di riduzione di 8. Ciò significa che un’immagine di forma (3, 512, 512) diventa (3, 64, 64) nello spazio latente, il che richiede 8 × 8 = 64 volte meno memoria.

Ecco perché è possibile generare immagini di dimensione 512 × 512 così rapidamente, anche su GPU Colab da 16 GB!

Stable Diffusion durante l’inferenza

Mettendo tutto insieme, analizziamo ora più da vicino come funziona il modello nell’inferenza illustrando il flusso logico.

Il modello di diffusione stabile prende sia un seed latente che un prompt di testo come input. Il seed latente viene quindi utilizzato per generare rappresentazioni casuali di immagini latenti di dimensione 64 × 64, mentre il prompt di testo viene trasformato in embedding di testo di dimensione 77 × 768 tramite l’encoder di testo di CLIP.

Successivamente, la U-Net denoisa iterativamente le rappresentazioni casuali di immagini latenti mentre viene condizionata dagli embedding di testo. L’output della U-Net, che è il residuo di rumore, viene utilizzato per calcolare una rappresentazione di immagine latente denoisa tramite un algoritmo di scheduling. Per questo calcolo possono essere utilizzati molti diversi algoritmi di scheduling, ognuno con i suoi vantaggi e svantaggi. Per Stable Diffusion, raccomandiamo di utilizzare uno dei seguenti:

Scheduler PNDM (usato per impostazione predefinita)
Scheduler DDIM
Scheduler K-LMS

La teoria su come funziona l’algoritmo di scheduling esce dallo scopo di questo notebook, ma in breve bisogna ricordare che calcolano la rappresentazione di immagine denoisa prevista dalla rappresentazione di rumore precedente e dal residuo di rumore previsto. Per ulteriori informazioni, consigliamo di consultare Elucidating the Design Space of Diffusion-Based Generative Models.

Il processo di denoising viene ripetuto circa 50 volte per recuperare passo dopo passo migliori rappresentazioni di immagini latenti. Una volta completato, la rappresentazione di immagine latente viene decodificata dalla parte decoder dell’autoencoder variazionale.

Dopo questa breve introduzione a Latent e Stable Diffusion, vediamo come fare un uso avanzato della libreria 🤗 Hugging Face diffusers!

Scrittura della propria pipeline di inferenza

Infine, mostriamo come è possibile creare pipeline di diffusione personalizzate con diffusers. Scrivere una pipeline di inferenza personalizzata è un utilizzo avanzato della libreria diffusers che può essere utile per sostituire determinati componenti, come il VAE o lo scheduler spiegato in precedenza.

Ad esempio, mostriamo come utilizzare Stable Diffusion con uno scheduler diverso, ovvero lo scheduler K-LMS di Katherine Crowson aggiunto in questo PR.

Il modello pre-addestrato include tutti i componenti necessari per configurare una pipeline di diffusione completa. Sono archiviati nelle seguenti cartelle:

text_encoder: Stable Diffusion utilizza CLIP, ma altri modelli di diffusione possono utilizzare altri encoder come BERT.
tokenizer. Deve corrispondere a quello utilizzato dal modello text_encoder.
scheduler: l’algoritmo di scheduling utilizzato per aggiungere progressivamente rumore all’immagine durante l’addestramento.
unet: il modello utilizzato per generare la rappresentazione latente dell’input.
vae: modulo dell’autoencoder che utilizzeremo per decodificare le rappresentazioni latenti in immagini reali.

Possiamo caricare i componenti facendo riferimento alla cartella in cui sono stati salvati, utilizzando l’argomento subfolder per from_pretrained.

from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

# 1. Carica il modello dell'autoencoder che verrà utilizzato per decodificare i latenti nello spazio delle immagini.
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

# 2. Carica il tokenizer e l'encoder di testo per tokenizzare e codificare il testo.
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

# 3. Il modello UNet per generare i latenti.
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")

Ora invece di caricare il scheduler predefinito, carichiamo il K-LMS scheduler con alcuni parametri di fitting.

from diffusers import LMSDiscreteScheduler

scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)

Successivamente, spostiamo i modelli sulla GPU.

torch_device = "cuda"
vae.to(torch_device)
text_encoder.to(torch_device)
unet.to(torch_device)

Definiamo ora i parametri che utilizzeremo per generare le immagini.

Si noti che guidance_scale è definito in modo analogo al peso w dell’equazione (2) nell’articolo Imagen. guidance_scale == 1 corrisponde a una guida senza classificatore. Qui lo impostiamo su 7.5 come fatto precedentemente.

A differenza degli esempi precedenti, impostiamo num_inference_steps a 100 per ottenere un’immagine ancora più definita.

prompt = ["una fotografia di un astronauta che cavalca un cavallo"]

height = 512                        # altezza predefinita di Stable Diffusion
width = 512                         # larghezza predefinita di Stable Diffusion

num_inference_steps = 100           # Numero di passi di denoising

guidance_scale = 7.5                # Scala per la guida senza classificatore

generator = torch.manual_seed(0)    # Generatore del seed per creare il rumore latente iniziale

batch_size = len(prompt)

Per prima cosa, otteniamo gli text_embeddings per il prompt passato. Questi embeddings verranno utilizzati per condizionare il modello UNet e guidare la generazione dell’immagine verso qualcosa che dovrebbe assomigliare al prompt di input.

text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")

text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

Otteniamo anche gli embeddings di testo non condizionati per la guida senza classificatore, che sono semplicemente gli embeddings per il token di padding (testo vuoto). Devono avere la stessa forma degli embeddings condizionati text_embeddings (batch_size e seq_length).

max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
    [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

Per la guida senza classificatore, dobbiamo effettuare due passaggi in avanti: uno con l’input condizionato (text_embeddings) e un altro con gli embeddings non condizionati (uncond_embeddings). In pratica, possiamo concatenare entrambi in un’unica batch per evitare di fare due passaggi in avanti.

text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

Successivamente, generiamo il rumore casuale iniziale.

latents = torch.randn(
    (batch_size, unet.in_channels, height // 8, width // 8),
    generator=generator,
)
latents = latents.to(torch_device)

Se esaminiamo i latents in questa fase vedremo che la loro forma è torch.Size([1, 4, 64, 64]), molto più piccola dell’immagine che vogliamo generare. Il modello trasformerà questa rappresentazione latente (rumore puro) in un’immagine 512 × 512 in seguito.

In primo luogo, inizializziamo lo scheduler con il nostro num_inference_steps scelto. Questo calcolerà i valori sigmas e i valori esatti dell’intervallo di tempo da utilizzare durante il processo di denoising.

scheduler.set_timesteps(num_inference_steps)

Lo scheduler K-LMS deve moltiplicare i latents per i suoi valori di sigma. Facciamolo qui:

latents = latents * scheduler.init_noise_sigma

Siamo pronti per scrivere il ciclo di denoising.

from tqdm.auto import tqdm

scheduler.set_timesteps(num_inference_steps)

for t in tqdm(scheduler.timesteps):
    # espandiamo i latents se stiamo utilizzando una guida senza classificatore per evitare di fare due passaggi in avanti.
    latent_model_input = torch.cat([latents] * 2)

    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

    # predici il residuo del rumore
    with torch.no_grad():
        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

    # effettua la guida
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

    # calcola il campione rumoroso precedente x_t -> x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample

Ora utilizziamo il vae per decodificare i latents generati nell’immagine.

# scala e decodifica i latents dell'immagine con il vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

E infine, convertiamo l’immagine in PIL in modo da poterla visualizzare o salvarla.

image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [Image.fromarray(image) for image in images]
pil_images[0]

Siamo passati dall’uso di base di Stable Diffusion utilizzando 🤗 Hugging Face Diffusers a utilizzi più avanzati della libreria, e abbiamo cercato di presentare tutte le parti di un sistema di diffusione moderno. Se ti è piaciuto questo argomento e vuoi saperne di più, ti consigliamo le seguenti risorse:

Il nostro notebook Colab .
Il notebook Introduzione ai Diffusers, che fornisce una panoramica più ampia sui sistemi di diffusione.
L’articolo del blog Annotated Diffusion Model.
Il nostro codice su GitHub, dove saremmo più che felici se lasciassi una ⭐ se diffusers ti è utile!

Citazione:

@article{patil2022stable,
  author = {Patil, Suraj and Cuenca, Pedro and Lambert, Nathan and von Platen, Patrick},
  title = {Stable Diffusion with 🧨 Diffusers},
  journal = {Hugging Face Blog},
  year = {2022},
  note = {[https://huggingface.co/blog/rlhf](https://huggingface.co/blog/stable_diffusion)},
}

Diffusione stabile con 🧨 diffusori