Repository: huggingface/diffusers

Issue:
[Training] Make officially maintained training scripts compatible with `torch.compile()`
We recently made the DreamBooth LoRA SDXL training script fully compatible with `torch.compile()` in the following PRs: https://github.com/huggingface/diffusers/pull/6483, https://github.com/huggingface/diffusers/pull/6511. 

Since there might be speedup benefits from `torch.compile()`, it makes sense to propagate these changes to the other training scripts too. We'd appreciate help from the community in this matter. 

Below is a list of the training scripts where we'd like to integrate this:

- DreamBooth LoRA
  - [x] [Regular](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py) 
  - [x] [LoRA](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora.py)
- ControlNet
  - [x] [SD](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py)
  - [x] [SDXL](https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet_sdxl.py)
- Text-to-image
  - [x] [SD](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py)
  - [x] [SD LoRA](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py)
  - [x] [SDXL](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_sdxl.py)
  - [x] [SDXL LoRA](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora_sdxl.py)
- InstructPix2Pix
  - [x] [SD](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix.py)
  - [x] [SDXL](https://github.com/huggingface/diffusers/blob/main/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py)
- T2I Adapters
  - [x] [SDXL](https://github.com/huggingface/diffusers/blob/main/examples/t2i_adapter/train_t2i_adapter_sdxl.py)

https://github.com/huggingface/diffusers/pull/6483 and https://github.com/huggingface/diffusers/pull/6511 can serve as a reference. 

Feel free to open PRs targeting only one example at a time and tag me :) 


Discussion:
@sayakpaul could you please explain how PR #6483 is connected with `torch.compile()` ? I can see that mostly the logic for predicting the noise residual has been changed and the function `encode_prompt()`.
That makes `torch.compile()` work _during_ training. Otherwise, it errors out. That is why it's connected. Your entry point to using `torch.compile()` during training is `accelerate`. While running `accelerate config` you ask it to use `torch.compile()`. See my `accelerate` configuration [here](https://github.com/huggingface/diffusers/pull/6483#issue-2069743558). 
Hi @sayakpaul , I would like to work on this issue
Feel free to pick one from the list and submit a PR. Thank you for showing interest! 
I'm currently doing the one for ControlNet, but wanted to just confirm the fact that we're delegating the actual trigger of `torch.compile()` to `accelerate`, right? As such, we do not invoke `torch.compile()` explicitly anywhere in the training script, correct? I'm following #6483 and #6511 as mentioned.
Correct!

Resolving patch:
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -55,6 +55,7 @@
 from diffusers.training_utils import compute_snr
 from diffusers.utils import check_min_version, is_wandb_available
 from diffusers.utils.import_utils import is_xformers_available
+from diffusers.utils.torch_utils import is_compiled_module
 
 
 if is_wandb_available():
@@ -129,15 +130,12 @@ def log_validation(
     if vae is not None:
         pipeline_args["vae"] = vae
 
-    if text_encoder is not None:
-        text_encoder = accelerator.unwrap_model(text_encoder)
-
     # create pipeline (note: unet and vae are loaded again in float32)
     pipeline = DiffusionPipeline.from_pretrained(
         args.pretrained_model_name_or_path,
         tokenizer=tokenizer,
         text_encoder=text_encoder,
-        unet=accelerator.unwrap_model(unet),
+        unet=unet,
         revision=args.revision,
         variant=args.variant,
         torch_dtype=weight_dtype,
@@ -794,6 +792,7 @@ def encode_prompt(text_encoder, input_ids, attention_mask, text_encoder_use_atte
     prompt_embeds = text_encoder(
         text_input_ids,
         attention_mask=attention_mask,
+        return_dict=False,
     )
     prompt_embeds = prompt_embeds[0]
 
@@ -931,11 +930,16 @@ def main(args):
         args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, variant=args.variant
     )
 
+    def unwrap_model(model):
+        model = accelerator.unwrap_model(model)
+        model = model._orig_mod if is_compiled_module(model) else model
+        return model
+
     # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
     def save_model_hook(models, weights, output_dir):
         if accelerator.is_main_process:
             for model in models:
-                sub_dir = "unet" if isinstance(model, type(accelerator.unwrap_model(unet))) else "text_encoder"
+                sub_dir = "unet" if isinstance(model, type(unwrap_model(unet))) else "text_encoder"
                 model.save_pretrained(os.path.join(output_dir, sub_dir))
 
                 # make sure to pop weight so that corresponding model is not saved again
@@ -946,7 +950,7 @@ def load_model_hook(models, input_dir):
             # pop models so that they are not loaded again
             model = models.pop()
 
-            if isinstance(model, type(accelerator.unwrap_model(text_encoder))):
+            if isinstance(model, type(unwrap_model(text_encoder))):
                 # load transformers style into model
                 load_model = text_encoder_cls.from_pretrained(input_dir, subfolder="text_encoder")
                 model.config = load_model.config
@@ -991,15 +995,12 @@ def load_model_hook(models, input_dir):
         " doing mixed precision training. copy of the weights should still be float32."
     )
 
-    if accelerator.unwrap_model(unet).dtype != torch.float32:
-        raise ValueError(
-            f"Unet loaded as datatype {accelerator.unwrap_model(unet).dtype}. {low_precision_error_string}"
-        )
+    if unwrap_model(unet).dtype != torch.float32:
+        raise ValueError(f"Unet loaded as datatype {unwrap_model(unet).dtype}. {low_precision_error_string}")
 
-    if args.train_text_encoder and accelerator.unwrap_model(text_encoder).dtype != torch.float32:
+    if args.train_text_encoder and unwrap_model(text_encoder).dtype != torch.float32:
         raise ValueError(
-            f"Text encoder loaded as datatype {accelerator.unwrap_model(text_encoder).dtype}."
-            f" {low_precision_error_string}"
+            f"Text encoder loaded as datatype {unwrap_model(text_encoder).dtype}." f" {low_precision_error_string}"
         )
 
     # Enable TF32 for faster training on Ampere GPUs,
@@ -1246,7 +1247,7 @@ def compute_text_embeddings(prompt):
                         text_encoder_use_attention_mask=args.text_encoder_use_attention_mask,
                     )
 
-                if accelerator.unwrap_model(unet).config.in_channels == channels * 2:
+                if unwrap_model(unet).config.in_channels == channels * 2:
                     noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)
 
                 if args.class_labels_conditioning == "timesteps":
