[rfc][compile] compile method for DiffusionPipeline

Author: anijain2305

src/diffusers/models/transformers/transformer_flux.py

@@ -286,6 +286,8 @@ def __init__(
 
         self.gradient_checkpointing = False
 
+        self.compile_region_classes = (FluxTransformerBlock, FluxSingleTransformerBlock)
+
     @property
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
     def attn_processors(self) -> Dict[str, AttentionProcessor]:

src/diffusers/models/transformers/transformer_wan.py

@@ -403,6 +403,8 @@ def __init__(
 
         self.gradient_checkpointing = False
 
+        self.compile_region_classes = (WanTransformerBlock,)
+
     def forward(
         self,
         hidden_states: torch.Tensor,

src/diffusers/pipelines/pipeline_utils.py

@@ -2027,6 +2027,39 @@ def _maybe_raise_error_if_group_offload_active(
                 return True
         return False
 
+    def compile(
+        self,
+        compile_regions_for_transformer: bool = True,
+        transformer_module_name: str = "transformer",
+        other_modules_names: List[str] = [],
+        **compile_kwargs,
+    ):
+        transformer = getattr(self, transformer_module_name, None)
+        if transformer is None:
+            raise ValueError(
+                f"{transformer_module_name} not found in the pipeline. Set `transformer_module_name` to the correct module name."
+            )
+
+        if compile_regions_for_transformer:
+            compile_region_classes = getattr(transformer, "compile_region_classes", None)
+            if compile_region_classes is None:
+                raise ValueError(
+                    f"{transformer_module_name} does not have `compile_region_classes` attribute. Set `compile_regions_for_transformer` to False."
+                )
+
+            for submod in transformer.modules():
+                if isinstance(submod, compile_region_classes):
+                    submod.compile(**compile_kwargs)
+        else:
+            transformer.compile(**compile_kwargs)
+
+        for module_name in other_modules_names:
+            module = getattr(self, module_name, None)
+            if module is None:
+                raise ValueError(
+                    f"{module_name} not found in the pipeline. Set `other_modules_names` to the correct module names."
+                )
+
 
 class StableDiffusionMixin:
     r"""