Repository: huggingface/diffusers

Issue:
Checkpoint Merger Pipeline giving error: "Expected all tensors to be on the same device"
### Describe the bug

Hi,

I'm following the three-way merge example given at https://github.com/huggingface/diffusers/tree/main/examples/community#checkpoint-merger-pipeline, and am getting the following reproducible error with `diffusers 0.11.1`:

```
RuntimeError: Expected all tensors to be on the same device, but found at least 
two devices, cuda:0 and cpu!
```

### Reproduction

This is reproduced by a literal copy of one of the examples given in README.md:
```py
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline="checkpoint_merger")
merged_pipe = pipe.merge(["CompVis/stable-diffusion-v1-4",
                          "hakurei/waifu-diffusion",
                          "prompthero/openjourney"],
                         force = True,
                         interp = "add_difference",
                         alpha = 0.4)
image = merged_pipe('An astronaut riding a horse on Mars').images[0]
image.save('test.png')
```

### Logs

```shell
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /data/lstein/InvokeAI/merge.py:4 in <module>                                 │
│                                                                              │
│   1 from diffusers import DiffusionPipeline                                  │
│   2                                                                          │
│   3 pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4" │
│ ❱ 4 merged_pipe = pipe.merge(["CompVis/stable-diffusion-v1-4","hakurei/waifu │
│   5 image = merged_pipe('An astronaut riding a horse on Mars').images[0]     │
│   6 image.save('test.png')                                                   │
│   7                                                                          │
│                                                                              │
│ /home/lstein/invokeai/.venv/lib/python3.9/site-packages/torch/autograd/grad_ │
│ mode.py:34 in decorate_context                                               │
│                                                                              │
│    31 │   │   @functools.wraps(func)                                         │
│    32 │   │   def decorate_context(*args, **kwargs):                         │
│    33 │   │   │   with self.clone():                                         │
│ ❱  34 │   │   │   │   return func(*args, **kwargs)                           │
│    35 │   │   return cast(F, decorate_context)                               │
│    36 │                                                                      │
│    37 │   def _wrap_generator(self, func):                                   │
│                                                                              │
│ /home/lstein/invokeai/models/modules/diffusers_modules/git/checkpoint_merger │
│ .py:226 in merge                                                             │
│                                                                              │
│   223 │   │   │   │   │   if theta_2:                                        │
│   224 │   │   │   │   │   │   theta_0[key] = theta_func(theta_0[key], theta_ │
│   225 │   │   │   │   │   else:                                              │
│ ❱ 226 │   │   │   │   │   │   theta_0[key] = theta_func(theta_0[key], theta_ │
│   227 │   │   │   │                                                          │
│   228 │   │   │   │   del theta_1                                            │
│   229 │   │   │   │   del theta_2                                            │
│                                                                              │
│ /home/lstein/invokeai/models/modules/diffusers_modules/git/checkpoint_merger │
│ .py:239 in weighted_sum                                                      │
│                                                                              │
│   236 │                                                                      │
│   237 │   @staticmethod                                                      │
│   238 │   def weighted_sum(theta0, theta1, theta2, alpha):                   │
│ ❱ 239 │   │   return ((1 - alpha) * theta0) + (alpha * theta1)               │
│   240 │                                                                      │
│   241 │   # Smoothstep (https://en.wikipedia.org/wiki/Smoothstep)            │
│   242 │   @staticmethod                                                      │
╰──────────────────────────────────────────────────────────────────────────────╯
```
```


### System Info

- `diffusers` version: 0.11.1
- Platform: Linux-5.15.0-56-generic-x86_64-with-glibc2.31
- Python version: 3.9.5
- PyTorch version (GPU?): 2.0.0+gitd1123c9 (True)
- Huggingface_hub version: 0.11.1
- Transformers version: 4.25.1
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no



Discussion:

Resolving patch:
diff --git a/examples/community/checkpoint_merger.py b/examples/community/checkpoint_merger.py
--- a/examples/community/checkpoint_merger.py
+++ b/examples/community/checkpoint_merger.py
@@ -32,6 +32,7 @@ class CheckpointMergerPipeline(DiffusionPipeline):
     """
 
     def __init__(self):
+        self.register_to_config()
         super().__init__()
 
     def _compare_model_configs(self, dict0, dict1):
@@ -167,6 +168,7 @@ def merge(self, pretrained_model_name_or_path_list: List[Union[str, os.PathLike]
         final_pipe = DiffusionPipeline.from_pretrained(
             cached_folders[0], torch_dtype=torch_dtype, device_map=device_map
         )
+        final_pipe.to(self.device)
 
         checkpoint_path_2 = None
         if len(cached_folders) > 2:
@@ -202,9 +204,9 @@ def merge(self, pretrained_model_name_or_path_list: List[Union[str, os.PathLike]
                     theta_0 = theta_0()
 
                     update_theta_0 = getattr(module, "load_state_dict")
-                    theta_1 = torch.load(checkpoint_path_1)
+                    theta_1 = torch.load(checkpoint_path_1, map_location="cpu")
 
-                    theta_2 = torch.load(checkpoint_path_2) if checkpoint_path_2 else None
+                    theta_2 = torch.load(checkpoint_path_2, map_location="cpu") if checkpoint_path_2 else None
 
                     if not theta_0.keys() == theta_1.keys():
                         print("SKIPPING ATTR ", attr, " DUE TO MISMATCH")


Test patch: