[Community Pipeline] Checkpoint merging

[osanseviero]

Intro

Community Pipelines are introduced in diffusers==0.4.0 with the idea of allowing the community to quickly add, integrate, and share their custom pipelines on top of diffusers.

You can find a guide about Community Pipelines here. You can also find all the community examples under examples/community/. If you have questions about the Community Pipelines feature, please head to the parent issue.

Idea: Checkpoint Merging

This pipeline aims to merge two checkpoints into one via interpolation.

[osanseviero]

@apolinario any resources you would suggest for this pipeline?

[apolinario]

Yes, so checkpoint merging is an idea implemented in the AUTOMATIC1111/stable-diffusion-webui repo, it interpolates the weights of models (can be fine-tuned or different versions of models), potentially creating cool results:

This is how it is implemented in the original repo

[Abhinay1997]

Would the following be a good start ?

# STEP 1:
# Verify that the checkpoints have the same dimensions/modules etc.
# STEP 2:
# Find the mergeable modules from both the checkpoints. ( vae, unet, safety checker etc.. )
# STEP 3:
# For each mergeable component,  use the interpolation technique on the component weights and update the weights
# STEP 4:
# Return a pipeline with the merged weights.

So you would run it like this:-

pipe = CheckpointMergerPipeline.from_pretrained(chkpt0 =  "sample/checkpoint-1", chkpt1 = "sample/checkpoint-2", alpha = 0.2, interp = "sigmoid")

pipe.save_pretrained()

prompt = "A cat riding a skateboard in an 18th century street at night, moon in the background"

pipe.to("cuda")
pipe(prompt).images[0]

What do you think ? @apolinario @osanseviero

[patrickvonplaten]

That seems sensible to me

[vvvm23]

Hey @Abhinay1997, are you still working on this? If not, @patrickvonplaten may I please give this a shot? 🙂

[Abhinay1997]

Hi @vvvm23, Currently working on a solution. Give me 24 hrs to make a PR. I'll let you know if my approach fails and maybe you can pick it up ?

[vvvm23]

No worries! Please take your time!

[Abhinay1997]

@patrickvonplaten, @apolinario

I ran into the following issue and wanted your thoughts:-

When building a custom pipeline inheriting from DiffusionPipeline, for checkpoint merging I have to do it this way:-

pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline = "checkpoint_merger_pipeline")
## Merge v1.4 weights with v1.2 weights of stable diffusion.
pipe.merge("CompVis/stable-diffusion-v1-2")

The problem here is that the custom_pipeline is loaded via the from_pretrained method of DiffusionPipeline. If we have to ensure that any kind of checkpoints are mergeable, we need to keep the modules dynamic. However the from_pretrained
method expects the custom_pipeline class to declare the kwargs in advance and results in an empty pipeline otherwise.

One solution is to pass the original checkpoint again in the merge method like:-

pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline = "checkpoint_merger_pipeline")
## Merge v1.4 weights with v1.2 weights of stable diffusion.
## SEE EXTRA ARGS HERE
pipe.merge("CompVis/stable-diffusion-v1-4", "CompVis/stable-diffusion-v1-2")

But it would result in the from_pretrained call being redundant other than to return an instance of the custom_pipeline class and result in confusion.

What would you suggest ?

[patrickvonplaten]

Really up to you as it's a community pipeline :-) In general we cannot assume to be able to merge weights at a low level, so I'm also not 100% sure of the usefulness here - @apolinario do you have a good idea?

[vvvm23]

Just my random thoughts, but why make this a pipeline instead of something functional? For example, a function with args module, ckpt1, ckpt2, alpha?

[Abhinay1997]

Update:- I've reached a point in the solution where validation of the checkpoints is complete and the individual interpolation steps being complete. The only thing that needs to be done is returning a DiffusionPipeline object with the updated weights.

Should be done by the weekend.

[Abhinay1997]

Hey team. Sorry it took a while. But here's the colab notebook with my code. Need your help in figuring out how to cut down on memory issues. Currently it crashes for the scenario in the notebook ( 12 GB RAM ) Will try it out on Kaggle ( 16 GB RAM ) and see if it's any better.