Removes duplicated args.

Author: andsteing

examples/controlnet/train_controlnet_flax.py

@@ -221,11 +221,6 @@ def parse_args():
         default=None,
         help="Revision of controlnet model identifier from huggingface.co/models.",
     )
-    parser.add_argument(
-        "--controlnet_from_pt",
-        action="store_true",
-        help="Load the controlnet model from a PyTorch checkpoint.",
-    )
     parser.add_argument(
         "--profile_steps",
         type=int,
@@ -248,12 +243,6 @@ def parse_args():
         default=None,
         help="Enables compilation cache.",
     )
-    parser.add_argument(
-        "--controlnet_revision",
-        type=str,
-        default=None,
-        help="Revision of controlnet model identifier from huggingface.co/models.",
-    )
     parser.add_argument(
         "--controlnet_from_pt",
         action="store_true",
@@ -1058,16 +1047,18 @@ def l2(xs):
         )
         # train
         for batch in train_dataloader:
-            if args.profile_steps and global_step == 0:
-                jax.profiler.start_trace(args.output_dir)
-            if args.profile_steps and args.profile_steps == global_step:
-                jax.profiler.stop_trace()
 
             batch = shard(batch)
             with jax.profiler.StepTraceAnnotation("train", step_num=global_step):
                 state, train_metric, train_rngs = p_train_step(
                     state, unet_params, text_encoder_params, vae_params, batch, train_rngs
                 )
+            if args.profile_steps and global_step == 0:
+                train_metric['loss'].block_until_ready()
+                jax.profiler.start_trace(args.output_dir)
+            if args.profile_steps and args.profile_steps == global_step:
+                train_metric['loss'].block_until_ready()
+                jax.profiler.stop_trace()
             train_metrics.append(train_metric)
 
             train_step_progress_bar.update(1)