1 files changed, 3 insertions, 13 deletions

diff --git a/dreambooth_plus.py b/dreambooth_plus.py
index 9e482b3..7996bc2 100644
--- a/dreambooth_plus.py
+++ b/dreambooth_plus.py
@@ -112,7 +112,7 @@ def parse_args():
     parser.add_argument(
         "--max_train_steps",
         type=int,
-        default=3000,
+        default=1600,
         help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
     )
     parser.add_argument(
@@ -129,13 +129,13 @@ def parse_args():
     parser.add_argument(
         "--learning_rate_unet",
         type=float,
-        default=1e-5,
+        default=5e-6,
         help="Initial learning rate (after the potential warmup period) to use.",
     )
     parser.add_argument(
         "--learning_rate_text",
         type=float,
-        default=1e-4,
+        default=5e-4,
         help="Initial learning rate (after the potential warmup period) to use.",
     )
     parser.add_argument(
@@ -222,12 +222,6 @@ def parse_args():
         ),
     )
     parser.add_argument(
-        "--local_rank",
-        type=int,
-        default=-1,
-        help="For distributed training: local_rank"
-    )
-    parser.add_argument(
         "--sample_frequency",
         type=int,
         default=100,
@@ -293,10 +287,6 @@ def parse_args():
             args = parser.parse_args(
                 namespace=argparse.Namespace(**json.load(f)["args"]))
 
-    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
-    if env_local_rank != -1 and env_local_rank != args.local_rank:
-        args.local_rank = env_local_rank
-
     if args.train_data_file is None:
         raise ValueError("You must specify --train_data_file")