Implement DREAM from http://arxiv.org/abs/2312.00210

RossM commented 6 months ago

Describe your changes

This implements DREAM (Diffusion Rectification and Estimation-Adaptive Models) from http://arxiv.org/abs/2312.00210. DREAM helps the model to correct for errors in previous generation steps, at the cost of an additional forward pass during training. Based on my experiments so far it gives greatly improved composition and realism but takes 30%-50% longer per training step.

Example with DREAM:

Control, same training steps and parameters, without DREAM:

Issue ticket number and link (if applicable)

Checklist before requesting a review

[ + ] This is based on the /dev branch (Or a fork of it)
[ + ] This was created or at least validated using a proper IDE
[ + ] I have tested this code and validated any modified functions
[ + ] I have added the appropriate documentation and hint strings if adding or changing a user-facing feature

RossM commented 6 months ago

I am seeing some of the loss of detail mentioned in the paper when using lambda = 1. More testing needed.

RossM commented 6 months ago

Added a slider to control lambda. I named it "detail preservation" to better explain what it does. Still testing to make sure it has the expected effect.

RossM commented 6 months ago

Sorry for marking this as draft so many times, I get nervous whenever a training run fails even if turns out to be for unrelated reasons.

RossM commented 6 months ago

More comparisons. Left: without DREAM, right: with DREAM. Both are trained for 10k steps starting from SD1.5.

comparison 1

comparison 2

d8ahazard / sd_dreambooth_extension