From 30093594ed16fe5e1c7dbc586c57c5c1ceea1d60 Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Wed, 24 Jan 2024 16:52:01 -0600
Subject: [PATCH] Fix Pruning

---
 modules/merging/merge.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/modules/merging/merge.py b/modules/merging/merge.py
index 555966a5d..d1e6129bc 100644
--- a/modules/merging/merge.py
+++ b/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():