add push_to_hub

.gitignore

@@ -1,5 +1,5 @@
 .vscode
 __pycache__
 .ipynb_checkpoints
-data
+data*
 ddpm-ema-audio-*

src/audio_to_images.py

@@ -58,6 +58,8 @@ def main(args):
         )
         dsd = DatasetDict({"train": ds})
         dsd.save_to_disk(os.path.join(args.output_dir))
+        if args.push_to_hub:
+            dsd.push_to_hub(args.push_to_hub)
 
 
 if __name__ == "__main__":
@@ -68,5 +70,6 @@ if __name__ == "__main__":
     parser.add_argument("--output_dir", type=str, default="data")
     parser.add_argument("--resolution", type=int, default=256)
     parser.add_argument("--hop_length", type=int, default=512)
+    parser.add_argument("--push_to_hub", type=str, default=None)
     args = parser.parse_args()
     main(args)

src/train_unconditional.py

@@ -80,7 +80,18 @@ def main(args):
     )
 
     if args.dataset_name is not None:
-        dataset = load_from_disk(args.dataset_name, args.dataset_config_name)["train"]
+        if os.path.exists(args.dataset_name):
+            dataset = load_from_disk(args.dataset_name, args.dataset_config_name)[
+                "train"
+            ]
+        else:
+            dataset = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                cache_dir=args.cache_dir,
+                use_auth_token=True if args.use_auth_token else None,
+                split="train",
+            )
     else:
         dataset = load_dataset(
             "imagefolder",
@@ -203,11 +214,14 @@ def main(args):
                 accelerator.trackers[0].writer.add_images(
                     "test_samples", images_processed, epoch
                 )
-                for image in images_processed:
+                for _, image in enumerate(images_processed):
                     image = Image.fromarray(np.mean(image, axis=0).astype("uint8"))
                     audio = mel.image_to_audio(image)
                     accelerator.trackers[0].writer.add_audio(
-                        "test_samples", audio, epoch, sample_rate=mel.get_sample_rate()
+                        f"test_audio_{_}",
+                        audio,
+                        epoch,
+                        sample_rate=mel.get_sample_rate(),
                     )
 
             if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: