update CLI default arguments with working ones

Author: Ipsedo

music_diffusion/__main__.py

@@ -60,25 +60,22 @@ def main() -> None:
     # Model hyper parameters
     model_parser = sub_command.add_parser("model")
 
-    model_parser.add_argument("--steps", type=int, default=1024)
-    model_parser.add_argument("--beta-1", type=float, default=1e-4)
-    model_parser.add_argument("--beta-t", type=float, default=2e-2)
+    model_parser.add_argument("--steps", type=int, default=4096)
+    model_parser.add_argument("--beta-1", type=float, default=2.5e-5)
+    model_parser.add_argument("--beta-t", type=float, default=5e-3)
     model_parser.add_argument("--channels", type=int, default=2)
     model_parser.add_argument(
         "--unet-channels",
         type=_channels,
         default=[
-            (8, 16),
-            (16, 24),
-            (24, 32),
-            (32, 40),
-            (40, 48),
-            (48, 56),
-            (56, 64),
+            (16, 32),
+            (32, 48),
+            (48, 64),
+            (64, 80),
         ],
     )
-    model_parser.add_argument("--time-size", type=int, default=32)
-    model_parser.add_argument("--norm-groups", type=int, default=8)
+    model_parser.add_argument("--time-size", type=int, default=48)
+    model_parser.add_argument("--norm-groups", type=int, default=16)
     model_parser.add_argument("--cuda", action="store_true")
 
     # Sub command run {train, generate}
@@ -92,7 +89,7 @@ def main() -> None:
     train_parser.add_argument("run_name", type=str)
 
     train_parser.add_argument("-i", "--input-dataset", type=str, required=True)
-    train_parser.add_argument("--batch-size", type=int, default=8)
+    train_parser.add_argument("--batch-size", type=int, default=6)
     train_parser.add_argument("--step-batch-size", type=int, default=1)
     train_parser.add_argument("--epochs", type=int, default=1000)
     train_parser.add_argument("--learning-rate", type=float, default=2e-4)