Eseguendo IF con diffusori 🧨 su un Google Colab Free Tier

Eseguendo IF su Google Colab Free Tier.

TL;DR: Mostriamo come eseguire uno dei modelli open-source più potenti di testo-immagine IF su una versione gratuita di Google Colab con diffusori 🧨.

Puoi esplorare le capacità del modello direttamente nello Spazio Hugging Face.

Immagine compressa dal repository GitHub ufficiale di IF.

Introduzione

IF è un modello di generazione di immagini basato su pixel ed è stato rilasciato alla fine di aprile 2023 da DeepFloyd. L’architettura del modello è fortemente ispirata a Imagen di Google, che è un modello chiuso.

IF ha due vantaggi distinti rispetto ai modelli esistenti di testo-immagine come Stable Diffusion:

  • Il modello opera direttamente nello “spazio dei pixel” (cioè sulle immagini non compresse) invece di eseguire il processo di denoising nello spazio latente come fa Stable Diffusion.
  • Il modello è addestrato sugli output di T5-XXL, un encoder di testo più potente rispetto a CLIP, utilizzato da Stable Diffusion come encoder di testo.

Come risultato, IF è in grado di generare immagini con dettagli ad alta frequenza (ad esempio, volti umani e mani) ed è il primo modello open-source di generazione di immagini che può generare immagini con testo in modo affidabile.

Il lato negativo dell’operare nello spazio dei pixel e l’utilizzo di un encoder di testo più potente è che IF ha un numero significativamente maggiore di parametri. T5, l’UNet di IF per la generazione di testo-immagine e l’UNet di IF per l’upscaling hanno rispettivamente 4.5 miliardi, 4.3 miliardi e 1.2 miliardi di parametri. Rispetto all’encoder di testo e all’UNet di Stable Diffusion 2.1, che hanno solo 400 milioni e 900 milioni di parametri rispettivamente.

Tuttavia, è possibile eseguire IF su hardware per consumatori ottimizzando il modello per un utilizzo a bassa memoria. Mostreremo come fare ciò con i diffusori 🧨 in questo post sul blog.

Nel punto 1), spiegheremo come utilizzare IF per la generazione di testo-immagine, e nei punti 2) e 3), esamineremo le capacità di variazione dell’immagine e di inpainting di IF.

💡 Nota: Qui stiamo scambiando guadagni in memoria con guadagni in velocità per rendere possibile l’esecuzione di IF su una versione gratuita di Google Colab. Se hai accesso a GPU di fascia alta come A100, ti consigliamo di lasciare tutti i componenti del modello sulla GPU per ottenere la massima velocità, come viene fatto nella demo ufficiale di IF.

💡 Nota: Alcune delle immagini più grandi sono state compresse per caricare più velocemente nel formato del blog. Utilizzando il modello ufficiale, dovrebbero avere una qualità ancora migliore!

Cominciamo 🚀!

Capacità di generazione di testo di IF

Indice

  • Accettare la licenza
  • Ottimizzare IF per l’esecuzione su hardware con restrizioni di memoria
  • Risorse disponibili
  • Installare le dipendenze
  • Generazione di testo-immagine
  • Variazione dell’immagine
  • Inpainting

Accettare la licenza

Prima di poter utilizzare IF, è necessario accettare le sue condizioni d’uso. Per farlo:

    1. Assicurati di avere un account Hugging Face e di essere connesso
    1. Accetta la licenza sulla scheda del modello di DeepFloyd/IF-I-XL-v1.0. Accettare la licenza sulla scheda del modello di Stage I accetterà automaticamente anche per gli altri modelli di IF.
    1. Assicurati di aver effettuato l’accesso localmente. Installa huggingface_hub
pip install huggingface_hub --upgrade

esegui la funzione di login in una shell di Python

from huggingface_hub import login

login()

e inserisci il tuo token di accesso Hugging Face Hub.

Ottimizzazione di IF per l’esecuzione su hardware con restrizioni di memoria

Lo stato dell’arte dell’Apprendimento Automatico non dovrebbe essere solo nelle mani di pochi eletti. Democratizzare l’Apprendimento Automatico significa rendere i modelli disponibili per l’esecuzione su hardware non solo di ultima generazione.

La comunità del deep learning ha creato strumenti di classe mondiale per eseguire modelli ad alta intensità di risorse su hardware di consumo:

  • 🤗 accelerate fornisce utilità per lavorare con modelli di grandi dimensioni.
  • bitsandbytes rende la quantizzazione a 8 bit disponibile per tutti i modelli PyTorch.
  • 🤗 safetensors garantisce non solo l’esecuzione del codice di salvataggio, ma accelera anche notevolmente il tempo di caricamento dei modelli di grandi dimensioni.

Diffusers integra in modo trasparente le librerie sopra indicate per consentire un’API semplice durante l’ottimizzazione di modelli di grandi dimensioni.

La versione gratuita di Google Colab è limitata sia in termini di RAM CPU (13 GB di RAM) che di VRAM GPU (15 GB di RAM per T4), il che rende difficile l’esecuzione dell’intero modello IF superiore a 10 miliardi!

Analizziamo le dimensioni dei componenti del modello IF in precisione float32:

  • Codificatore di testo T5-XXL: 20 GB
  • Stage 1 UNet: 17,2 GB
  • Stage 2 Super Resolution UNet: 2,5 GB
  • Stage 3 Super Resolution Model: 3,4 GB

Non possiamo eseguire il modello in float32 in quanto i pesi di T5 e Stage 1 UNet sono entrambi più grandi della RAM CPU disponibile.

In float16, le dimensioni dei componenti sono rispettivamente 11 GB, 8,6 GB e 1,25 GB per T5, Stage1 e Stage2 UNet, il che è fattibile per la GPU, ma continuiamo a riscontrare errori di overflow della memoria CPU durante il caricamento di T5 (alcuna CPU è occupata da altri processi).

Pertanto, riduciamo ulteriormente la precisione di T5 utilizzando la quantizzazione a 8 bit di bitsandbytes, che consente di salvare il checkpoint di T5 con soli 8 GB.

Ora che ogni componente si adatta individualmente alla memoria CPU e GPU, dobbiamo assicurarci che i componenti abbiano tutta la memoria CPU e GPU a loro disposizione quando necessario.

Diffusers supporta il caricamento modulare dei singoli componenti, ad esempio possiamo caricare il codificatore di testo senza caricare l’UNet. Questo caricamento modulare garantirà che carichiamo solo il componente di cui abbiamo bisogno in un determinato passaggio del processo per evitare di esaurire la RAM CPU e la VRAM GPU disponibili.

Proviamoci 🚀

Risorse disponibili

La versione gratuita di Google Colab ha circa 13 GB di RAM CPU:

!grep MemTotal /proc/meminfo

MemTotal:       13297192 kB

E un NVIDIA T4 con 15 GB di VRAM:

!nvidia-smi

Sun Apr 23 23:14:19 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   72C    P0    32W /  70W |   1335MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                                
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

Installazione delle dipendenze

Alcune ottimizzazioni possono richiedere versioni aggiornate delle dipendenze. Se riscontri problemi, verifica e aggiorna le versioni.

! pip install --upgrade \
  diffusers~=0.16 \
  transformers~=4.28 \
  safetensors~=0.3 \
  sentencepiece~=0.1 \
  accelerate~=0.18 \
  bitsandbytes~=0.38 \
  torch~=2.0 -q

1. Generazione di testo-immagine

Andremo passo dopo passo attraverso la generazione di testo-immagine con IF utilizzando Diffusers. Spiegheremo brevemente le API e le ottimizzazioni, ma spiegazioni più dettagliate possono essere trovate nella documentazione ufficiale di Diffusers, Transformers, Accelerate e bitsandbytes.

1.1 Carica l’encoder di testo

