Diffusers v0.19.0 cause torch.compile graph break

[LSC527]

Describe the bug

Diffusers v0.19.0 cause graph breaks when using torch.compile, which have been fixed in v0.18.2 (#3286)

Reproduction

sd pipeline from #3286

Logs

No response

System Info

• diffusers version: 0.19.0
• Platform: Linux-3.10.0-1062.9.1.el7.x86_64-x86_64-with-glibc2.35
• Python version: 3.10.6
• PyTorch version (GPU?): 2.1.0a0+4136153 (True)
• Huggingface_hub version: 0.16.4
• Transformers version: 4.32.0.dev0
• Accelerate version: 0.21.0
• xFormers version: 0.0.21+b14e67d.d20230724
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

@patrickvonplaten

[patrickvonplaten]

Do you have a reproducible code snippet by any chance? Which model do you use?

[patrickvonplaten]

I can reproduce it:

#!/usr/bin/env python3
from diffusers import StableDiffusionPipeline
import torch

path = "runwayml/stable-diffusion-v1-5"

run_compile = True  # Set True / False

pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe = pipe.to("cuda:0")
pipe.unet.to(memory_format=torch.channels_last)

if run_compile:
    print("Run torch compile")
    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

prompt = "ghibli style, a fantasy landscape with castles"

for _ in range(3):
    images = pipe(prompt=prompt).images

gives:

    return self.obj.call_method(
  File "/home/patrick/hf/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 522, in call_method
    return super().call_method(tx, name, args, kwargs)
  File "/home/patrick/hf/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 253, in call_method
    raise unimplemented(f"call_method {self} {name} {args} {kwargs}")
  File "/home/patrick/hf/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 71, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: call_method NNModuleVariable() _conv_forward [TensorVariable(), TensorVariable(), TensorVariable()] {}

from user code:
   File "/home/patrick/python_bin/diffusers/models/lora.py", line 94, in forward
    return super().forward(x)
  File "/home/patrick/hf/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 463, in forward
    return self._conv_forward(input, self.weight, self.bias)

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

[patrickvonplaten]

Related to LoRA changes I think cc @sayakpaul:

from user code:
   File "/home/patrick/python_bin/diffusers/models/lora.py", line 94, in forward
    return super().forward(x)
  File "/home/patrick/hf/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 463, in forward
    return self._conv_forward(input, self.weight, self.bias)

[patrickvonplaten]

Thanks a bunch for finding this issue so quickly @LSC527 - it should be fixed with: #4315

[patrickvonplaten]

@LSC527 can you double check if it works again now on current main? Think #4315 should have fixed it :-)

[LSC527]

@patrickvonplaten yes, it works. No graph breaks now.