Merge pull request Sygil-Dev#54 from ZeroCool940711/dev

Added a couple of lines to print the total number of parameters the model has after it has been loaded for the vae and maskgit training scripts as well as the infer_vae.py

infer_vae.py

@@ -413,6 +413,12 @@ def main():
     # move vae to device
     vae = vae.to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
 
+
+    # Use the parameters() method to get an iterator over all the learnable parameters of the model
+    total_params = sum(p.numel() for p in vae.parameters())
+
+    print(f"Total number of parameters: {format(total_params, ',d')}")
+
     # then you plug the vae and transformer into your MaskGit as so
 
     dataset = ImageDataset(

train_muse_maskgit.py

@@ -790,6 +790,12 @@ def main():
             accelerator.print("Initialized new empty MaskGit model.")
             current_step = 0
 
+
+    # Use the parameters() method to get an iterator over all the learnable parameters of the model
+    total_params = sum(p.numel() for p in maskgit.parameters())
+
+    print(f"Total number of parameters: {format(total_params, ',d')}")
+
     # Create the dataset objects
     with accelerator.main_process_first():
         if args.no_cache and args.train_data_dir:

train_muse_vae.py

@@ -541,6 +541,11 @@ def main():
 
         current_step = 0
 
+    # Use the parameters() method to get an iterator over all the learnable parameters of the model
+    total_params = sum(p.numel() for p in vae.parameters())
+
+    print(f"Total number of parameters: {format(total_params, ',d')}")
+
     dataset = ImageDataset(
         dataset,
         args.image_size,