@@ -1256,8 +1257,8 @@ def compute_text_embeddings(prompt):
 
                 # Predict the noise residual
                 model_pred = unet(
-                    noisy_model_input, timesteps, encoder_hidden_states, class_labels=class_labels
-                ).sample
+                    noisy_model_input, timesteps, encoder_hidden_states, class_labels=class_labels, return_dict=False
+                )[0]
 
                 if model_pred.shape[1] == 6:
                     model_pred, _ = torch.chunk(model_pred, 2, dim=1)
@@ -1350,9 +1351,9 @@ def compute_text_embeddings(prompt):
 
                     if args.validation_prompt is not None and global_step % args.validation_steps == 0:
                         images = log_validation(
-                            text_encoder,
+                            unwrap_model(text_encoder) if text_encoder is not None else text_encoder,
                             tokenizer,
-                            unet,
+                            unwrap_model(unet),
                             vae,
                             args,
                             accelerator,
@@ -1375,14 +1376,14 @@ def compute_text_embeddings(prompt):
         pipeline_args = {}
 
         if text_encoder is not None:
-            pipeline_args["text_encoder"] = accelerator.unwrap_model(text_encoder)
+            pipeline_args["text_encoder"] = unwrap_model(text_encoder)
 
         if args.skip_save_text_encoder:
             pipeline_args["text_encoder"] = None
 
         pipeline = DiffusionPipeline.from_pretrained(
             args.pretrained_model_name_or_path,
-            unet=accelerator.unwrap_model(unet),
+            unet=unwrap_model(unet),
             revision=args.revision,
             variant=args.variant,
             **pipeline_args,


Test patch: