Addestra il tuo ControlNet con diffusori

Addestra il ControlNet con diffusori.

Introduzione

ControlNet è una struttura di rete neurale che consente un controllo fine dei modelli di diffusione aggiungendo condizioni extra. La tecnica è stata introdotta con l’articolo “Adding Conditional Control to Text-to-Image Diffusion Models” e ha rapidamente conquistato la comunità di diffusione open-source grazie al rilascio di 8 diverse condizioni per controllare Stable Diffusion v1-5, tra cui stime di posa, mappe di profondità, bordi di Canny, schizzi e altro ancora.

In questo post del blog andremo attraverso ogni passaggio in dettaglio su come abbiamo addestrato il modello Uncanny Faces, un modello sulle pose facciali basato su volti sintetici 3D (gli uncanny faces sono in realtà una conseguenza non intenzionata, rimanete sintonizzati per vedere come è successo).

Primi passi con l’addestramento del tuo ControlNet per Stable Diffusion

Per addestrare il tuo ControlNet sono necessari 3 passaggi:

• Apprendimento federato utilizzando Hugging Face e Flower
• Accelerazione dell’Inferenza di Diffusione Stabile su CPU Intel
• Inferenza rapida su modelli di linguaggio estesi BLOOMZ su acceleratore Habana Gaudi2

1. Pianificare la tua condizione: ControlNet è abbastanza flessibile da poter gestire Stable Diffusion per molte attività. I modelli pre-addestrati mostrano una vasta gamma di condizioni e la comunità ne ha costruite altre, come la condizionamento su palette di colori pixelati.

2. Costruire il tuo dataset: Una volta decisa una condizione, è il momento di costruire il dataset. Puoi costruire un dataset da zero o utilizzare un sottoinsieme di un dataset esistente. Hai bisogno di tre colonne nel tuo dataset per addestrare il modello: un’immagine di ground truth, un’immagine di conditioning e una prompt.

3. Addestrare il modello: Una volta che il tuo dataset è pronto, è il momento di addestrare il modello. Questa è la parte più facile grazie allo script di addestramento dei diffusori. Avrai bisogno di una GPU con almeno 8GB di VRAM.

1. Pianificare la tua condizione

Per pianificare la tua condizione, è utile pensare a due domande:

1. Che tipo di condizionamento voglio utilizzare?
2. Esiste già un modello esistente che può convertire immagini “regolari” nella mia condizione?

Per il nostro esempio, abbiamo pensato di utilizzare un condizionamento basato sui landmark facciali. Il nostro ragionamento è stato: 1. il ControlNet condizionato dai landmark generali funziona bene. 2. I landmark facciali sono una tecnica abbastanza diffusa e ci sono diversi modelli che calcolano i landmark facciali su immagini regolari. 3. Potrebbe essere divertente “domare” Stable Diffusion per seguire un determinato landmark facciale o imitare la propria espressione facciale.

2. Costruire il tuo dataset

Ok! Quindi abbiamo deciso di fare un condizionamento Stable Diffusion basato sui landmark facciali. Quindi, per preparare il dataset abbiamo bisogno di:

• L’immagine di ground truth: in questo caso, immagini di volti
• L’immagine di conditioning: in questo caso, immagini in cui i landmark facciali sono visualizzati
• La caption: una didascalia che descrive le immagini utilizzate

Per questo progetto, abbiamo deciso di utilizzare il dataset FaceSynthetics di Microsoft: è un dataset che contiene 100.000 volti sintetici. Altri dataset di ricerca sui volti con volti reali come Celeb-A HQ e FFHQ – ma abbiamo deciso di utilizzare volti sintetici per questo progetto.

Il dataset FaceSynthetics sembrava un ottimo punto di partenza: contiene immagini di volti ground truth, landmark facciali annotati nel formato dei 68 landmark di iBUG e un’immagine segmentata del volto.

Perfetto. Giusto? Purtroppo, non proprio. Ricordate la seconda domanda nel passaggio “pianificare la tua condizione” – che avremmo dovuto avere modelli in grado di convertire immagini regolari nel formato di landmark annotato di questo dataset. Risulta che non esista un modello noto che possa trasformare volti nel formato di landmark annotato di questo dataset.

Quindi abbiamo deciso di seguire un’altra strada:

• Utilizzare le vere immagini di volti del dataset FaceSynthetics
• Utilizzare un modello noto in grado di convertire qualsiasi immagine di un volto nel formato dei 68 landmark facciali di iBUG (nel nostro caso abbiamo utilizzato il modello SPIGA SOTA)
• Utilizzare un codice personalizzato che converte i landmark facciali in una maschera illustrata da utilizzare come conditioning_image
• Salvare tutto ciò come un Hugging Face Dataset

Qui puoi trovare il codice utilizzato per convertire le immagini vere del dataset FaceSynthetics in maschere illustrate e salvarle come un Hugging Face Dataset.

Ora, con l’immagine vera image e la conditioning_image nel dataset, manca solo un passaggio: una didascalia per ogni immagine. Questo passaggio è altamente consigliato, ma puoi sperimentare con prompt vuoti e riportare i tuoi risultati. Poiché non avevamo didascalie per il dataset FaceSynthetics, lo abbiamo sottoposto a un captioning BLIP. Puoi controllare il codice utilizzato per la didascalia di tutte le immagini qui

In questo modo, siamo arrivati al nostro dataset finale! Il Face Synthetics SPIGA con didascalie contiene un’immagine vera, una segmentazione e una didascalia per le 100.000 immagini del dataset FaceSynthetics. Siamo pronti per addestrare il modello!

3. Addestramento del modello

Con il nostro dataset pronto, è ora di addestrare il modello! Anche se questa era supposta essere la parte più difficile del processo, con lo script di addestramento di Diffusers, si è rivelata la più facile. Abbiamo utilizzato una singola GPU A100 noleggiata per 1,10 dollari all’ora su LambdaLabs.

La nostra esperienza di addestramento

Abbiamo addestrato il modello per 3 epoche (ciò significa che il batch di 100.000 immagini è stato mostrato al modello 3 volte) e con una dimensione di batch di 4 (ogni passaggio mostra 4 immagini al modello). Questo si è rivelato eccessivo e ha causato un overfitting (quindi ha dimenticato concetti che si discostano leggermente da un volto reale, ad esempio “Shrek” o “un gatto” nel prompt non avrebbe generato un volto di Shrek o un gatto, ma piuttosto una persona, e ha iniziato a ignorare anche gli stili).

Con solo 1 epoca (quindi dopo che il modello ha “visto” 100.000 immagini), si è già adattato a seguire le pose e non ha causato overfitting. Quindi ha funzionato, ma… poiché abbiamo utilizzato il dataset di volti sintetici, il modello ha imparato a generare volti in 3D che sembrano spettrali, invece di volti realistici. Questo ha senso dato che abbiamo utilizzato un dataset di volti sintetici anziché reali, e può essere utilizzato per scopi divertenti/memici. Ecco il modello uncannyfaces_25K.

In questa tabella interattiva puoi giocare con la manopola sottostante per vedere quante iterazioni di addestramento il modello ha completato e come influisce sul processo di addestramento. Intorno alle 15.000 iterazioni, ha iniziato a imparare le pose. E intorno alle 25.000 iterazioni si è consolidato. Qui

Come abbiamo fatto l’addestramento

Tutto ciò che dovevamo fare era installare le dipendenze:

pip install git+https://github.com/huggingface/diffusers.git transformers accelerate xformers==0.0.16 wandb
huggingface-cli login
wandb login 

E quindi eseguire il codice train_controlnet.py

!accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

Suddividiamo alcune delle impostazioni e passiamo anche alcuni consigli di ottimizzazione per utilizzare solo 8GB di VRAM per l’addestramento.

