mosaicml
Preliminary ControlNet PR (WIP)
Adds ControlNet to the diffusion repo for both SDXL and SD2 style models.
Some of the highlights of the work here:
-
Custom callback to handle initializing a ControlNet from a Composer UNet Checkpoint: Sample Workflow: a. We train a custom SDXL style model, with a checkpoint of UNet, VAE, text encoder. b. We want to initialize a ControlNet from this Composer checkpoint c. Challenge: FSDP wrapping happens before the checkpoint is loaded, meaning that we will get NaNs if we load directly d. Solution: Need to all gather the weights before restoring the checkpoint and duplicating weights
-
Class that handles both ControlNet style generation and forward (training passes) -- for both the forward and the generate, this involves generating the control image and feeding directly into the UNet.
-
Model classes for the two architectures that support loading directly from YAML files.
To-dos: Show inference outputs here + add comments