Caricheremo T5 utilizzando la quantizzazione a 8 bit. Transformers supporta direttamente bitsandbytes tramite il flag load_in_8bit.

Il flag variant="8bit" scaricherà i pesi pre-quantizzati.

Useremo anche il flag device_map per consentire a transformers di spostare i livelli del modello sulla CPU o sul disco. Transformers big modeling supporta mappe di dispositivi arbitrarie, che possono essere utilizzate per caricare separatamente i parametri del modello direttamente sui dispositivi disponibili. Passando "auto" verrà creata automaticamente una mappa dei dispositivi. Consultare la documentazione di transformers per ulteriori informazioni.

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder", 
    device_map="auto", 
    load_in_8bit=True, 
    variant="8bit"
)

1.2 Crea gli embedding del testo

Le API di Diffusers per accedere ai modelli di diffusione sono la classe DiffusionPipeline e le sue sottoclassi. Ogni istanza di DiffusionPipeline è un insieme completamente autonomo di metodi e modelli per l’esecuzione di reti di diffusione. Possiamo sostituire i modelli che utilizza passando istanze alternative come argomenti di parola chiave a from_pretrained.

In questo caso, passiamo None per l’argomento unet, quindi non verrà caricato alcun UNet. Questo ci consente di eseguire la parte di embedding del testo del processo di diffusione senza caricare l’UNet in memoria.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=text_encoder, # passa l'encoder di testo a 8 bit precedentemente istanziato
    unet=None, 
    device_map="auto"
)

IF viene fornito anche con una pipeline di super risoluzione. Salveremo gli embedding del prompt in modo da poterli passare direttamente alla pipeline di super risoluzione. In questo modo la pipeline di super risoluzione può essere caricata senza un encoder di testo.

Invece di un astronauta che cavalca un cavallo, diamogli anche un cartello!

Definiamo un prompt adeguato:

prompt = "una fotografia di un astronauta che cavalca un cavallo tenendo un cartello che dice Pixel's in space"

e passiamolo attraverso il modello T5 quantizzato a 8 bit:

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

1.3 Liberare la memoria

Una volta creati gli embedding del prompt, non abbiamo più bisogno dell’encoder di testo. Tuttavia, è ancora in memoria sulla GPU. Dobbiamo rimuoverlo in modo da poter caricare l’UNet.

È complesso liberare la memoria di PyTorch. Dobbiamo raccogliere i rifiuti degli oggetti Python che puntano alla memoria effettivamente allocata sulla GPU.

Prima, usa la parola chiave Python del per eliminare tutti gli oggetti Python che fanno riferimento alla memoria GPU allocata

del text_encoder
del pipe

Eliminare l’oggetto Python non è sufficiente per liberare la memoria della GPU. La garbage collection è quando la memoria GPU effettiva viene liberata.

Inoltre, chiameremo torch.cuda.empty_cache(). Questo metodo non è strettamente necessario poiché la memoria cuda memorizzata nella cache sarà immediatamente disponibile per ulteriori allocazioni. Svuotando la cache possiamo verificare nell’interfaccia utente di Colab che la memoria sia disponibile.

Utilizzeremo una funzione ausiliaria flush() per liberare la memoria.

import gc
import torch

def flush():
    gc.collect()
    torch.cuda.empty_cache()

e eseguilo

flush()

1.4 Fase 1: Il processo principale di diffusione

Con la memoria GPU ora disponibile, possiamo ricaricare il DiffusionPipeline solo con il modello UNet per eseguire il processo principale di diffusione.

I flag variant e torch_dtype vengono utilizzati dai Diffusers per scaricare e caricare i pesi in formato floating point a 16 bit.

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

Spesso, passiamo direttamente il prompt di testo a DiffusionPipeline.__call__. Tuttavia, abbiamo precedentemente calcolato i nostri embedding di testo che possiamo passare al suo posto.

