Autodetect checkpoint type

hydra_node/models.py

@@ -279,7 +279,7 @@ class StableDiffusionModel(nn.Module):
         pl_sd = torch.load(ckpt, map_location="cpu")
         if "global_step" in pl_sd:
             print(f"Global Step: {pl_sd['global_step']}")
-        if self.config.basedformer == "1":
+        if self.config.basedformer == "1" or 'state_dict' in pl_sd:
             sd = pl_sd
         else:
             sd = pl_sd["state_dict"]