[Torch.compile] Fixes torch compile graph break

[patrickvonplaten]

For some reason one cannot create a parent class of nn.Conv2d and then call super().forward(...) in it with torch.compile.

See issue: #4305

[patrickvonplaten]

@Chillee to me this looks like something that should be fixed in downstream PyTorch.

This is more of a quick'n'dirty fix.

To reproduce:

pip install diffusers==0.19.0 transformers

and then run:

#!/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

which 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

which is related to the changes done in this PR. This PR fixes the problem, but it's not clear why the previous code did not work out of the box for PyTorch.

[HuggingFaceDocBuilderDev]

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

[MikeHanKK]

It looks like this fix will break the seamless texture solution by setting conv2d.init padding to "circular".
see #556
@patrickvonplaten

In lora.py, forward(), when lora_layer is None, F.conv() doesn't handle padding as super().forward does.

[kolar1988]

It looks like this fix will break the seamless texture solution by setting conv2d.init padding to "circular". see #556 @patrickvonplaten

In lora.py, forward(), when lora_layer is None, F.conv() doesn't handle padding as super().forward does.

yes, "F.conv()" causes the problem the seamless texture. cannot set the parameter "padding_mode". @patrickvonplaten

[patrickvonplaten]

Aiii ok, we can probably revert this change once Torch 2.1 is out as this has been natively fixed in PyTorch now: pytorch/pytorch#106402

Can you maybe open a new issue here so that we can track the issue?