Support torch._subclksses.EulerDiscreteScheduler scheduler

[thiagocrepaldi]

Describe the bug

Slightly related to huggingface/accelerate#2494 in the sense it applies the suggested workaround to avoid the big modeling feature (not supported by FakeTensorMode).

Once the model is initialized, it hits an issue within EulerDiscreteScheduler when numpy.array is called to create sigmas tensor. FakeTensorMode does not support numpy.array so ideally we should rewrite the code without using this call whener is possible.

Reproduction

The repro also has another workaround (PR #318 that was rejected, but a better solution never found) for now

from huggingface_hub import hf_hub_download
import torch
from diffusers import DiffusionPipeline
import safetensors
import transformers

# Monkey-patch for https://github.com/huggingface/safetensors/pull/318
class ONNXTorchPatcher:

    def __init__(self):
        def safetensors_load_file_wrapper(filename, device="cpu"):
            result = {}
            with safetensors.torch.safe_open(  # type: ignore[attr-defined]
                filename, framework="pt", device=device
            ) as f:
                for k in f.keys():
                    fake_mode = torch._guards.detect_fake_mode()
                    if not fake_mode:
                        result[k] = f.get_tensor(k)
                    else:
                        empty_tensor = f.get_slice(k)
                        result[k] = torch.empty(
                            tuple(empty_tensor.get_shape()),
                            dtype=safetensors.torch._getdtype(
                                empty_tensor.get_dtype()
                            ),
                        )
            return result

        self.safetensors_torch_load_file = safetensors.torch.load_file
        self.safetensors_torch_load_file_wrapper = safetensors_load_file_wrapper
        self.transformers_modeling_utils_safe_load_file = (
            transformers.modeling_utils.safe_load_file
        )

    def __enter__(self):
        safetensors.torch.load_file = self.safetensors_torch_load_file_wrapper
        transformers.modeling_utils.safe_load_file = (
            self.safetensors_torch_load_file_wrapper
        )

    def __exit__(self, exc_type, exc_value, traceback):
        safetensors.torch.load_file = self.safetensors_torch_load_file
        transformers.modeling_utils.safe_load_file = (
            self.transformers_modeling_utils_safe_load_file
        )

with torch._subclasses.FakeTensorMode():
    fake_model = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", low_cpu_mem_usage=False)

The expected behavior is that DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", low_cpu_mem_usage=False) should succeed under torch._subclasses.fake_tensor.FakeTensorMode

Logs

No response

System Info

diffuser main branch

Who can help?

No response

[sayakpaul]

Why is there a need to use torch._subclasses.fake_tensor.FakeTensorMode here?

[thiagocrepaldi]

Why is there a need to use torch._subclasses.fake_tensor.FakeTensorMode here?

PyTorch's FakeTensor is a feature to load model without weights, which is needed for large models.
The PyTorch ONNX exporter leverages fake tensor to load any torch model without allocating memory with a single API/entry point, including HF's models

[sayakpaul]

We cannot change one scheduler implementation to support this as this breaks consistency across the schedulers.

So, would like to take @yiyixuxu’s opinions as well as @echarlaix’s.

[thiagocrepaldi]

We could change other schedulers incrementally.
maybe we could contrib all schedulers withut use of numpy. The one issue I have found for and which justified not doing all at once is that some schedulers use np.interp which is not available in pytorch. Refer to this comment

[yiyixuxu]

@sayakpaul
I'm totally ok with this change here https://github.com/huggingface/diffusers/pull/7151/files
why not? I don't think it makes any difference, no? it seems to make the code simpler actually

[sayakpaul]

We cannot change one scheduler implementation to support this as this breaks consistency across the schedulers.

[yiyixuxu]

what consistency?
the results would be the same

[yiyixuxu]

not sure about other schedulers though because it seems like we need to add a bunch of code
I prefer not to do anything for other schedulers if we don't have to

[sayakpaul]

If someone wants to use DPM with SDXL for the use case, it would break similarly without the solution from #7151.

[yiyixuxu]

If someone wants to use DPM with SDXL for the use case, it would break similarly without the solution from #7151.

we can worry about that later.
I think that change made the code better even if it wasn't for this - we convert torch into numpy and immediately convert back, did not have to do that, no?

[sayakpaul]

Yeah sure. I am okay with the change then :-)