PR from nachoal: https://github.com/JoePenna/Dreambooth-Stable-Diffus…

…ion/pull/166/files

.gitignore

@@ -17,3 +17,6 @@ venv
 *.ckpt
 /joepenna-dreambooth-configs/
 /src/
+
+dreambooth_joepenna
+src/taming-transformers

ldm/util.py

@@ -77,7 +77,12 @@ def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
 
     pl_sd = torch.load(ckpt, map_location="cpu")
-    sd = pl_sd["state_dict"]
+    if "state_dict" in pl_sd:
+        sd = pl_sd["state_dict"]
+    else:
+        print(f"Warning: 'state_dict' key not found in the checkpoint file {ckpt}. Attempting to load the entire checkpoint as the model state.")
+        sd = pl_sd
+
     config["model"]["params"]["ckpt_path"] = ckpt
 
     print("")