huggingface / huggingface/diffusers

Merge pipelines and/or checkpoints

Open
#8,977 19 comments 0 reactions 0 assignees View on GitHub
stale
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Is your feature request related to a problem? Please describe.
I can't find a way to merge pipelines and/or checkpoints with each other.

### Describe the solution you'd like.
I'd like to merge different Stable Diffusion models using `diffusers`.

AFAIK, `diffusers` doesn't provide that feature yet.

Right now I merge my LORAs with my models in `diffusers`, then save them to `ckpt`, and then use the script below to merge my models.

It would save me a lot of time if I could just merge straight from a `StableDiffusionXLPipeline` and I could skip the step in between.

But, right now, even just adding this script (or a variation of it) to [diffusers/scripts](https://github.com/huggingface/diffusers/tree/main/scripts) would already be an improvement.

### Describe alternatives you've considered.
I currently use the following code from [eyriewow/merge-models](https://github.com/eyriewow/merge-models/) :

```python
import os
import argparse
import torch
from tqdm import tqdm

parser = argparse.ArgumentParser(description="Merge two models")
parser.add_argument("model_0", type=str, help="Path to model 0")
parser.add_argument("model_1", type=str, help="Path to model 1")
parser.add_argument("--alpha", type=float, help="Alpha value, optional, defaults to 0.5", default=0.5, required=False)
parser.add_argument("--output", type=str, help="Output file name, without extension", default="merged", required=False)
parser.add_argument("--device", type=str, help="Device to use, defaults to cpu", default="cpu", required=False)
parser.add_argument("--without_vae", action="store_true", help="Do not merge VAE", required=False)

args = parser.parse_args()

device = args.device
model_0 = torch.load(args.model_0, map_location=device)
model_1 = torch.load(args.model_1, map_location=device)
theta_0 = model_0["state_dict"]
theta_1 = model_1["state_dict"]
alpha = args.alpha

output_file = f'{args.output}-{str(alpha)[2:] + "0"}.ckpt'

# check if output file already exists, ask to overwrite
if os.path.isfile(output_file):
print("Output file already exists. Overwrite? (y/n)")
while True:
overwrite = input()
if overwrite == "y":
break
elif overwrite == "n":
print("Exiting...")
exit()
else:
print("Please enter y or n")

for key in tqdm(theta_0.keys(), desc="Stage 1/2"):
# skip VAE model parameters to get better results(tested for anime models)
# for anime model,with merging VAE model, the result will be worse (dark and blurry)
if args.without_vae and "first_stage_model" in key:
continue

if "model" in key and key in theta_1:
theta_0[key] = (1 - alpha) * theta_0[key] + alpha * theta_1[key]

for key in tqdm(theta_1.keys(), desc="Stage 2/2"):
if "model" in key and key not in theta_0:
theta_0[key] = theta_1[key]

print("Saving...")

torch.save({"state_dict": theta_0}, output_file)

print("Done!")
```

It usually gets the job done for `ckpt` files that are the same size, but that's it.

Contributor guide

Open the contributing guide

Research direction

The request points to diffusers/scripts and StableDiffusionXLPipeline; start by reviewing those entry points and the existing pipeline and checkpoint handling. Compare the proposed merge behavior with the linked merge-models script, then define the supported inputs and validation needed for direct merging, including what successful output should contain.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.