stable diffusion fine-tuning

[patil-suraj]

This PR adds fine-tuning script for StableDiffusion

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[minimaxir]

Is it a good idea to finetune the entire unet? It seems like there is a risk of catastrophic forgetting, particularly if the input data is small and/or uniform. (this was an issue when I finetuned ruDALL-E on Pokemon; too long training and it ignored any specificity in the prompts)

It might be worth it to allow parametric freezing of embeddings and/or early layers of the unet as CLI args down the line, depending on how well this finetuning works.

[patil-suraj]

Hey @minimaxir ! That's a good point. We are doing some initial runs now and and we could indeed allow parametric freezing of certain layers if we run into these issues.
Also, some community members have tried fine-tuning the full model and observed that the model can still retain its ‘general’ knowledge. cc @justinpinkney

[justinpinkney]

@minimaxir After too long fine-tuning catastrophic forgetting def does happen, but long enough that you can represent the target domain well before forgetting everything. I also think borrowing some of the training tricks from Dream Booth is probably a good way to try and retain the generality.

[ezhang7423]

Hi, could we get a sample command that we could use to try this script out?

[minimaxir]

@justinpinkney thanks for the info! Will be curious to see how things pan out (and of course thanks to the well-commented code it's pretty easy for the user to manually freeze layers if needed)

[ghost]

I am eager to do some fine-tuning, but I guess this is still non-functional/in early stages? I am only able to produce black images, even after bypassing the safety filters. I will have to do a lot of catching up before I would be able to meaningfully contribute towards progress on this, but I'll keep an eye out on progress for this !

[zcrypt0]

Does the default configuration save a checkpoint on any type of interval in case of interruptions?

I looked through the script and I don't see any config for it, but I'm not familiar with the accelerator training lib.

[patil-suraj]

It's not added yet, but will add support for it in a follow-up PR.

[patrickvonplaten]

This is good to merge for me. @anton-l @pcuenca do you want to take a final look?

[pcuenca]

Looks good to me, I just have some questions about the types of datasets we support.

[j-min]

Hi, thanks for adding the finetuning script!
Could you please guide how I should edit the script to jointly fine-tune both unet and text encoder? I'm not familiar with the accelerator package.

[tengshaofeng]

@patil-suraj @HuggingFaceDocBuilderDev @minimaxir @justinpinkney @ezhang7423 , Thanks for your ideas. I have learned so much from that. I still have a question. How many steps should I finetune the stable diffusion model with 10000images? To observe the loss changing, and when loss is not descend, then it is time to end the training? But I do not think the smaller the loss, the better performs.

[aihao2000]

@patil-suraj I have some doubts as to why model.train() should be added at the beginning of each epoch.