Fix bf16 support

Fix scheduler.step call to latest
huggingface · Jan 12, 2024 · e0dab26 · e0dab26
1 parent 6b4c5f3
commit e0dab26
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/examples/controlnet/README.md b/examples/controlnet/README.md
@@ -33,6 +33,7 @@ python text_to_image_generation_canny.py \
     --image_save_dir /tmp/controlnet_images \
     --use_habana \
     --use_hpu_graphs \
-    --gaudi_config Habana/stable-diffusion
+    --gaudi_config Habana/stable-diffusion \
+    --bf16
 ```
 
diff --git a/examples/controlnet/text_to_image_generation_canny.py b/examples/controlnet/text_to_image_generation_canny.py
@@ -191,7 +191,8 @@ def main():
 
     # Initialize the scheduler and the generation pipeline
     scheduler = GaudiDDIMScheduler.from_pretrained(args.model_name_or_path, subfolder="scheduler")
-    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
+    model_dtype = torch.bfloat16 if args.bf16 else None
+    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path, torch_dtype=model_dtype)
     kwargs = {
         "scheduler": scheduler,
         "use_habana": args.use_habana,

diff --git a/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py b/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py
@@ -1209,7 +1209,7 @@ def __call__(
 
                     # compute the previous noisy sample x_t -> x_t-1
                     latents_batch = self.scheduler.step(
-                        noise_pred, latents_batch, **extra_step_kwargs, return_dict=False
+                        noise_pred, timestep, latents_batch, **extra_step_kwargs, return_dict=False
                     )[0]
 
                     if not self.use_hpu_graphs: