[Text-to-video] Add torch.compile() compatibility

[sayakpaul]

What does this PR do?

Fixes #3915

Description

torch.compile() for the repeat_interleave() function was added in a nightly build. See: pytorch/pytorch#99929.

So, once I upgraded to Torch 2.1 nightly, the issue went away. However, there were other issues which are fixed in this PR. The PR takes inspiration from #3313.

Even though we're able to successfully compile the model, it takes a hefty amount of time after torch.compile() is called on the UNet:

import torch
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video
from PIL import Image


pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16)
pipe.to("cuda")
pipe.enable_vae_slicing()

pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

prompt = "Darth Vader is surfing on waves"
video_frames = pipe(prompt, num_inference_steps=40, height=320, width=576, num_frames=36).frames
video_path = export_to_video(video_frames, output_video_path="video_576_darth_vader_36.mp4")

The first call pipe is really time-consuming which is understandable because that is when the compiled UNet model is also used for the first time. But even in the subsequent calls, the timing doesn't seem to improve much. In my experiments, I actually found the runtime to be performing much better without torch.compile().

Let me know if anything is unclear.

I leaving the outputs of the progress bars:

With torch.compile()

Without torch.compile()

My explorations can be found in this Colab Notebook.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Great job! Looks clean - also ok for me to not add a test at the moment, given the problems we have with memory leaks. Good to merge for me

[sayakpaul]

Do we have a handle on why these leaks happen? Why does torch.compile() perform so poorly when we 3D inputs like videos?