Allenamento Stable Diffusion con Dreambooth utilizzando Diffusori

Allenamento Stable Diffusion con Dreambooth e Diffusori.

Dreambooth è una tecnica per insegnare nuovi concetti alla Stable Diffusion utilizzando una forma specializzata di fine-tuning. Alcune persone lo stanno utilizzando con alcune loro foto per inserirsi in situazioni fantastiche, mentre altri lo stanno utilizzando per incorporare nuovi stili. 🧨 I Diffusers forniscono uno script di addestramento Dreambooth. Non ci vuole molto tempo per addestrare, ma è difficile selezionare il giusto set di iperparametri ed è facile fare overfitting.

Abbiamo condotto molti esperimenti per analizzare l’effetto di diverse impostazioni in Dreambooth. Questo post presenta le nostre scoperte e alcuni consigli per migliorare i risultati quando si fa fine-tuning della Stable Diffusion con Dreambooth.

Prima di iniziare, si prega di essere consapevoli che questo metodo non dovrebbe mai essere utilizzato per scopi maliziosi, per generare danno in alcun modo, o per impersonare persone senza il loro consenso. I modelli addestrati con esso sono comunque vincolati dalla licenza CreativeML Open RAIL-M che regola la distribuzione dei modelli Stable Diffusion.

Nota: una versione precedente di questo post è stata pubblicata come rapporto W&B.

  • Dreambooth tende a fare overfitting rapidamente. Per ottenere immagini di buona qualità, dobbiamo trovare un punto di equilibrio tra il numero di passaggi di addestramento e il tasso di apprendimento. Consigliamo di utilizzare un basso tasso di apprendimento e aumentare progressivamente il numero di passaggi finché i risultati non sono soddisfacenti.
  • Dreambooth ha bisogno di più passaggi di addestramento per i volti. Nei nostri esperimenti, 800-1200 passaggi hanno funzionato bene utilizzando una dimensione del batch di 2 e un tasso di apprendimento di 1e-6.
  • La conservazione delle informazioni precedenti è importante per evitare l’overfitting durante l’addestramento sui volti. Per altri soggetti, non sembra fare una grande differenza.
  • Se vedete che le immagini generate sono rumorose o la qualità è degradata, probabilmente significa che c’è overfitting. Prima di tutto, provate i passaggi sopra per evitarlo. Se le immagini generate sono ancora rumorose, utilizzate il programma di pianificazione DDIM o eseguite più passaggi di inferenza (~100 hanno funzionato bene nei nostri esperimenti).
  • L’addestramento dell’encoder di testo oltre all’UNet ha un grande impatto sulla qualità. I nostri migliori risultati sono stati ottenuti utilizzando una combinazione di fine-tuning dell’encoder di testo, basso tasso di apprendimento e un numero adeguato di passaggi. Tuttavia, il fine-tuning dell’encoder di testo richiede più memoria, quindi una GPU con almeno 24 GB di RAM è ideale. Utilizzando tecniche come Adam a 8 bit, addestramento fp16 o accumulo del gradiente, è possibile addestrare su GPU da 16 GB come quelle fornite da Google Colab o Kaggle.
  • Il fine-tuning con o senza EMA ha prodotto risultati simili.
  • Non è necessario utilizzare la parola sks per addestrare Dreambooth. Una delle prime implementazioni l’ha utilizzata perché era un token raro nel vocabolario, ma in realtà è una specie di fucile. I nostri esperimenti e quelli di @nitrosocke, ad esempio, mostrano che è ok selezionare termini che usereste naturalmente per descrivere il vostro obiettivo.

Impatto del tasso di apprendimento

Dreambooth fa overfitting molto velocemente. Per ottenere buoni risultati, regolate il tasso di apprendimento e il numero di passaggi di addestramento in modo che abbiano senso per il vostro dataset. Nei nostri esperimenti (dettagliati di seguito), abbiamo fatto fine-tuning su quattro diversi dataset con tassi di apprendimento alti e bassi. In tutti i casi, abbiamo ottenuto risultati migliori con un basso tasso di apprendimento.

Impostazioni degli esperimenti

Tutti i nostri esperimenti sono stati condotti utilizzando lo script train_dreambooth.py con l’ottimizzatore AdamW su 2x 40GB A100s. Abbiamo utilizzato lo stesso seed e mantenuto tutti gli iperparametri uguali tra le esecuzioni, ad eccezione di LR, numero di passaggi di addestramento e l’uso della conservazione delle informazioni precedenti.

Per i primi 3 esempi (oggetti vari), abbiamo fatto fine-tuning del modello con una dimensione del batch di 4 (2 per GPU) per 400 passaggi. Abbiamo utilizzato un alto tasso di apprendimento di 5e-6 e un basso tasso di apprendimento di 2e-6. Non è stata utilizzata alcuna conservazione delle informazioni precedenti.

L’ultimo esperimento cerca di aggiungere un soggetto umano al modello. Abbiamo utilizzato la conservazione delle informazioni precedenti con una dimensione del batch di 2 (1 per GPU), 800 e 1200 passaggi in questo caso. Abbiamo utilizzato un alto tasso di apprendimento di 5e-6 e un basso tasso di apprendimento di 2e-6.

Notate che potete utilizzare Adam a 8 bit, addestramento fp16 o accumulo del gradiente per ridurre i requisiti di memoria ed eseguire esperimenti simili su GPU con 16 GB di memoria.

Giocattolo per gatti

Tasso di apprendimento elevato ( 5e-6 )

Tasso di apprendimento basso ( 2e-6 )

Testa di maiale

Tasso di apprendimento elevato ( 5e-6 ). Notare che gli artefatti di colore sono residui di rumore; eseguire più passaggi di inferenza potrebbe aiutare a risolvere alcuni di questi dettagli.

Tasso di apprendimento basso ( 2e-6 )

Signor Potato Head

Tasso di apprendimento elevato ( 5e-6 ). Notare che gli artefatti di colore sono residui di rumore; eseguire più passaggi di inferenza potrebbe aiutare a risolvere alcuni di questi dettagli.

Tasso di apprendimento basso ( 2e-6 )

Volto umano

Abbiamo cercato di incorporare il personaggio Kramer di Seinfeld in Stable Diffusion. Come precedentemente menzionato, abbiamo allenato per più passaggi con una dimensione del batch più piccola. Tuttavia, i risultati non sono stati eccezionali. Per motivi di brevità, abbiamo omesso queste immagini campione e rimandiamo il lettore alle sezioni successive, in cui l’allenamento dei volti è diventato il focus dei nostri sforzi.

Riepilogo dei risultati iniziali

Per ottenere buoni risultati nell’allenamento di Stable Diffusion con Dreambooth, è importante regolare il tasso di apprendimento e i passaggi di allenamento per il proprio dataset.

  • Tassi di apprendimento elevati e troppi passaggi di allenamento porteranno all’overfitting. Il modello genererà principalmente immagini dai dati di allenamento, indipendentemente dal prompt utilizzato.
  • Tassi di apprendimento bassi e troppi pochi passaggi porteranno all’underfitting: il modello non sarà in grado di generare il concetto che si stava cercando di incorporare.

I volti sono più difficili da allenare. Nei nostri esperimenti, un tasso di apprendimento di 2e-6 con 400 passaggi di allenamento funziona bene per gli oggetti, ma i volti richiedono 1e-6 (o 2e-6 ) con ~1200 passaggi.

La qualità dell’immagine si degrada molto se il modello fa overfitting, e questo accade se:

  • Il tasso di apprendimento è troppo elevato.
  • Si eseguono troppi passaggi di allenamento.
  • Nel caso dei volti, quando non viene utilizzata alcuna preservazione precedente, come mostrato nella sezione successiva.

Utilizzare la preservazione precedente durante l’allenamento dei volti

La preservazione precedente è una tecnica che utilizza immagini aggiuntive della stessa classe che stiamo cercando di allenare come parte del processo di messa a punto. Ad esempio, se cerchiamo di incorporare una nuova persona nel modello, la classe che vorremmo preservare potrebbe essere “persona”. La preservazione precedente cerca di ridurre l’overfitting utilizzando foto della nuova persona combinate con foto di altre persone. La cosa interessante è che possiamo generare quelle immagini di classe aggiuntive utilizzando il modello Stable Diffusion stesso! Lo script di allenamento si occupa automaticamente di ciò se lo si desidera, ma è anche possibile fornire una cartella con le proprie immagini di preservazione precedente.

Preservazione precedente, 1200 passaggi, lr= 2e-6 .

Nessuna preservazione precedente, 1200 passaggi, lr= 2e-6 .

Come puoi vedere, i risultati sono migliori quando viene utilizzata la preservazione precedente, ma ci sono comunque macchie rumorose. È ora di utilizzare alcuni trucchi aggiuntivi!

Effetto degli schedulers

Negli esempi precedenti, abbiamo utilizzato lo scheduler PNDM per campionare le immagini durante il processo di inferenza. Abbiamo osservato che quando il modello sovraadatta, DDIM di solito funziona molto meglio rispetto a PNDM e LMSDiscrete. Inoltre, la qualità può essere migliorata eseguendo l’inferenza per più passaggi: 100 sembra essere una buona scelta. I passaggi aggiuntivi aiutano a risolvere alcune delle patch di rumore nei dettagli dell’immagine.

PNDM, volto di Kramer

LMSDiscrete, volto di Kramer. I risultati sono terribili!

DDIM, volto di Kramer. Molto meglio

Un comportamento simile può essere osservato per altri soggetti, sebbene in misura minore.

PNDM, Potato Head

LMSDiscrete, Potato Head

DDIM, Potato Head

Raffinamento dell’encoder di testo

Il documento originale di Dreambooth descrive un metodo per raffinare il componente UNet del modello, ma mantiene l’encoder di testo congelato. Tuttavia, abbiamo osservato che il raffinamento dell’encoder produce risultati migliori. Abbiamo sperimentato questo approccio dopo averlo visto utilizzato in altre implementazioni di Dreambooth e i risultati sono sorprendenti!

Encoder di testo congelato

Encoder di testo raffinato

Il raffinamento dell’encoder di testo produce i migliori risultati, specialmente con i volti. Genera immagini più realistiche, è meno incline all’overfitting e raggiunge anche una migliore interpretazione della richiesta, essendo in grado di gestire richieste più complesse.

Epilogo: Inversione testuale + Dreambooth

Abbiamo anche eseguito un ultimo esperimento in cui abbiamo combinato l’inversione testuale con Dreambooth. Entrambe le tecniche hanno un obiettivo simile, ma i loro approcci sono diversi.

In questo esperimento abbiamo prima eseguito l’inversione testuale per 2000 passaggi. Da quel modello, abbiamo quindi eseguito Dreambooth per ulteriori 500 passaggi utilizzando un learning rate di 1e-6. Questi sono i risultati:

Pensiamo che i risultati siano molto migliori rispetto a fare semplicemente Dreambooth, ma non così buoni come quando raffiniamo l’intero encoder di testo. Sembra copiare un po’ lo stile delle immagini di addestramento, quindi potrebbe attenersi troppo ad esse. Non abbiamo approfondito questa combinazione, ma potrebbe essere un’interessante alternativa per migliorare Dreambooth e comunque adattare il processo a una GPU da 16 GB. Sentiti libero di esplorare e raccontarci dei tuoi risultati!