Fix Pruning

modules/merging/merge.py

@@ -1,7 +1,7 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Set
 import safetensors.torch
 import torch
 from tensordict import TensorDict
@@ -49,14 +49,14 @@ def fix_clip(model: Dict) -> Dict:
     return model
 
 
-def prune_sd_model(model: Dict) -> Dict:
+def prune_sd_model(model: Dict, keyset: Set) -> Dict:
     keys = list(model.keys())
     for k in keys:
         if (
             not k.startswith("model.diffusion_model.")
             # and not k.startswith("first_stage_model.")
             and not k.startswith("cond_stage_model.")
-        ):
+        ) or k not in keyset:
             del model[k]
     return model
 
@@ -78,10 +78,10 @@ def load_thetas(
     device: torch.device,
     precision: str,
 ) -> Dict:
+    thetas = {k: TensorDict.from_dict(read_state_dict(m, "cpu")) for k, m in models.items()}
     if prune:
-        thetas = {k: prune_sd_model(TensorDict.from_dict(read_state_dict(m, "cpu"))) for k, m in models.items()}
-    else:
-        thetas = {k: TensorDict.from_dict(read_state_dict(m, device)) for k, m in models.items()}
+        keyset = set.intersection(*[set(m.keys()) for m in thetas.values() if len(m.keys())])
+        thetas = {k: prune_sd_model(m, keyset) for k, m in thetas.items()}
 
     for model_key, model in thetas.items():
         for key, block in model.items():