Small QOL when using the validation_folder_at_end_of_epoch argument, …

…we now append the epoch number to the filename so we can identify at what epoch the image was saved. (Sygil-Dev#82)

- Added experimental t5 offloading to CPU.

muse_maskgit_pytorch/trainers/maskgit_trainer.py

@@ -58,6 +58,7 @@ def __init__(
         clear_previous_experiments=False,
         validation_image_scale: float = 1.0,
         only_save_last_checkpoint=False,
+        t5_offloading=False,
         args=None,
     ):
         super().__init__(
@@ -88,6 +89,8 @@ def __init__(
         # we are going to use them later to save them to a config file.
         self.args = args
 
+        self.t5_offloading = t5_offloading
+
         # maskgit
         maskgit.vae.requires_grad_(False)
         maskgit.transformer.t5.requires_grad_(False)
@@ -161,7 +164,7 @@ def train(self):
                 if not text_embeds:
                     with torch.no_grad():
                         text_embeds = t5_encode_text_from_encoded(
-                            input_ids, attn_mask, self.model.transformer.t5, self.accelerator.device
+                            input_ids, attn_mask, self.model.transformer.t5, self.accelerator.device if not self.t5_offloading else 'cpu'
                         )
 
                 with self.accelerator.accumulate(self.model), self.accelerator.autocast():

muse_maskgit_pytorch/utils.py

@@ -202,7 +202,7 @@ def vae_folder_validation(accelerator, vae, dataset, args=None, checkpoint_name=
                 now = datetime.now().strftime("%m-%d-%Y_%H-%M-%S")
                 hash = hashlib.sha1(input_image.tobytes()).hexdigest()
 
-                filename = f"{hash}_{now}{'-' + epoch if epoch else ''}-{os.path.basename(checkpoint_name)}.png"
+                filename = f"{str(hash)}_{str(now)}{'-'}{'E' + str(epoch) if epoch is not None else ''}-{str(os.path.basename(checkpoint_name))}.png"
                 grid_image.save(f"{output_dir}/{filename}", format="PNG")
 
                 if not save_originals:

train_muse_maskgit.py

@@ -454,6 +454,12 @@ def decompress_pickle(file):
     default="",
     help="The path to save or load embeds",
 )
+parser.add_argument(
+    "--t5_offloading",
+    action="store_true",
+    default=False,
+    help="Wheter to offload the t5 model to cpu instad of loading it on GPU. Should help with loading bigger models but will affect performance.",
+)
 
 
 @dataclass
@@ -529,6 +535,7 @@ class Arguments:
     attention_type: str = "flash"
     precompute: bool = False
     precompute_path: str = ""
+    t5_offloading = False
 
 
 def main():
@@ -1016,6 +1023,7 @@ def main():
         validation_image_scale=args.validation_image_scale,
         only_save_last_checkpoint=args.only_save_last_checkpoint,
         num_epochs=args.num_epochs,
+        t5_offloading=args.t5_offloading,
         args=args,
     )