pruned ckpt support

hydra_node/models.py

@@ -154,7 +154,11 @@ class StableDiffusionModel(nn.Module):
     def from_folder(self, folder):
         folder = Path(folder)
         model_config = OmegaConf.load(folder / "config.yaml")
-        model = self.load_model_from_config(model_config, folder / "model.ckpt")
+        if (folder / "pruned.ckpt").is_file():
+            model_path = folder / "pruned.ckpt"
+        else:
+            model_path = folder / "model.ckpt"
+        model = self.load_model_from_config(model_config, folder / "pruned.ckpt")
         return model, model_config
 
     def load_model_from_config(self, config, ckpt, verbose=False):