logging

hydra_node/models.py

@@ -168,9 +168,11 @@ class StableDiffusionModel(nn.Module):
         self.config = config
         self.premodules = None
         if Path(self.config.model_path).is_dir():
+            config.logger.info(f"Loading model from folder {self.config.model_path}")
             model, model_config = self.from_folder(config.model_path)
         
         elif Path(self.config.model_path).is_file():
+            config.logger.info(f"Loading model from file {self.config.model_path}")
             model, model_config = self.from_file(config.model_path)
         
         else:
@@ -195,10 +197,12 @@ class StableDiffusionModel(nn.Module):
             model.first_stage_model = model.first_stage_model.float()
             del ckpt
             del loss
+            config.logger.info(f"Using VAE from {self.config.vae_path}")
 
         if self.config.penultimate == "1":
             model.cond_stage_model.return_layer = -2
             model.cond_stage_model.do_final_ln = True
+            config.logger.info(f"CLIP: Using penultimate layer")
             
         model.cond_stage_model.inference_mode = True
         self.k_model = K.external.CompVisDenoiser(model)
@@ -210,6 +214,9 @@ class StableDiffusionModel(nn.Module):
         self.ema_manager = self.model.ema_scope
         if self.config.enable_ema == "0":
             self.ema_manager = contextlib.nullcontext
+            config.logger.info("Disabling EMA")
+        else:
+            config.logger.info(f"Using EMA")
         self.sampler_map = {
             'plms': self.plms.sample,
             'ddim': self.ddim.sample,
@@ -241,7 +248,6 @@ class StableDiffusionModel(nn.Module):
         model = self.load_model_from_config(model_config, file)
         return model, model_config
 
-
     def load_model_from_config(self, config, ckpt, verbose=False):
         print(f"Loading model from {ckpt}")
         pl_sd = torch.load(ckpt, map_location="cpu")