Merge remote-tracking branch 'optimum-habana/main' into port_svd_cn

examples/stable-diffusion/image_to_video_generation.py

@@ -228,8 +228,12 @@ def main():
         from optimum.habana.diffusers import GaudiStableVideoDiffusionControlNetPipeline
         from optimum.habana.diffusers.models import ControlNetSDVModel, UNetSpatioTemporalConditionControlNetModel
 
-        controlnet = ControlNetSDVModel.from_pretrained(args.controlnet_model_name_or_path, subfolder="controlnet", **kwargs)
-        unet = UNetSpatioTemporalConditionControlNetModel.from_pretrained(args.model_name_or_path, subfolder="unet", **kwargs)
+        controlnet = ControlNetSDVModel.from_pretrained(
+            args.controlnet_model_name_or_path, subfolder="controlnet", **kwargs
+        )
+        unet = UNetSpatioTemporalConditionControlNetModel.from_pretrained(
+            args.model_name_or_path, subfolder="unet", **kwargs
+        )
         pipeline = GaudiStableVideoDiffusionControlNetPipeline.from_pretrained(
             args.model_name_or_path, controlnet=controlnet, unet=unet, **kwargs
         )

optimum/habana/diffusers/models/controlnet_sdv.py

@@ -551,8 +551,7 @@ def from_unet(
         return controlnet
 
     # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
-    def set_attn_processor(
-        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
         r"""
         Sets the attention processor to use to compute attention.
 

tests/test_diffusers.py

@@ -109,13 +109,13 @@
     GaudiStableDiffusionXLImg2ImgPipeline,
     GaudiStableDiffusionXLInpaintPipeline,
     GaudiStableDiffusionXLPipeline,
-    GaudiStableVideoDiffusionPipeline,
     GaudiStableVideoDiffusionControlNetPipeline,
+    GaudiStableVideoDiffusionPipeline,
     GaudiTextToVideoSDPipeline,
 )
 from optimum.habana.diffusers.models import (
-  ControlNetSDVModel,
-  UNetSpatioTemporalConditionControlNetModel,
+    ControlNetSDVModel,
+    UNetSpatioTemporalConditionControlNetModel,
 )
 from optimum.habana.utils import set_seed
 
@@ -2832,6 +2832,7 @@ def test_stable_video_diffusion_no_throughput_regression_bf16(self):
         if IS_GAUDI2:
             self.assertGreaterEqual(outputs.throughput, 0.95 * 0.012)
 
+
 class GaudiStableVideoDiffusionControlNetPipelineTester(TestCase):
     """
     Tests the StableVideoDiffusionControlNetPipeline for Gaudi.
@@ -2866,7 +2867,7 @@ def get_dummy_components(self):
         )
         torch.manual_seed(0)
         controlnet = ControlNetSDVModel(
-            sample_size=96,#
+            sample_size=96,  #
             in_channels=8,
             out_channels=4,
             down_block_types=(
@@ -2887,7 +2888,7 @@ def get_dummy_components(self):
             layers_per_block=2,
             cross_attention_dim=1024,
             transformer_layers_per_block=1,
-            num_attention_heads=(5, 10, 20, 20),#
+            num_attention_heads=(5, 10, 20, 20),  #
             num_frames=14,
             conditioning_channels=3,
             conditioning_embedding_out_channels=(16, 32, 96, 256),
@@ -2912,10 +2913,7 @@ def get_dummy_components(self):
         vae = AutoencoderKLTemporalDecoder(
             in_channels=3,
             out_channels=3,
-            down_block_types=('DownEncoderBlock2D',
-                              'DownEncoderBlock2D',
-                              'DownEncoderBlock2D',
-                              'DownEncoderBlock2D'),
+            down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"),
             block_out_channels=(128, 256, 512, 512),
             layers_per_block=2,
             latent_channels=4,
@@ -2984,6 +2982,7 @@ def test_stable_video_diffusion_single_video(self):
         components = self.get_dummy_components()
         gaudi_config = GaudiConfig(use_torch_autocast=True)
         sd_pipe = GaudiStableVideoDiffusionControlNetPipeline(use_habana=True, gaudi_config=gaudi_config, **components)
+
         def _get_image_from_pipeline(pipeline, device=device):
             for component in pipeline.components.values():
                 if hasattr(component, "set_default_attn_processor"):
@@ -2997,13 +2996,43 @@ def _get_image_from_pipeline(pipeline, device=device):
             self.assertEqual(len(outputs), 1)
             self.assertEqual(image.shape, (14, 3, 64, 64))
             return image[0::6, -3:, -3:, -1]
-        expected_slice = np.array([0.5463, 0.5432, 0.5576, 0.4749, 0.4736, 0.4888, 0.3865, 0.3912, 0.3869,
-                                   0.5402, 0.5389, 0.5326, 0.4747, 0.4920, 0.4760, 0.4115, 0.4248, 0.4222,
-                                   0.5359, 0.5335, 0.5221, 0.4835, 0.5044, 0.4884, 0.4266, 0.4409, 0.4295])
+
+        expected_slice = np.array(
+            [
+                0.5463,
+                0.5432,
+                0.5576,
+                0.4749,
+                0.4736,
+                0.4888,
+                0.3865,
+                0.3912,
+                0.3869,
+                0.5402,
+                0.5389,
+                0.5326,
+                0.4747,
+                0.4920,
+                0.4760,
+                0.4115,
+                0.4248,
+                0.4222,
+                0.5359,
+                0.5335,
+                0.5221,
+                0.4835,
+                0.5044,
+                0.4884,
+                0.4266,
+                0.4409,
+                0.4295,
+            ]
+        )
         image_slice = _get_image_from_pipeline(sd_pipe)
 
         self.assertLess(np.abs(image_slice.flatten() - expected_slice).max(), 1e-1)
 
+
 class GaudiStableDiffusionInstructPix2PixPipelineTests(TestCase):
     """
     Tests the class StableDiffusionInstructPix2PixPipeline for Gaudi.