Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Preliminary ControlNet PR (WIP) #163

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

rishab-partha
Copy link
Contributor

Adds ControlNet to the diffusion repo for both SDXL and SD2 style models.

Some of the highlights of the work here:

  1. 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

  2. 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.

  3. Model classes for the two architectures that support loading directly from YAML files.

To-dos:
Show inference outputs here + add comments

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant