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!
Thanks for your contributions, everyone. Let's keep this coming :) 

Resolving patch:
diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py
--- a/examples/text_to_image/train_text_to_image_lora.py
+++ b/examples/text_to_image/train_text_to_image_lora.py
@@ -46,6 +46,7 @@
 from diffusers.training_utils import compute_snr
 from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available
 from diffusers.utils.import_utils import is_xformers_available
+from diffusers.utils.torch_utils import is_compiled_module
 
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
@@ -598,6 +599,11 @@ def tokenize_captions(examples, is_train=True):
         ]
     )
 
+    def unwrap_model(model):
+        model = accelerator.unwrap_model(model)
+        model = model._orig_mod if is_compiled_module(model) else model
+        return model
+
     def preprocess_train(examples):
         images = [image.convert("RGB") for image in examples[image_column]]
         examples["pixel_values"] = [train_transforms(image) for image in images]
@@ -731,7 +737,7 @@ def collate_fn(examples):
                 noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
                 # Get the text embedding for conditioning
-                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
+                encoder_hidden_states = text_encoder(batch["input_ids"], return_dict=False)[0]
 
                 # Get the target for loss depending on the prediction type
                 if args.prediction_type is not None:
@@ -746,7 +752,7 @@ def collate_fn(examples):
                     raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
                 # Predict the noise residual and compute loss
-                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]
 
                 if args.snr_gamma is None:
                     loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
@@ -811,7 +817,7 @@ def collate_fn(examples):
                         save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                         accelerator.save_state(save_path)
 
-                        unwrapped_unet = accelerator.unwrap_model(unet)
+                        unwrapped_unet = unwrap_model(unet)
                         unet_lora_state_dict = convert_state_dict_to_diffusers(
                             get_peft_model_state_dict(unwrapped_unet)
                         )
@@ -839,7 +845,7 @@ def collate_fn(examples):
                 # create pipeline
                 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,
                     torch_dtype=weight_dtype,
@@ -880,7 +886,7 @@ def collate_fn(examples):
     if accelerator.is_main_process:
         unet = unet.to(torch.float32)
 
-        unwrapped_unet = accelerator.unwrap_model(unet)
+        unwrapped_unet = unwrap_model(unet)
         unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unwrapped_unet))
         StableDiffusionPipeline.save_lora_weights(
             save_directory=args.output_dir,


Test patch: