Skip to content

Commit e1901a8

Browse files
raikonenfnuStanley
andauthored
[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>
1 parent 7d0cbd8 commit e1901a8

2 files changed

Lines changed: 13 additions & 3 deletions

File tree

shark/examples/shark_inference/stable_diffusion/main.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -147,9 +147,10 @@ def end_profiling(device):
147147
latents = latents * scheduler.init_noise_sigma
148148
avg_ms = 0
149149

150-
for i, t in tqdm(enumerate(scheduler.timesteps)):
150+
for i, t in tqdm(enumerate(scheduler.timesteps), disable=args.hide_steps):
151151
step_start = time.time()
152-
print(f"i = {i} t = {t}", end="")
152+
if args.hide_steps == False:
153+
print(f"i = {i} t = {t}", end="")
153154
timestep = torch.tensor([t]).to(dtype).detach().numpy()
154155
latent_model_input = scheduler.scale_model_input(latents, t)
155156
if cpu_scheduling:
@@ -177,7 +178,8 @@ def end_profiling(device):
177178
step_time = time.time() - step_start
178179
avg_ms += step_time
179180
step_ms = int((step_time) * 1000)
180-
print(f" ({step_ms}ms)")
181+
if args.hide_steps == False:
182+
print(f" ({step_ms}ms)")
181183

182184
avg_ms = 1000 * avg_ms / args.steps
183185
print(f"Average step time: {avg_ms}ms/it")

shark/examples/shark_inference/stable_diffusion/stable_args.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,4 +170,12 @@
170170
action=argparse.BooleanOptionalAction,
171171
help="flag for inserting debug frames between iterations for use with rgp.",
172172
)
173+
174+
p.add_argument(
175+
"--hide_steps",
176+
default=False,
177+
action=argparse.BooleanOptionalAction,
178+
help="flag for hiding the details of iteration/sec for each step.",
179+
)
180+
173181
args = p.parse_args()

0 commit comments

Comments
 (0)