IF offre anche un processo di diffusione per la super risoluzione. Impostando output_type="pt" verranno restituiti tensori PyTorch grezzi invece di un’immagine PIL. In questo modo possiamo mantenere i tensori PyTorch sulla GPU e passarli direttamente alla pipeline di super risoluzione di stage 2.

Definiamo un generatore casuale e avviamo il processo di diffusione di stage 1.

generator = torch.Generator().manual_seed(1)
image = pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Convertiamo manualmente i tensori grezzi in immagini PIL e diamo un’occhiata al risultato finale. L’output di stage 1 è un’immagine 64×64.

from diffusers.utils import pt_to_pil

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

E ancora una volta, rimuoviamo il puntatore Python e liberiamo la memoria della CPU e della GPU:

del pipe
flush()

1.5 Stage 2: Super Risoluzione 64×64 a 256×256

IF dispone di un processo di diffusione separato per l’upscaling.

Eseguiamo ogni processo di diffusione con una pipeline separata.

La pipeline di super risoluzione può essere caricata con un text encoder se necessario. Tuttavia, di solito avremo embedding di testo pre-calcolati dalla prima pipeline di IF. In tal caso, carichiamo la pipeline senza il text encoder.

Crea la pipeline

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, # nessun utilizzo del text encoder => risparmio di memoria!
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

e avviamola riutilizzando gli embedding di testo pre-calcolati

image = pipe(
    image=image, 
    prompt_embeds=prompt_embeds, 
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Di nuovo possiamo ispezionare i risultati intermedi.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

E ancora una volta, eliminiamo il puntatore Python e liberiamo la memoria

del pipe
flush()

1.6 Stage 3: Super Risoluzione 256×256 a 1024×1024

Il secondo modello di super risoluzione per IF è il precedente x4 Upscaler di Stability AI.

Creiamo la pipeline e carichiamola direttamente sulla GPU con device_map="auto".

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

🧨 diffusers permette di comporre facilmente modelli di diffusione sviluppati indipendentemente, in quanto le pipeline possono essere concatenate. Qui possiamo semplicemente prendere l’output del tensore PyTorch precedente e passarlo alla pipeline di stage 3 come image=image.

💡 Nota: L’Upscaler x4 non utilizza T5 e ha il suo encoder di testo. Pertanto, non possiamo utilizzare gli embedding dei prompt precedentemente creati e invece dobbiamo passare il prompt originale.

pil_image = pipe(prompt, generator=generator, image=image).images

A differenza delle pipeline IF, la filigrana IF non verrà aggiunta per impostazione predefinita agli output della pipeline di upscaling Stable Diffusion x4.

Possiamo invece applicare manualmente la filigrana.

from diffusers.pipelines.deepfloyd_if import IFWatermarker

watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

Visualizza l’immagine di output

pil_image[0]

Et voilà! Una bellissima immagine 1024×1024 su Google Colab di livello gratuito.

Abbiamo mostrato come 🧨 diffusers rende facile decomporre e caricare modularmente modelli di diffusione intensivi in termini di risorse.

💡 Nota: Non consigliamo di utilizzare la configurazione sopra descritta in produzione. La quantizzazione a 8 bit, la deallocazione manuale dei pesi del modello e lo scaricamento su disco scambiano tutti la memoria per il tempo (cioè la velocità di inferenza). Questo può essere particolarmente evidente se la pipeline di diffusione viene riutilizzata. In produzione, consigliamo di utilizzare un A100 da 40 GB con tutti i componenti del modello lasciati sulla GPU. Consulta la demo ufficiale IF.

2. Variazione dell’immagine

Gli stessi checkpoint IF possono essere utilizzati anche per la variazione guidata dal testo dell’immagine e il completamento. Il processo di diffusione centrale è lo stesso della generazione di testo-immagine, tranne che l’immagine iniziale con rumore viene creata dall’immagine da variare o completare.

Per eseguire la variazione dell’immagine, carica gli stessi checkpoint con IFImg2ImgPipeline.from_pretrained() e IFImg2ImgSuperResolution.from_pretrained().

Le API per l’ottimizzazione della memoria sono tutte le stesse!

Liberiamo la memoria dalla sezione precedente.

del pipe
flush()

Per la variazione dell’immagine, partiamo da un’immagine iniziale che vogliamo adattare.

In questa sezione, adatteremo il famoso meme “Slaps Roof of Car”. Scarichiamolo da internet.

import requests

url = "https://i.kym-cdn.com/entries/icons/original/000/026/561/car.jpg"
response = requests.get(url)

e caricamolo in un’immagine PIL

from PIL import Image
from io import BytesIO

original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((768, 512))
original_image

La pipeline di variazione dell’immagine accetta sia immagini PIL che tensori grezzi. Consulta le docstring per una documentazione più approfondita sugli input attesi, qui .

2.1 Codificatore di testo

La variazione dell’immagine è guidata dal testo, quindi possiamo definire un prompt e codificarlo con il Text Encoder di T5.

Di nuovo, carichiamo il codificatore di testo con precisione a 8 bit.

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder", 
    device_map="auto", 
    load_in_8bit=True, 
    variant="8bit"
)

Per la variazione dell’immagine, carichiamo il checkpoint con IFImg2ImgPipeline. Quando si utilizza DiffusionPipeline.from_pretrained(...), i checkpoint vengono caricati nella pipeline predefinita. La pipeline predefinita per l’IF è la text-to-image IFPipeline. Quando si caricano checkpoint con una pipeline non predefinita, la pipeline deve essere esplicitamente specificata.

from diffusers import IFImg2ImgPipeline

pipe = IFImg2ImgPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=text_encoder, 
    unet=None, 
    device_map="auto"
)

Trasformiamo il nostro venditore in un personaggio anime.

prompt = "stile anime"

Come prima, creiamo i text embeddings con T5

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

e liberiamo la memoria della GPU e della CPU.

Prima, rimuoviamo i puntatori Python

del text_encoder
del pipe

e poi liberiamo la memoria

flush()

2.2 Fase 1: Il processo di diffusione principale

Successivamente, carichiamo solo i pesi stage 1 UNet nell’oggetto pipeline, proprio come abbiamo fatto nella sezione precedente.

pipe = IFImg2ImgPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

La pipeline di variazione dell’immagine richiede sia l’immagine originale che i text embeddings.

Possiamo opzionalmente utilizzare l’argomento strength per configurare l’ammontare di variazione. strength controlla direttamente l’ammontare di rumore aggiunto. Una maggiore forza significa più rumore, il che significa più variazione.

generator = torch.Generator().manual_seed(0)
image = pipe(
    image=original_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Controlliamo nuovamente l’immagine intermedia 64×64.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

Sembra buono! Possiamo liberare la memoria e aumentare nuovamente la dimensione dell’immagine.

del pipe
flush()

2.3 Fase 2: Super Risoluzione

Per la super risoluzione, carichiamo il checkpoint con IFImg2ImgSuperResolutionPipeline e lo stesso checkpoint di prima.

from diffusers import IFImg2ImgSuperResolutionPipeline

pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

💡 Nota : La pipeline di super risoluzione della variazione dell’immagine richiede l’immagine generata e l’immagine originale.

Puoi anche utilizzare l’upscaler Stable Diffusion x4 su questa immagine. Sentiti libero di provarlo utilizzando i frammenti di codice nella sezione 1.6.

image = pipe(
    image=image,
    original_image=original_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    generator=generator,
).images[0]
image

Bel lavoro! Liberiamo la memoria e guardiamo le pipeline finali di inpainting.

del pipe
flush()

3. Inpainting

La pipeline di inpainting IF è la stessa della variazione dell’immagine, tranne che viene denoizzata solo un’area selezionata dell’immagine.

Specifichiamo l’area da inpaintare con una maschera dell’immagine.

Mostriamo le incredibili capacità di “generazione di lettere” di IF. Possiamo sostituire questo testo del segno con uno slogan diverso.

Prima scarichiamo l’immagine

import requests

url = "https://i.imgflip.com/5j6x75.jpg"
response = requests.get(url)

e la trasformiamo in un’immagine PIL

from PIL import Image
from io import BytesIO

original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((512, 768))
original_image

Maskeremo il segno in modo da poter sostituire il testo.

Per comodità, abbiamo generato in precedenza la maschera e l’abbiamo caricata in un dataset HF.

Andiamolo a scaricare.

from huggingface_hub import hf_hub_download

mask_image = hf_hub_download("diffusers/docs-images", repo_type="dataset", filename="if/sign_man_mask.png")
mask_image = Image.open(mask_image)

mask_image

💡 Nota: È possibile creare le maschere manualmente creando un’immagine in scala di grigi.

from PIL import Image
import numpy as np

height = 64
width = 64

example_mask = np.zeros((height, width), dtype=np.int8)

# Imposta i pixel mascherati a 255
example_mask[20:30, 30:40] = 255

# Assicurarsi di creare l'immagine in modalità 'L'
# che significa scala di grigi a un solo canale
example_mask = Image.fromarray(example_mask, mode='L')

example_mask

Ora possiamo iniziare l’inpainting 🎨🖌

3.1. Codificatore di testo

Di nuovo, carichiamo prima il codificatore di testo

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder", 
    device_map="auto", 
    load_in_8bit=True, 
    variant="8bit"
)

Questa volta, inizializziamo la pipeline di inpainting IFInpaintingPipeline con i pesi del codificatore di testo.

from diffusers import IFInpaintingPipeline

pipe = IFInpaintingPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=text_encoder, 
    unet=None, 
    device_map="auto"
)

Okay, facciamo pubblicità all’uomo per più strati.

prompt = 'il testo, "basta impilare più strati"'

Dopo aver definito il prompt, possiamo creare i prompt embeddings

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

Come prima, liberiamo la memoria

del text_encoder
del pipe
flush()

3.2 Fase 1: Il processo principale di diffusione

Come prima, carichiamo ora la pipeline della fase 1 solo con UNet.

pipe = IFInpaintingPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

Ora, dobbiamo passare l’immagine di input, l’immagine della maschera e i prompt embeddings.

image = pipe(
    image=original_image,
    mask_image=mask_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Diamo un’occhiata all’output intermedio.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

Sembra buono! Il testo è abbastanza coerente!

Liberiamo la memoria per poter ingrandire l’immagine

del pipe
flush()

3.3 Fase 2: Super Risoluzione

Per la super risoluzione, carica il checkpoint con IFInpaintingSuperResolutionPipeline.

from diffusers import IFInpaintingSuperResolutionPipeline

pipe = IFInpaintingSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

La pipeline di inpainting con super risoluzione richiede l’immagine generata, l’immagine originale, l’immagine della maschera e gli embedding del prompt.

Facciamo un’ultima esecuzione di denoising.

image = pipe(
    image=image,
    original_image=original_image,
    mask_image=mask_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    generator=generator,
).images[0]
image

Bravo, il modello ha generato il testo senza commettere nemmeno un errore di ortografia!

Conclusioni

IF in precisione a virgola mobile a 32 bit utilizza un totale di 40 GB di pesi. Abbiamo mostrato come, utilizzando solo modelli e librerie open source, IF possa essere eseguito su un’istanza gratuita di Google Colab.

L’ecosistema di ML beneficia profondamente dalla condivisione di strumenti e modelli aperti. Questo notebook ha utilizzato modelli di DeepFloyd, StabilityAI e Google. Le librerie utilizzate – Diffusers, Transformers, Accelerate e bitsandbytes – beneficiano tutte di innumerevoli contributori provenienti da diverse organizzazioni.

Un enorme ringraziamento al team di DeepFloyd per la creazione e la condivisione di IF e per il contributo alla democratizzazione dell’apprendimento automatico di qualità 🤗.