• pretrained_model_name_or_path : Il modello di base Stable Diffusion che si desidera utilizzare (abbiamo scelto v2-1 qui in quanto può renderizzare meglio i volti)
• output_dir : La directory in cui si desidera salvare il modello
• dataset_name : Il dataset che verrà utilizzato per l’addestramento. Nel nostro caso, Face Synthetics SPIGA con didascalie
• conditioning_image_column : Il nome della colonna nel dataset che contiene l’immagine di condizionamento (nel nostro caso spiga_seg)
• image_column : Il nome della colonna nel dataset che contiene l’immagine di riferimento (nel nostro caso image)
• caption_column : Il nome della colonna nel dataset che contiene la didascalia dell’immagine (nel nostro caso image_caption)
• resolution : La risoluzione delle immagini di condizionamento e di riferimento (nel nostro caso 512x512)
• learning_rate : Il tasso di apprendimento. Abbiamo scoperto che 1e-5 funziona bene per questi esempi, ma è possibile sperimentare con valori diversi compresi tra 1e-4 e 2e-6, ad esempio.
• validation_image : Questo serve per dare uno sguardo durante l’addestramento! Le immagini di convalida verranno eseguite per ogni quantità di validation_steps in modo da poter vedere come sta procedendo l’addestramento. Inserisci qui un percorso locale per un numero arbitrario di immagini di condizionamento
• validation_prompt : Una frase da eseguire insieme all’immagine di convalida. Può essere qualsiasi cosa che possa testare se il modello si sta addestrando bene
• train_batch_size : Questa è la dimensione del batch di addestramento da adattare alla GPU. Possiamo permetterci 4 grazie a un’unità A100, ma se si dispone di una GPU con una VRAM più bassa, si consiglia di ridurre questo valore a 1.
• num_train_epochs : Ogni epoca corrisponde a quante volte le immagini nel set di addestramento verranno “viste” dal modello. Abbiamo sperimentato con 3 epoche, ma si è scoperto che i migliori risultati richiedevano poco più di 1 epoca, con 3 epoche il nostro modello si sovrapponeva.
• checkpointing_steps : Salva un checkpoint intermedio ogni x passaggi (nel nostro caso 5000). Ogni 5000 passaggi, veniva salvato un checkpoint intermedio.
• validation_steps : Ogni x passaggi, vengono eseguiti validation_prompt e validation_image.
• report_to : dove segnalare il tuo addestramento. Qui abbiamo usato Weights and Biases, che ci ha fornito questa bella relazione . Ma ridurre la train_batch_size da 4 a 1 potrebbe non essere sufficiente affinché l’addestramento si adatti a una piccola GPU, ecco alcuni parametri aggiuntivi da aggiungere per ogni dimensione della VRAM della GPU:
• push_to_hub : un parametro per caricare il modello addestrato finale nell’Hugging Face Hub.

Adattamento su una GPU con 16GB di VRAM

pip install bitsandbytes

--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam

La combinazione di una dimensione del batch di 1 con 4 passaggi di accumulo del gradiente è equivalente all’utilizzo della dimensione del batch originale di 4 utilizzata nel nostro esempio. Inoltre, abbiamo abilitato il checkpointing del gradiente e l’Adam a 8 bit per ulteriori risparmi di memoria.

Adattamento su una GPU con 12GB di VRAM

--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam
--set_grads_to_none

Addestramento su una GPU con 8GB di VRAM

Si prega di seguire la nostra guida qui

4. Conclusioni!

Questa esperienza di addestramento di un ControlNet è stata molto divertente. Abbiamo addestrato con successo un modello che può seguire le posizioni reali del viso – tuttavia ha imparato a creare volti 3D insoliti invece di volti 3D reali perché è stato addestrato su questo dataset, il che ha il suo fascino e il suo stile.

Prova il nostro Hugging Face Space:

Per quanto riguarda i prossimi passi per noi – al fine di creare volti dall’aspetto realistico, pur non utilizzando ancora un dataset di volti reali, un’idea è far passare l’intero dataset FaceSynthetics attraverso Stable Diffusion Image2Image, convertendo i volti dall’aspetto 3D in volti dall’aspetto realistico, e poi addestrare un altro ControlNet.

E rimanete sintonizzati, perché presto avremo un evento di addestramento di ControlNet! Seguite Hugging Face su Twitter o unitevi al nostro Discord per rimanere aggiornati su questo.