search

huggingface / diffusers

🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
https://huggingface.co/docs/diffusers
Apache License 2.0
25.03k stars 5.17k forks source link

StableDiffusionPipeline.from_single_file() with CKPT fails: "UnpicklingError: Weights only load failed" #8874

Open whydna opened 1 month ago

whydna commented 1 month ago

Describe the bug

Trying to load local CKPT file using the "from_single_file()" method fails. Works fine with .safetensors file from same repo (Runway ML SD).

Reproduction

from diffusers import StableDiffusionPipeline
device = 'cpu'
pipe = StableDiffusionPipeline.from_single_file("models/checkpoints/v1-5-pruned-emaonly.ckpt")
pipe = pipe.to(device)

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

Logs

UnpicklingError                           Traceback (most recent call last)
File ~/Development/ai-flow-backend/venv/lib/python3.12/site-packages/diffusers/models/model_loading_utils.py:108, in load_state_dict(checkpoint_file, variant)
    107         weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {}
--> 108         return torch.load(
    109             checkpoint_file,
    110             map_location="cpu",
    111             **weights_only_kwarg,
    112         )
    113 except Exception as e:

File ~/Development/ai-flow-backend/venv/lib/python3.12/site-packages/torch/serialization.py:1024, in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
   1023     except RuntimeError as e:
-> 1024         raise pickle.UnpicklingError(UNSAFE_MESSAGE + str(e)) from None
   1025 return _load(opened_zipfile,
   1026              map_location,
   1027              pickle_module,
   1028              overall_storage=overall_storage,
   1029              **pickle_load_args)

UnpicklingError: Weights only load failed. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution.Do it only if you get the file from a trusted source. WeightsUnpickler error: Unsupported class pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint

During handling of the above exception, another exception occurred:

UnicodeDecodeError                        Traceback (most recent call last)
...
--> 128     raise OSError(
    129         f"Unable to load weights from checkpoint file for '{checkpoint_file}' " f"at '{checkpoint_file}'. "
    130     )

OSError: Unable to load weights from checkpoint file for 'models/checkpoints/v1-5-pruned-emaonly.ckpt' at 'models/checkpoints/v1-5-pruned-emaonly.ckpt'.

System Info

Who can help?

@yiyixuxu @asomoza @DN

tolgacangoz commented 1 month ago

It says you can change weights_only to False at diffusers/models/model_loading_utils.py:107; but be careful because it is not a .safetensors file, it might run malicious codes. If the .safetensors equivalent works fine; then I would prefer it. For more info about safetensors.

whydna commented 1 month ago

doesn't look like there is a way to set that from the pipeline class.

asomoza commented 1 month ago

it's so rare nowadays that someone wants to load a ckpt on purpose that probably no one noticed this.

This is a safety feature in the code because a generative model shouldn't need to execute arbitrary code with diffusers:

https://github.com/huggingface/diffusers/blob/3b37fefee99425286984a9d5fa4f1850064d01eb/src/diffusers/models/model_loading_utils.py#L107

If you still want to load it, you can change it to False in the code but we probably won't do it as a library because of the risk. You can also downgrade your torch version.

I did it as a test and now it's asking my to install torch lighting which is weird, never loaded a ckpt with diffusers before so I don't know if that was a requirement or not.

cc: @DN6 for awareness

whydna commented 1 month ago

I think the bug is the fact that it's not able to load the weights using weights_only=True from the .ckpt checkpoint (https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned.ckpt).

If weights_only is set to False then it will try to load the model modules as well, which depending on the implementation may require pytorch lightning.

The expected behavior for the ckpt should just be the same as loading the .safetensors - which seems to be working fine.

asomoza commented 1 month ago

If you consider that a bug then you probably should report it in the model repo, probably it was saved with some weird call to pytorch lighting because of the training.

which depending on the implementation may require pytorch lightning

The safetensors alternative works because it doesn't store any library requirement/implementation and it's a fact that you don't need to install pytorch lighting to load any stable diffusion checkpoints.

But yeah, you're right. You can easily reproduce this error with this code:

import torch

state_dict = torch.load("models/checkpoints/v1-5-pruned-emaonly.ckpt", weights_only=True)

but this is not a diffusers issue then.

JemiloII commented 1 month ago

It's a change they made. 0.27.2 this worked

pipe.vae = AutoencoderKL.from_single_file(
    path.join("./assets/vae/orangemix.vae.pt"),
    local_files_only=True,
    torch_dtype=torch.bfloat16
)

But now in 0.28.0+ it is broken. Can we stop making changes that completely breaks functionality without noting it? @sayakpaul many of us still use pickle files.

Also how can I make sure all weights are loaded? Some weights of the model checkpoint were not used when initializing CLIPTextModel I'm now receiving this message. It's annoying.

yiyixuxu commented 1 month ago

we started to add weights_only=True for safety concern https://github.com/huggingface/diffusers/pull/7393

I wonder if we would allow the user to explicitly set weights_only=False when they need it cc @DN6