diff --git a/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py b/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py index 81e151f723..57675812f7 100644 --- a/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py +++ b/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py @@ -496,10 +496,11 @@ def __call__( hb_profiler.start() # 8. Denoising loop + throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3) for j in self.progress_bar(range(num_batches)): # The throughput is calculated from the 3rd iteration # because compilation occurs in the first two iterations - if j == kwargs.get("throughput_warmup_steps", 3): + if j == throughput_warmup_steps: t1 = time.time() latents_batch = latents_batches[0] @@ -615,7 +616,9 @@ def __call__( speed_measures = speed_metrics( split=speed_metrics_prefix, start_time=t0, - num_samples=num_batches * batch_size if t1 == t0 else (num_batches - 2) * batch_size, + num_samples=num_batches * batch_size + if t1 == t0 + else (num_batches - throughput_warmup_steps) * batch_size, num_steps=num_batches, start_time_after_warmup=t1, ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 8f412f0c20..bf5d90df51 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -477,10 +477,11 @@ def __call__( self._num_timesteps = len(timesteps) # 8. Denoising loop + throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3) for j in self.progress_bar(range(num_batches)): # The throughput is calculated from the 3rd iteration # because compilation occurs in the first two iterations - if j == kwargs.get("throughput_warmup_steps", 3): + if j == throughput_warmup_steps: t1 = time.time() latents_batch = latents_batches[0] @@ -560,7 +561,9 @@ def __call__( speed_measures = speed_metrics( split=speed_metrics_prefix, start_time=t0, - num_samples=num_batches * batch_size if t1 == t0 else (num_batches - 2) * batch_size, + num_samples=num_batches * batch_size + if t1 == t0 + else (num_batches - throughput_warmup_steps) * batch_size, num_steps=num_batches, start_time_after_warmup=t1, ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py index 0fbcd4a274..ffde6d08a7 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py @@ -339,10 +339,11 @@ def __call__( t1 = t0 # 8. Denoising loop + throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3) for j in self.progress_bar(range(num_batches)): # The throughput is calculated from the 3rd iteration # because compilation occurs in the first two iterations - if j == kwargs.get("throughput_warmup_steps", 3): + if j == throughput_warmup_steps: t1 = time.time() latents_batch = latents_batches[0] @@ -401,7 +402,9 @@ def __call__( speed_measures = speed_metrics( split=speed_metrics_prefix, start_time=t0, - num_samples=num_batches * batch_size if t1 == t0 else (num_batches - 2) * batch_size, + num_samples=num_batches * batch_size + if t1 == t0 + else (num_batches - throughput_warmup_steps) * batch_size, num_steps=num_batches, start_time_after_warmup=t1, ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index a574746e38..c4f7a5d245 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -436,10 +436,11 @@ def __call__( t1 = t0 # 10. Denoising loop + throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3) for j in self.progress_bar(range(num_batches)): # The throughput is calculated from the 3rd iteration # because compilation occurs in the first two iterations - if j == kwargs.get("throughput_warmup_steps", 3): + if j == throughput_warmup_steps: t1 = time.time() latents_batch = latents_batches[0] @@ -522,7 +523,9 @@ def __call__( speed_measures = speed_metrics( split=speed_metrics_prefix, start_time=t0, - num_samples=num_batches * batch_size if t1 == t0 else (num_batches - 2) * batch_size, + num_samples=num_batches * batch_size + if t1 == t0 + else (num_batches - throughput_warmup_steps) * batch_size, num_steps=num_batches, start_time_after_warmup=t1, ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py index 72c132847f..406f46753b 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py @@ -678,10 +678,11 @@ def __call__( self._num_timesteps = len(timesteps) # 8.3 Denoising loop + throughput_warmup_steps = kwargs.get("throughput_warmup_steps", 3) for j in self.progress_bar(range(num_batches)): # The throughput is calculated from the 3rd iteration # because compilation occurs in the first two iterations - if j == kwargs.get("throughput_warmup_steps", 3): + if j == throughput_warmup_steps: t1 = time.time() latents_batch = latents_batches[0] @@ -778,7 +779,9 @@ def __call__( speed_measures = speed_metrics( split=speed_metrics_prefix, start_time=t0, - num_samples=num_batches * batch_size if t1 == t0 else (num_batches - 2) * batch_size, + num_samples=num_batches * batch_size + if t1 == t0 + else (num_batches - throughput_warmup_steps) * batch_size, num_steps=num_batches, start_time_after_warmup=t1, )