From 42f5bd4e12e7418d77b0e7ea6bd64ebfb9523794 Mon Sep 17 00:00:00 2001
From: spezialspezial <75758219+spezialspezial@users.noreply.github.com>
Date: Sat, 3 Dec 2022 22:01:43 +0100
Subject: [PATCH] Account for flat models

Merged models from auto11 merge board are flat for some reason. Current behavior of invoke is not changed by this modification.
---
 ldm/invoke/model_cache.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ldm/invoke/model_cache.py b/ldm/invoke/model_cache.py
index 7d1654718a..f937b7eae5 100644
--- a/ldm/invoke/model_cache.py
+++ b/ldm/invoke/model_cache.py
@@ -227,7 +227,9 @@ class ModelCache(object):
         model_hash  = self._cached_sha256(weights,weight_bytes)
         sd = torch.load(io.BytesIO(weight_bytes), map_location='cpu')
         del weight_bytes
-        sd = sd['state_dict']
+        # merged models from auto11 merge board are flat for some reason
+        if 'state_dict' in sd:
+            sd = sd['state_dict']
         model = instantiate_from_config(omega_config.model)
         model.load_state_dict(sd, strict=False)