Training Pipelines keras_cv/training/*?

[innat-asj]

# current
keras_cv/training/contrastive/
 
# what if?
keras_cv/training/drembooth/
keras_cv/training/texualinversion/

Inspiration: https://huggingface.co/docs/diffusers/training/dreambooth

[jbischof]

Not a bad idea @innat-asj, but the issue is that we haven't trained anything ourselves for Stable Diffusion. What training scripts did you have in mind?

[innat]

This is mostly inspired from diffusers and for current keras_cv/training/contrastive/ pipelines. The key idea is to provide model class i.e. DreamBoothTrainer, StableDiffusionFineTuner in order to fine-tune.

Thoughg I'm still thinking how to formulate it (keras_cv/training/) in a better way, but working with diffusers is pretty convinient as it supprts various pipelines and all interact to each other in most part. Not to mention, diffusers support LoRA for dreambooth, integrate trained texual inversion model to dreambooth's text-encoder pipelines.

[jbischof]

Yes, I have been talking with @sayakpaul about refactoring the SD codebase to make it more modular like the rest of our new API. I think this will require some thought.

[sayakpaul]

Not a bad idea @innat-asj, but the issue is that we haven't trained anything ourselves for Stable Diffusion.

WDYM? There's a Textual Inversion tutorial on keras.io authored by the Keras team.

The results shown in the other two tutorials (DreamBooth and regular fine-tuning) are inline with what the community found out.

My thoughts of trainer modules for diffusion model training:

• I think a standard API might not be able to suit different flavors of fine-tuning Diffusion models (saying from my experience of maintaining a large chunk of https://github.com/huggingface/diffusers/blob/main/examples).
• I think ^ more so because DreamBooth and vanilla text-to-image fine-tuning overlap in some portions of their training methodology.
• But the above doesn't hold true for textual inversion since it only operates with the embedding layer of the text encoder.
• Some techniques (such as Custom Diffusion) only care about optimizing the cross-attention maps and optionally the embedding layers of the text encoder.
• Same goes for prompt-to-prompt.

This is why I believe the API design could easily become very messy for the above-mentioned use cases. If the model / backbone classes themselves are implemented nicely, then overriding the train_step() for a trainer class becomes a piece of cake really (I think this becomes evident when you check out the keras.io tutorials).

So, at best, there could be a base trainer class for Diffusion models that basically implements abstract methods on top of tf.keras.Model, but I cannot immediately envision how that might look like, though.

So off-the-top ideas for making the model / backbone classes more hackable:

• Allow users to easily fetch the attention layers (both self and cross) and access their parameters. This will help you easily unlock training with Custom Diffusion strategy.
• The above will also allow users to implement techniques like Attend and Excite, Pix2Pix Zero easily since all of them operate on the cross-attention maps.
• Support for more / better schedulers (like DPMMultiStepSolver).