Add options and local inference

Added options to:
- Disable Inference (it consumes about 2gb of VRAM even when not active)
- Disable wandb

and:
- if no hftoken is provided it just fills it with nothing so it doesn't argues
- if wandb is not enabled, save the inference outputs to a local folder along with information about it

trainer/diffusers_trainer.py

@@ -84,6 +84,8 @@ parser.add_argument('--clip_penultimate', type=bool_t, default='False', help='Us
 parser.add_argument('--output_bucket_info', type=bool_t, default='False', help='Outputs bucket information and exits')
 parser.add_argument('--resize', type=bool_t, default='False', help="Resizes dataset's images to the appropriate bucket dimensions.")
 parser.add_argument('--use_xformers', type=bool_t, default='False', help='Use memory efficient attention')
+parser.add_argument('--wandb', dest='enablewandb', type=str, default='True', help='Enable WeightsAndBiases Reporting')
+parser.add_argument('--inference', dest='enableinference', type=str, default='True', help='Enable Inference during training (Consumes 2GB of VRAM)')
 args = parser.parse_args()
 
 def setup():
@@ -520,7 +522,11 @@ def main():
 
     if rank == 0:
         os.makedirs(args.output_path, exist_ok=True)
+        
+        if args.enablewandb:
             run = wandb.init(project=args.project_id, name=args.run_name, config=vars(args), dir=args.output_path+'/wandb')
+        else:
+            run = wandb.init(project=args.project_id, name=args.run_name, config=vars(args), dir=args.output_path+'/wandb', mode="disabled")
 
         # Inform the user of host, and various versions -- useful for debugging issues.
         print("RUN_NAME:", args.run_name)
@@ -534,8 +540,12 @@ def main():
         print("RESOLUTION:", args.resolution)
 
     if args.hf_token is None:
+        try:
             args.hf_token = os.environ['HF_API_TOKEN']
             print('It is recommended to set the HF_API_TOKEN environment variable instead of passing it as a command line argument since WandB will automatically log it.')
+        except Exception:
+            print("No HF Token detected in arguments or enviroment variable, setting it to none (as in string)")
+            args.hf_token = "none"
 
     device = torch.device('cuda')
 
@@ -744,6 +754,7 @@ def main():
                 if global_step % args.save_steps == 0:
                     save_checkpoint(global_step)
 
+                if args.enableinference:
                     if global_step % args.image_log_steps == 0:
                         if rank == 0:
                             # get prompt from random batch
@@ -770,17 +781,35 @@ def main():
                                 feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32"),
                             ).to(device)
                             # inference
+                            if args.enablewandb:
                                 images = []
+                            else:
+                                saveInferencePath = args.output_path + "/inference"
+                                os.makedirs(saveInferencePath, exist_ok=True)
                             with torch.no_grad():
                                 with torch.autocast('cuda', enabled=args.fp16):
                                     for _ in range(args.image_log_amount):
+                                        if args.enablewandb:
                                             images.append(
                                                 wandb.Image(pipeline(
                                                     prompt, num_inference_steps=args.image_log_inference_steps
                                                 ).images[0],
                                                 caption=prompt)
                                             )
+                                        else:
+                                            from datetime import datetime
+                                            images = pipeline(prompt, num_inference_steps=args.image_log_inference_steps).images[0]
+                                            filenameImg = str(time.time_ns()) + ".png"
+                                            filenameTxt = str(time.time_ns()) + ".txt"
+                                            images.save(saveInferencePath + "/" + filenameImg)
+                                            with open(saveInferencePath + "/" + filenameTxt, 'a') as f:
+                                                f.write('Used prompt: ' + prompt + '\n')
+                                                f.write('Generated Image Filename: ' + filenameImg + '\n')
+                                                f.write('Generated at: ' + str(global_step) + ' steps' + '\n')
+                                                f.write('Generated at: ' + str(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))+ '\n')
+
                             # log images under single caption
+                            if args.enablewandb:
                                 run.log({'images': images}, step=global_step)
 
                             # cleanup so we don't run out of memory