When doing inference on the VAE the output images are now saved on th…

…e outputs folder on the results_dir.
ZeroCool940711 · Mar 9, 2023 · 8677417 · 8677417
1 parent 6dc15f0
commit 8677417
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/infer_vae.py b/infer_vae.py
@@ -192,13 +192,13 @@ def main():
 
     image_id = 0 if not args.random_image else random.randint(0, len(dataset))
 
-    os.makedirs(args.results_dir, exist_ok=True)
+    os.makedirs(f"{args.results_dir}/outputs", exist_ok=True)
 
-    save_image(dataset[image_id], f"{args.results_dir}/input.png")
+    save_image(dataset[image_id], f"{args.results_dir}/outputs/input.png")
 
     _, ids, _ = vae.encode(dataset[image_id][None].to(accelerator.device))
     recon = vae.decode_from_ids(ids)
-    save_image(recon, f"{args.results_dir}/output.png")
+    save_image(recon, f"{args.results_dir}/outputs/output.png")
 
 if __name__ == "__main__":
     main()