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

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.