How to merge a checkpoint with the main model ?

[simbrams]

Hey community,

I've trained my own stable diffusion model, however it turns out the best version of it is a checkpoint that I have to load separetely.

But the optimizer.bin of the checkpoint is 6GB heavy, so I would like to know If I can "merge" this configuration a one single model to avoid multiple loads ?

Here is an example of my code:

from diffusers import StableDiffusionPipeline
import torch

device = "cuda"

# load model
model_path = "me/my-model"
pipe = StableDiffusionPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
)

# load checkpoint
pipe.unet.load_attn_procs(
    model_path, 
    subfolder="checkpoint-10000", 
    weight_name="optimizer.bin"
)

I tried to pipe.save_pretrained() but the model created is incomplete, unet is half the size it should be.

Any idea how to make it one single model or make my checkpoint the "main" model ?

thank you in advance for your time!

[tolgacangoz]

Hey Simon,
When you set pipe.save_pretrained(..., variant='fp32'), do you get unet as you expect?

[simbrams]

Hey Simon, When you set pipe.save_pretrained(..., variant='fp32'), do you get unet as you expect?

Hey @StandardAI, thanks a lot for the tip!

I ended up having an error KeyError: 'shortest_edge' when loading the merged model, so I tried to update transformers to the latest version and it worked. (4.41.0)

I also tried to merge the model using default fp16 and it works like a charm 👌

pipe = StableDiffusionPipeline.from_pretrained(
        'me/my-model', 
        torch_dtype=torch.float16,
        local_files_only=True,
        safety_checker=None
    ).to("cuda")

pipe.unet.load_attn_procs(
    'me/my-model', 
    subfolder="checkpoint-10000", 
    weight_name="optimizer.bin",
    cache_dir=cache_path,
)

# Saved the the final model to the cache
pipe.save_pretrained('./models/merged-model-fp16')

Cheers