[Enhance] Support fast norm controlnet

diffengine/engine/hooks/fast_norm_hook.py

@@ -101,7 +101,12 @@ def before_train(self, runner) -> None:
             model = model.module
         if self.fuse_unet_ln:
             self._replace_ln(model.unet, "model", model.device)
+            if hasattr(model, "controlnet"):
+                self._replace_ln(model.controlnet, "model", model.device)
+
         self._replace_gn_forward(model.unet, "unet")
+        if hasattr(model, "controlnet"):
+            self._replace_gn_forward(model.controlnet, "unet")
 
         if self.fuse_text_encoder_ln:
             if hasattr(model, "text_encoder"):

tests/test_engine/test_hooks/test_fast_norm_hook.py

@@ -8,9 +8,11 @@
 from diffengine.engine.hooks import FastNormHook
 from diffengine.models.editors import (
     SDDataPreprocessor,
+    SDXLControlNetDataPreprocessor,
     SDXLDataPreprocessor,
     StableDiffusion,
     StableDiffusionXL,
+    StableDiffusionXLControlNet,
 )
 from diffengine.models.losses import L2Loss
 from diffengine.models.utils import WhiteNoise
@@ -31,6 +33,12 @@ def setUp(self) -> None:
             name="SDDataPreprocessor", module=SDDataPreprocessor)
         MODELS.register_module(
             name="SDXLDataPreprocessor", module=SDXLDataPreprocessor)
+        MODELS.register_module(
+            name="StableDiffusionXLControlNet",
+            module=StableDiffusionXLControlNet)
+        MODELS.register_module(
+            name="SDXLControlNetDataPreprocessor",
+            module=SDXLControlNetDataPreprocessor)
         MODELS.register_module(name="L2Loss", module=L2Loss)
         MODELS.register_module(name="WhiteNoise", module=WhiteNoise)
         return super().setUp()
@@ -40,6 +48,8 @@ def tearDown(self) -> None:
         MODELS.module_dict.pop("StableDiffusionXL")
         MODELS.module_dict.pop("SDDataPreprocessor")
         MODELS.module_dict.pop("SDXLDataPreprocessor")
+        MODELS.module_dict.pop("StableDiffusionXLControlNet")
+        MODELS.module_dict.pop("SDXLControlNetDataPreprocessor")
         MODELS.module_dict.pop("L2Loss")
         MODELS.module_dict.pop("WhiteNoise")
         return super().tearDown()
@@ -98,3 +108,37 @@ def test_before_train(self) -> None:
         assert isinstance(
             runner.model.text_encoder_two.text_model.encoder.layers[
                 0].layer_norm1, FusedLayerNorm)
+
+        # Test StableDiffusionXLControlNet
+        cfg = copy.deepcopy(self.epoch_based_cfg)
+        cfg.model.type = "StableDiffusionXLControlNet"
+        cfg.model.model = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+        cfg.model.controlnet_model = "hf-internal-testing/tiny-controlnet-sdxl"
+        runner = self.build_runner(cfg)
+        hook = FastNormHook(fuse_text_encoder_ln=True)
+        assert isinstance(
+            runner.model.unet.down_blocks[
+                1].attentions[0].transformer_blocks[0].norm1, nn.LayerNorm)
+        assert isinstance(
+            runner.model.controlnet.down_blocks[
+                1].attentions[0].transformer_blocks[0].norm1, nn.LayerNorm)
+        assert isinstance(
+            runner.model.text_encoder_one.text_model.encoder.layers[
+                0].layer_norm1, nn.LayerNorm)
+        assert isinstance(
+            runner.model.text_encoder_two.text_model.encoder.layers[
+                0].layer_norm1, nn.LayerNorm)
+        # replace norm
+        hook.before_train(runner)
+        assert isinstance(
+            runner.model.unet.down_blocks[
+                1].attentions[0].transformer_blocks[0].norm1, FusedLayerNorm)
+        assert isinstance(
+            runner.model.controlnet.down_blocks[
+                1].attentions[0].transformer_blocks[0].norm1, FusedLayerNorm)
+        assert isinstance(
+            runner.model.text_encoder_one.text_model.encoder.layers[
+                0].layer_norm1, FusedLayerNorm)
+        assert isinstance(
+            runner.model.text_encoder_two.text_model.encoder.layers[
+                0].layer_norm1, FusedLayerNorm)