Merge pull request #64 from ZeroCool940711/dev

Added several QOL improvements to the infer_vae.py script.
Sygil-Dev · Aug 10, 2023 · 562a4bf · 562a4bf
2 parents aa8e5fb + ab13cd8
commit 562a4bf
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 57 deletions.
diff --git a/infer_vae.py b/infer_vae.py
@@ -1,9 +1,5 @@
-import argparse
-import glob
-import hashlib
-import os
-import random
-import re
+import argparse, glob, hashlib
+import os, random, re, shutil
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional
@@ -184,6 +180,11 @@
     default=0,
     help="GPU to use in case we want to use a specific GPU for inference.",
 )
+parser.add_argument(
+    "--cpu",
+    action="store_true",
+    help="Use the CPU instead of the GPU, this will be really slow but can be useful for testing or if you dont have a good GPU.",
+)
 parser.add_argument(
     "--max_retries",
     type=int,
@@ -200,6 +201,11 @@
     action="store_true",
     help="Use PaintMind VAE..",
 )
+parser.add_argument(
+    "--save_originals",
+    action="store_true",
+    help="Save the original input.png and output.png images to a subfolder instead of deleting them after the grid is made.",
+)
 
 
 @dataclass
@@ -356,7 +362,7 @@ def main():
             channels=args.channels,
             layers=args.layers,
             discr_layers=args.discr_layers,
