add sdp_on_bf16 to controlnet (#1631)

* add sdp_on_bf16 to controlnet * Update pipeline_controlnet.py pass sdp_on_bf16 to controlnet_pipeline * Update text_to_image_generation.py * Update text_to_image_generation.py
huggingface · Dec 20, 2024 · 1e747dd · 1e747dd
1 parent 948476a
commit 1e747dd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/examples/stable-diffusion/training/train_controlnet.py b/examples/stable-diffusion/training/train_controlnet.py
@@ -120,6 +120,7 @@ def log_validation(
         use_habana=True,
         use_hpu_graphs=args.use_hpu_graphs,
         gaudi_config=gaudi_config,
+        sdp_on_bf16=args.sdp_on_bf16,
     )
     pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
     pipeline = pipeline.to(accelerator.device)
@@ -438,6 +439,12 @@ def parse_args(input_args=None):
             ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
         ),
     )
+    parser.add_argument(
+        "--sdp_on_bf16",
+        action="store_true",
+        default=False,
+        help="Allow pyTorch to use reduced precision in the SDPA math backend",
+    )
     parser.add_argument(
         "--bf16",
         action="store_true",

diff --git a/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py b/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py
@@ -100,13 +100,15 @@ def __init__(
         use_hpu_graphs: bool = False,
         gaudi_config: Union[str, GaudiConfig] = None,
         bf16_full_eval: bool = False,
+        sdp_on_bf16: bool = True,
     ):
         GaudiDiffusionPipeline.__init__(
             self,
             use_habana,
             use_hpu_graphs,
             gaudi_config,
             bf16_full_eval,
+            sdp_on_bf16,
         )
 
         StableDiffusionControlNetPipeline.__init__(