enable_sequential_cpu_offload followed by load_textual_inversion does not handle placement of loaded weights

[JeLuF]

Describe the bug

Using pipe.enable_sequential_cpu_offload() and pipe.load_textual_inversion() together crashes diffusers.

Reproduction

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",  
    torch_dtype=torch.float16,
)

pipe.enable_sequential_cpu_offload()
pipe.load_textual_inversion("./happy512.pt", token="happy512")
prompt = "girl, happy512"

image = pipe(prompt).images[0]

Logs

(D:\2.35\dev\stable-diffusion\env) D:\2.35\dev\models\embeddings>python x.py
safety_checker\model.safetensors not found
`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config["id2label"]` will be overriden.
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ D:\2.35\dev\models\embeddings\x.py:13 in <module>                                                │
│                                                                                                  │
│   10 pipe.load_textual_inversion("./happy512.pt", token="happy512")                              │
│   11 prompt = "girl, happy512"                                                                   │
│   12                                                                                             │
│ ❱ 13 image = pipe(prompt).images[0]                                                              │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\utils\_contextlib.py:115 in             │
│ decorate_context                                                                                 │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\diffusers\pipelines\stable_diffusion\pipeline │
│ _stable_diffusion.py:688 in __call__                                                             │
│                                                                                                  │
│   685 │   │   text_encoder_lora_scale = (                                                        │
│   686 │   │   │   cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not N   │
│   687 │   │   )                                                                                  │
│ ❱ 688 │   │   prompt_embeds = self._encode_prompt(                                               │
│   689 │   │   │   prompt,                                                                        │
│   690 │   │   │   device,                                                                        │
│   691 │   │   │   num_images_per_prompt,                                                         │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\diffusers\pipelines\stable_diffusion\pipeline │
│ _stable_diffusion.py:382 in _encode_prompt                                                       │
│                                                                                                  │
│   379 │   │   │   else:                                                                          │
│   380 │   │   │   │   attention_mask = None                                                      │
│   381 │   │   │                                                                                  │
│ ❱ 382 │   │   │   prompt_embeds = self.text_encoder(                                             │
│   383 │   │   │   │   text_input_ids.to(device),                                                 │
│   384 │   │   │   │   attention_mask=attention_mask,                                             │
│   385 │   │   │   )                                                                              │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\nn\modules\module.py:1501 in _call_impl │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\accelerate\hooks.py:165 in new_forward        │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\transformers\models\clip\modeling_clip.py:816 │
│ in forward                                                                                       │
│                                                                                                  │
│    813 │   │   """                                                                            │
│    814 │   │   return_dict = return_dict if return_dict is not None else self.config.use_return  │
│    815 │   │                                                                                     │
│ ❱  816 │   │   return self.text_model(                                                           │
│    817 │   │   │   input_ids=input_ids,                                                          │
│    818 │   │   │   attention_mask=attention_mask,                                                │
│    819 │   │   │   position_ids=position_ids,                                                    │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\nn\modules\module.py:1501 in _call_impl │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\transformers\models\clip\modeling_clip.py:712 │
│ in forward                                                                                       │
│                                                                                                  │
│    709 │   │   input_shape = input_ids.size()                                                    │
│    710 │   │   input_ids = input_ids.view(-1, input_shape[-1])                                   │
│    711 │   │                                                                                     │
│ ❱  712 │   │   hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)   │
│    713 │   │                                                                                     │
│    714 │   │   bsz, seq_len = input_shape                                                        │
│    715 │   │   # CLIP's text model uses causal mask, prepare it here.                            │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\nn\modules\module.py:1501 in _call_impl │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\accelerate\hooks.py:165 in new_forward        │
│                                                                                                  │
│   162 │   │   │   with torch.no_grad():                                                          │
│   163 │   │   │   │   output = old_forward(*args, **kwargs)                                      │
│   164 │   │   else:                                                                              │
│ ❱ 165 │   │   │   output = old_forward(*args, **kwargs)                                          │
│   166 │   │   return module._hf_hook.post_forward(module, output)                                │
│   167 │                                                                                          │
│   168 │   module.forward = new_forward                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\transformers\models\clip\modeling_clip.py:230 │
│ in forward                                                                                       │
│                                                                                                  │
│    227 │   │   │   inputs_embeds = self.token_embedding(input_ids)                               │
│    228 │   │                                                                                     │
│    229 │   │   position_embeddings = self.position_embedding(position_ids)                       │
│ ❱  230 │   │   embeddings = inputs_embeds + position_embeddings                                  │
│    231 │   │                                                                                     │
│    232 │   │   return embeddings                                                                 │
│    233                                                                                           │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_prims_common\wrappers.py:220 in _fn    │
│                                                                                                  │
│   217 │   │   │   │   │   if k not in kwargs:                                                    │
│   218 │   │   │   │   │   │   kwargs[k] = out_attr                                               │
│   219 │   │   │                                                                                  │
│ ❱ 220 │   │   │   result = fn(*args, **kwargs)                                                   │
│   221 │   │   │   assert (                                                                       │
│   222 │   │   │   │   isinstance(result, TensorLike)                                             │
│   223 │   │   │   │   and is_tensor                                                              │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_prims_common\wrappers.py:130 in _fn    │
│                                                                                                  │
│   127 │   │   │   }                                                                              │
│   128 │   │   │   bound.arguments.update(promoted_args)                                          │
│   129 │   │   │                                                                                  │
│ ❱ 130 │   │   │   result = fn(**bound.arguments)                                                 │
│   131 │   │   │                                                                                  │
│   132 │   │   │   if isinstance(result, TensorLike):                                             │
│   133 │   │   │   │   return _maybe_convert_to_dtype(result, result_dtype)                       │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_refs\__init__.py:975 in add            │
│                                                                                                  │
│    972 │   │   │   raise ValueError(msg)                                                         │
│    973 │   │   b = prims.mul(b, alpha)                                                           │
│    974 │                                                                                         │
│ ❱  975 │   return prims.add(a, b)                                                                │
│    976                                                                                           │
│    977                                                                                           │
│    978 # TODO: add docstring                                                                     │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_ops.py:287 in __call__                 │
│                                                                                                  │
│   284 │   │   )                                                                                  │
│   285 │                                                                                          │
│   286 │   def __call__(self, *args, **kwargs):                                                   │
│ ❱ 287 │   │   return self._op(*args, **kwargs or {})                                             │
│   288 │                                                                                          │
│   289 │   def __hash__(self):                                                                    │
│   290 │   │   return hash(self._op)                                                              │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_prims\__init__.py:346 in               │
│ _elementwise_meta                                                                                │
│                                                                                                  │
│    343 │   if args_with_fixed_dtypes is not None:                                                │
│    344 │   │   args_ = list(args_with_fixed_dtypes) + args_                                      │
│    345 │                                                                                         │
│ ❱  346 │   utils.check_same_device(*args_, allow_cpu_scalar_tensors=True)                        │
│    347 │   utils.check_same_shape(*args_, allow_cpu_scalar_tensors=True)                         │
│    348 │                                                                                         │
│    349 │   strides = utils.compute_elementwise_output_strides(*args_)                            │
│                                                                                                  │
│ D:\2.35\dev\stable-diffusion\env\lib\site-packages\torch\_prims_common\__init__.py:596 in        │
│ check_same_device                                                                                │
│                                                                                                  │
│    593 │   │   │   │   │   + str(device)                                                         │
│    594 │   │   │   │   │   + "!"                                                                 │
│    595 │   │   │   │   )                                                                         │
│ ❱  596 │   │   │   │   raise RuntimeError(msg)                                                   │
│    597 │   │   else:                                                                             │
│    598 │   │   │   msg = (                                                                       │
│    599 │   │   │   │   "Unexpected type when checking for same device, " + str(type(arg)) + "!"  │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Tensor on device cuda:0 is not on the expected device meta!

System Info

Diffusers 0.17.1
Python 3.8.5
Windows 10

Who can help?

@williamberman, @patrickvonplaten @sayakpaul

[sayakpaul]

Cc: @patrickvonplaten @pcuenca

From the usage example (here), even without CPU offloading, it seems to error out. Colab Notebook: https://colab.research.google.com/gist/sayakpaul/993db4af51e0e3203aff1507619b77ef/scratchpad.ipynb.

[JeLuF]

@sayakpaul

wget https://civitai.com/api/download/models/9857 -O charturnerv2.pt

9857 is the SD2.1 embedding, you can't use that one with an SD1.5 model. Use 8387 instead, that's the SD1.5 version of charturner_v2.

[pcuenca]

Hi @JeLuF, you are right, the textual inversion weights are not moved to the appropriate device when loaded. Thanks for spotting it!

One workaround while this is fixed is to reverse the call order, the following should work:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",  
    torch_dtype=torch.float16,
)

pipe.load_textual_inversion("./happy512.pt", token="happy512")
pipe.enable_sequential_cpu_offload()
prompt = "girl, happy512"

image = pipe(prompt).images[0]

In addition, I'd recommend using enable_model_cpu_offload() if you can. Just in case you didn't know about it, it's much faster but takes more memory than enable_sequential_cpu_offload(). It also is more robust to this particular scenario, there's no need to reverse the calls when using it.

[JeLuF]

The above code is not from our actual code base, but was intended to serve as a minimal example showing the issue.

The actual code we use is much more complex. We offer multiple modes that the user can switch between. The "Low VRAM usage" mode is optimized to work on GPUs with as little as 2GB. "Faster but uses more VRAM" is a different mode, in which load_textual_inversion() works fine.

Swapping the order would require a very complex architecture change. We load the model and then load the embeddings based on the user's prompt input. If the user adds more embeddings to the prompt, those will be added to the already loaded model. Enabling the CPU offloading after loading the embeddings would require a reload of the model every time the user adds one to the prompt. I'd like to avoid this, so I have to wait for a fix.

[patrickvonplaten]

I think we can fix both this and #3958 by make sure the newly added weights (whether it's textual inversion or LoRA) are correctly placed on the devices. It means that we have to add hooks to the newly added weights just like there were before. @sayakpaul do you want to give this a try?

[sayakpaul]

@patrickvonplaten pardon my unwisdom, but the cpu_offload() and cpu_offload_with_hook() methods from accelerate accept nn.Modules as inputs but these are weights (non nn.Modules) we're talking about when dealing with textual inversion or LoRA.

What do you have in mind?

I would suggest the following:

• we detect if the concerned module (text encoder for textual inversion) has _hf_hook and detect its type (accelerate.hooks.CpuOffload and accelerate.hooks.AlignDevicesHook).
• based on the type of offload we provide instructions in the error message.

It's simple and does the job, no?

[patrickvonplaten]

@sayakpaul do you think fundamentally it won't work here? Textual inversion should surely work with offloading - here I think we just have to do the following:

• 1.) When calling load_textual_inversion, we a) detect if there are offloading hooks, b) if this is the case we remove all the offloading hooks c) then we load all weights and d) We call cpu_offload() again

[patrickvonplaten]

• 2.) LoRA might be harder since we do monkey patching, but I think just removing offload hooks and re-applying them could/should work as well no? Also cc @williamberman here - could you maybe give this a look?

[williamberman]

The safest way to deal with accelerate hooks and adding new weights is going to be removing all hooks, adding the new weights, and then re-applying the hooks. This should be cheap so-long as the model stays on the cpu the whole time and we don't have to move any weights to or from the gpu (which I think is the way it works today).

This will work with the current LoRA weights as is because the patched just swaps out the linear nn.Module for the patched linear nn.Module, i.e. from the view of accelerate, there's nothing abnormal about the model.

@patrickvonplaten are you asking for just my opinion or should I implement this? If the latter, it should only be implemented for load_lora_weights and load_textual_inversion, yes?

[sayakpaul]

@williamberman I think Patrick was asking for your inputs. I will give this a try.

[patrickvonplaten]

Great then I think we should try to offload all hooks and later load them again when loading textual inversion and/or loading LoRA