Repository: huggingface/diffusers

Issue:
DownBlockMotion performs different with or without checkpointing 
### Describe the bug

DownBlockMotion, UpBlockMotion can be normally called without gradient_checkpointing but cant with it.
something seems weird here:
https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_3d_blocks.py
line 1034~1043
```python
            if self.training and self.gradient_checkpointing:
                ......
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(motion_module),
                    hidden_states.requires_grad_(),
                    temb,                                                                                                   
                    num_frames,
                )#has temb here
            else:
                hidden_states = resnet(hidden_states, temb, scale=scale)
                hidden_states = motion_module(hidden_states, num_frames=num_frames)[0] # no temb here
```
it's strange that motion_module is self-attention during non-checkpointed training,and is cross-attention during checkpointed training.

### Reproduction

import diffusers,torch
from diffusers.models.unet_3d_blocks import DownBlockMotion as B
b=B(32,32,32)
b(torch.zeros(1,32,4,4),torch.zeros(1,32))[0].flatten()[:10] # everything ok here
b.gradient_checkpointing=True
b(torch.zeros(1,32,4,4),torch.zeros(1,32))[0].flatten()[:10] # raise ValueError: not enough values to unpack (expected 3, got 2)

### Logs

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
 in 
     25 b.gradient_checkpointing=True
     26 
---> 27 b(torch.zeros(1,32,4,4),torch.zeros(1,32))[0].flatten()[:10]

c:\Users\Administrator\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

c:\Users\Administrator\anaconda3\lib\site-packages\diffusers\models\unet_3d_blocks.py in forward(self, hidden_states, temb, scale, num_frames)
   1032                         create_custom_forward(resnet), hidden_states, temb, scale
   1033                     )
-> 1034                 hidden_states = torch.utils.checkpoint.checkpoint(
   1035                     create_custom_forward(motion_module),
   1036                     hidden_states.requires_grad_(),

c:\Users\Administrator\anaconda3\lib\site-packages\torch\utils\checkpoint.py in checkpoint(function, use_reentrant, *args, **kwargs)
    247 
    248     if use_reentrant:
--> 249         return CheckpointFunction.apply(function, preserve, *args)
...
-> 1225         batch_size, sequence_length, _ = (
   1226             hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
   1227         )

ValueError: not enough values to unpack (expected 3, got 2)
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...
```


### System Info

- `diffusers` version: 0.25.1
- Platform: Windows-10-10.0.19041-SP0       
- Python version: 3.8.8
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Huggingface_hub version: 0.20.3
- Transformers version: 4.37.0
- Accelerate version: 0.26.1
- xFormers version: not installed
- Using GPU in script?: both are same
- Using distributed or parallel set-up in script?: Nope

### Who can help?

@DN6 @sayakpaul @patrickvonplaten


Discussion:
Would you want to open a PR to fix the issue? 
similar problems exist in CrossAttnDownBlockMotion & UpBlockMotion:
https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_3d_blocks.py line 1214:
```python
        for i, (resnet, attn, motion_module) in enumerate(blocks):
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
                )[0]
            else:
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
                )[0]
                hidden_states = motion_module(
                    hidden_states,
                    num_frames=num_frames,
                )[0]
```
WHERE ARE THE MOTION MODULES WHEN gradient_checkpointing????
I am asking you if you'd be open to fix that and create a PR. 

Resolving patch:
diff --git a/src/diffusers/models/unets/unet_3d_blocks.py b/src/diffusers/models/unets/unet_3d_blocks.py
--- a/src/diffusers/models/unets/unet_3d_blocks.py
+++ b/src/diffusers/models/unets/unet_3d_blocks.py
@@ -1031,16 +1031,10 @@ def custom_forward(*inputs):
                     hidden_states = torch.utils.checkpoint.checkpoint(
                         create_custom_forward(resnet), hidden_states, temb, scale
                     )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(motion_module),
-                    hidden_states.requires_grad_(),
-                    temb,
-                    num_frames,
-                )
 
             else:
                 hidden_states = resnet(hidden_states, temb, scale=scale)
-                hidden_states = motion_module(hidden_states, num_frames=num_frames)[0]
+            hidden_states = motion_module(hidden_states, num_frames=num_frames)[0]
 
             output_states = output_states + (hidden_states,)
 
@@ -1221,10 +1215,10 @@ def custom_forward(*inputs):
                     encoder_attention_mask=encoder_attention_mask,
                     return_dict=False,
                 )[0]
-                hidden_states = motion_module(
-                    hidden_states,
-                    num_frames=num_frames,
-                )[0]
+            hidden_states = motion_module(
+                hidden_states,
+                num_frames=num_frames,
+            )[0]
 
             # apply additional residuals to the output of the last pair of resnet and attention blocks
             if i == len(blocks) - 1 and additional_residuals is not None:
@@ -1425,10 +1419,10 @@ def custom_forward(*inputs):
                     encoder_attention_mask=encoder_attention_mask,
                     return_dict=False,
                 )[0]
-                hidden_states = motion_module(
-                    hidden_states,
-                    num_frames=num_frames,
-                )[0]
+            hidden_states = motion_module(
+                hidden_states,
+                num_frames=num_frames,
+            )[0]
 
         if self.upsamplers is not None:
             for upsampler in self.upsamplers:
@@ -1563,15 +1557,10 @@ def custom_forward(*inputs):
                     hidden_states = torch.utils.checkpoint.checkpoint(
                         create_custom_forward(resnet), hidden_states, temb
                     )
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(resnet),
-                    hidden_states,
-                    temb,
-                )
 
             else:
                 hidden_states = resnet(hidden_states, temb, scale=scale)
-                hidden_states = motion_module(hidden_states, num_frames=num_frames)[0]
+            hidden_states = motion_module(hidden_states, num_frames=num_frames)[0]
 
         if self.upsamplers is not None:
             for upsampler in self.upsamplers:


Test patch: