[SD][CL] Disable print at every iteration. (huggingface#664)

Printing might incur extra time to runtime. Hence, we add a flag to hide it. To disable printing please set this flag `--hide_steps`.

Co-authored-by: Stanley <stanley@MacStudio.lan>

Author: raikonenfnu

shark/examples/shark_inference/stable_diffusion/main.py

@@ -147,9 +147,10 @@ def end_profiling(device):
     latents = latents * scheduler.init_noise_sigma
     avg_ms = 0
 
-    for i, t in tqdm(enumerate(scheduler.timesteps)):
+    for i, t in tqdm(enumerate(scheduler.timesteps), disable=args.hide_steps):
         step_start = time.time()
-        print(f"i = {i} t = {t}", end="")
+        if args.hide_steps == False:
+            print(f"i = {i} t = {t}", end="")
         timestep = torch.tensor([t]).to(dtype).detach().numpy()
         latent_model_input = scheduler.scale_model_input(latents, t)
         if cpu_scheduling:
@@ -177,7 +178,8 @@ def end_profiling(device):
         step_time = time.time() - step_start
         avg_ms += step_time
         step_ms = int((step_time) * 1000)
-        print(f" ({step_ms}ms)")
+        if args.hide_steps == False:
+            print(f" ({step_ms}ms)")
 
     avg_ms = 1000 * avg_ms / args.steps
     print(f"Average step time: {avg_ms}ms/it")

shark/examples/shark_inference/stable_diffusion/stable_args.py

@@ -170,4 +170,12 @@
     action=argparse.BooleanOptionalAction,
     help="flag for inserting debug frames between iterations for use with rgp.",
 )
+
+p.add_argument(
+    "--hide_steps",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for hiding the details of iteration/sec for each step.",
+)
+
 args = p.parse_args()