# Saving and loading modified Unet

**URL:** <https://discuss.huggingface.co/t/saving-and-loading-modified-unet/64541>\
**Category:** 🧨 Diffusers\
**Created:** [December 4, 2023, 3:44pm UTC](https://discuss.huggingface.co/t/saving-and-loading-modified-unet/64541 "2023-12-04T15:44:58Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![brsunagad](https://avatars.discourse-cdn.com/v4/letter/b/34f0e0/32.png) [@brsunagad](https://discuss.huggingface.co/u/brsunagad)\
**Post date:** [December 4, 2023, 3:44pm UTC](https://discuss.huggingface.co/t/saving-and-loading-modified-unet/64541/1 "2023-12-04T15:44:58Z")

</div>

Hello I am loading a Unet model and adding additional class-embedding layers like this:  
unet = UNet2DConditionModel.from\_pretrained(  
args.pretrained\_model\_name\_or\_path, subfolder=“unet”, revision=args.revision  
)

```
camera_embedding = ToWeightsDType(TimestepEmbedding(16, 1280), weight_dtype)
unet.class_embedding = camera_embedding

```

These are my saving and loading methods:  
def save\_model\_hook(models, weights, output\_dir):  
for model in models:  
model\_type = str(type(model))

```
            print(f'model type: {model_type}')
            
            # Check for substring matches in the model type
            if 'UNet' in model_type:
                sub_dir = "unet"
            elif 'CLIPTextModel' in model_type:
                sub_dir = "text_encoder"
            else:
                print(f'Unknown model type: {model_type}. Skipping save.')
                continue
            model.save_pretrained(os.path.join(output_dir, sub_dir))

            # make sure to pop weight so that corresponding model is not saved again
            weights.pop()

    def load_model_hook(models, input_dir):
        while len(models) > 0:
            # pop models so that they are not loaded again
            model = models.pop()
            model_type = str(type(model))

            if 'CLIPTextModel' in model_type:
                # load transformers style into model
                load_model = text_encoder_cls.from_pretrained(
                    input_dir, subfolder="text_encoder")
                model.config = load_model.config
            elif 'UNet' in model_type:
                # load diffusers style into model
                load_model = UNet2DConditionModel.from_pretrained(
                    input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)
            else:
                print(f'Unknown model type: {model_type}. Skipping load.')
                continue

            model.load_state_dict(load_model.state_dict())
            del load_model

```

But when I try to load the saved model I get this error:

RuntimeError: Error(s) in loading state\_dict for UNet2DConditionModel:  
Missing key(s) in state\_dict: “class\_embedding.module.linear\_1.weight”, “class\_embedding.module.linear\_1.bias”, “class\_embedding.module.linear\_2.weight”, “class\_embedding.module.linear\_2.bias”.

I think the config is not modified during the save.  
Please let me know how to resolve this issue.

Thank you.

---

<div class="post-metadata">

**Author:** ![brsunagad](https://avatars.discourse-cdn.com/v4/letter/b/34f0e0/32.png) [@brsunagad](https://discuss.huggingface.co/u/brsunagad)\
**Post date:** [December 4, 2023, 7:06pm UTC](https://discuss.huggingface.co/t/saving-and-loading-modified-unet/64541/2 "2023-12-04T19:06:47Z")

</div>

Solved it by  
unet = UNet2DConditionModel.from\_pretrained(model, subfolder=“unet”, class\_embed\_type = “projection”,projection\_class\_embeddings\_input\_dim=16,revision=None, low\_cpu\_mem\_usage=False, device\_map=None)