-        ).to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+        ).to('cpu' if args.cpu else accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
 
         if args.latest_checkpoint:
             accelerator.print("Finding latest checkpoint...")
@@ -421,7 +427,7 @@ def main():
         args.seq_len = vae.get_encoded_fmap_size(args.image_size) ** 2
 
     # move vae to device
-    vae = vae.to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+    vae = vae.to('cpu' if args.cpu else 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())
@@ -452,7 +458,7 @@ def main():
         )
 
         _, ids, _ = vae.encode(
-            dataset[image_id][None].to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+            dataset[image_id][None].to('cpu' if args.cpu else accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
         )
         recon = vae.decode_from_ids(ids)
         save_image(recon, f"{args.results_dir}/outputs/output.{str(args.input_image).split('.')[-1]}")
@@ -465,7 +471,7 @@ def main():
         save_image(dataset[image_id], f"{args.results_dir}/outputs/input.png")
 
         _, ids, _ = vae.encode(
-            dataset[image_id][None].to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+            dataset[image_id][None].to('cpu' if args.cpu else accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
         )
         recon = vae.decode_from_ids(ids)
         save_image(recon, f"{args.results_dir}/outputs/output.png")
@@ -484,7 +490,7 @@ def main():
                     if not args.use_paintmind:
                         # encode
                         _, ids, _ = vae.encode(
-                            dataset[i][None].to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+                            dataset[i][None].to('cpu' if args.cpu else accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
                         )
                         # decode
                         recon = vae.decode_from_ids(ids)
@@ -493,7 +499,7 @@ def main():
                     else:
                         # encode
                         encoded, _, _ = vae.encode(
-                            dataset[i][None].to(accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
+                            dataset[i][None].to('cpu' if args.cpu else accelerator.device if args.gpu == 0 else f"cuda:{args.gpu}")
                         )
 
                         # decode
@@ -520,9 +526,15 @@ def main():
                     filename = f"{hash}_{now}-{os.path.basename(args.vae_path)}.png"
                     grid_image.save(f"{output_dir}/{filename}", format="PNG")
 
-                    # Remove input and output images after the grid was made.
-                    os.remove(f"{output_dir}/input.png")
-                    os.remove(f"{output_dir}/output.png")
+                    if not args.save_originals:
+                        # Remove input and output images after the grid was made.
+                        os.remove(f"{output_dir}/input.png")
+                        os.remove(f"{output_dir}/output.png")
+                    else:
+                        os.makedirs(os.path.join(output_dir, 'originals'), exist_ok=True)
+                        shutil.move(f"{output_dir}/input.png", f"{os.path.join(output_dir, 'originals')}/input_{now}.png")
+                        shutil.move(f"{output_dir}/output.png", f"{os.path.join(output_dir, 'originals')}/output_{now}.png")
+
 
                     del _
                     del ids

diff --git a/train_muse_maskgit.py b/train_muse_maskgit.py
@@ -361,11 +361,11 @@
 parser.add_argument(
     "--optimizer",
     type=str,
-    default="Adafactor",
+    default="Adam",
     help="Optimizer to use. Choose between: ['Adam', 'AdamW','Lion', 'Adafactor', "
     "'AdaBound', 'AdaMod', 'AccSGD', 'AdamP', 'AggMo', 'DiffGrad', 'Lamb', "
     "'NovoGrad', 'PID', 'QHAdam', 'QHM', 'RAdam', 'SGDP', 'SGDW', 'Shampoo', "
-    "'SWATS', 'Yogi']. Default: Lion",
+    "'SWATS', 'Yogi']. Default: Adam",
 )
 parser.add_argument(
     "--weight_decay",
@@ -427,22 +427,38 @@
 @dataclass
 class Arguments:
     total_params: Optional[int] = None
-    only_save_last_checkpoint: bool = False
-    validation_image_scale: float = 1.0
-    no_center_crop: bool = False
-    no_flip: bool = False
-    dataset_save_path: Optional[str] = None
-    clear_previous_experiments: bool = False
+    image_size: int = 256
     num_tokens: int = 256
+    num_train_steps: int = -1
+    num_epochs: int = 5
+    dim: int = 128
+    channels: int = 3
+    batch_size: int = 1
+    lr: float = 1e-4
+    gradient_accumulation_steps: int = 1
+    save_results_every: int = 100
+    save_model_every: int = 500
+    vq_codebook_size: int = 256
+    vq_codebook_dim: int = 256
+    lr_scheduler: str = "constant"
+    lr_warmup_steps: int = 0
     seq_len: int = 1024
     depth: int = 2
     dim_head: int = 64
     heads: int = 8
     ff_mult: int = 4
     t5_name: str = "t5-small"
+    mixed_precision: str = "no"
     cond_image_size: Optional[int] = None
     validation_prompt: str = "A photo of a dog"
     timesteps: int = 18
+    optimizer: str = "Lion"
+    only_save_last_checkpoint: bool = False
+    validation_image_scale: float = 1.0
+    no_center_crop: bool = False
+    no_flip: bool = False
+    dataset_save_path: Optional[str] = None
+    clear_previous_experiments: bool = False
     max_grad_norm: Optional[float] = None
     seed: int = 42
     valid_frac: float = 0.05
@@ -454,7 +470,6 @@ class Arguments:
     image_column: str = "image"
     caption_column: str = "caption"
     log_with: str = "wandb"
-    mixed_precision: str = "no"
     use_8bit_adam: bool = False
     results_dir: str = "results"
     logging_dir: Optional[str] = None
@@ -463,27 +478,13 @@ class Arguments:
     hf_split_name: Optional[str] = None
     streaming: bool = False
     train_data_dir: Optional[str] = None
-    num_train_steps: int = -1
-    num_epochs: int = 5
-    dim: int = 128
-    batch_size: int = 512
-    lr: float = 1e-4
-    gradient_accumulation_steps: int = 1
-    save_results_every: int = 100
-    save_model_every: int = 500
     checkpoint_limit: Union[int, str] = None
-    vq_codebook_size: int = 256
-    vq_codebook_dim: int = 256
     cond_drop_prob: float = 0.5
-    image_size: int = 256
-    lr_scheduler: str = "constant"
     scheduler_power: float = 1.0
-    lr_warmup_steps: int = 0
     num_cycles: int = 1
     resume_path: Optional[str] = None
     taming_model_path: Optional[str] = None
     taming_config_path: Optional[str] = None
-    optimizer: str = "Lion"
     weight_decay: float = 0.0
     cache_path: Optional[str] = None
     no_cache: bool = False

diff --git a/train_muse_vae.py b/train_muse_vae.py
@@ -285,8 +285,11 @@
 parser.add_argument(
     "--optimizer",
     type=str,
-    default="Lion",
-    help="Optimizer to use. Choose between: ['Adam', 'AdamW','Lion']. Default: Lion",
+    default="Adam",
+    help="Optimizer to use. Choose between: ['Adam', 'AdamW','Lion', 'Adafactor', "
+    "'AdaBound', 'AdaMod', 'AccSGD', 'AdamP', 'AggMo', 'DiffGrad', 'Lamb', "
+    "'NovoGrad', 'PID', 'QHAdam', 'QHM', 'RAdam', 'SGDP', 'SGDW', 'Shampoo', "
+    "'SWATS', 'Yogi']. Default: Adam",
 )
 parser.add_argument(
     "--cache_path",
@@ -320,6 +323,24 @@
 @dataclass
 class Arguments:
     total_params: Optional[int] = None
+    image_size: int = 256
+    num_train_steps: int = -1
+    num_epochs: int = 5
+    batch_size: int = 512
+    lr: float = 1e-5
+    lr_warmup_steps: int = 0
+    lr_scheduler: str = "constant"
+    gradient_accumulation_steps: int = 1
+    save_results_every: int = 100
+    save_model_every: int = 500
+    dim: int = 128
+    vq_codebook_size: int = 256
+    vq_codebook_dim: int = 256
+    num_tokens: int = 256
+    seq_len: int = 1024
+    channels: int = 3
+    scheduler_power: float = 1.0
+    num_cycles: int = 1
     only_save_last_checkpoint: bool = False
     validation_image_scale: float = 1.0
     no_center_crop: bool = False
@@ -329,8 +350,6 @@ class Arguments:
     clear_previous_experiments: bool = False
     max_grad_norm: Optional[float] = None
     discr_max_grad_norm: Optional[float] = None
-    num_tokens: int = 256
-    seq_len: int = 1024
     seed: int = 42
     valid_frac: float = 0.05
     use_ema: bool = False
@@ -349,23 +368,8 @@ class Arguments:
     dataset_name: Optional[str] = None
     streaming: bool = False
     train_data_dir: Optional[str] = None
-    num_train_steps: int = -1
-    num_epochs: int = 5
-    dim: int = 128
-    batch_size: int = 512
-    lr: float = 1e-5
-    gradient_accumulation_steps: int = 1
-    save_results_every: int = 100
-    save_model_every: int = 500
     checkpoint_limit: Union[int, str] = None
-    vq_codebook_size: int = 256
-    vq_codebook_dim: int = 256
     cond_drop_prob: float = 0.5
-    image_size: int = 256
-    lr_scheduler: str = "constant"
-    scheduler_power: float = 1.0
-    lr_warmup_steps: int = 0
-    num_cycles: int = 1
     taming_model_path: Optional[str] = None
     taming_config_path: Optional[str] = None
     optimizer: str = "Lion"
@@ -539,6 +543,8 @@ def main():
         )
         args.num_tokens = vae.codebook_size
         args.seq_len = vae.get_encoded_fmap_size(args.image_size) ** 2
+
+        current_step = 0
     else:
         print("Initialising empty VAE")
         vae = VQGanVAE(