FLUX IPAdapter fails when transformers are quantized

[vladmandic]

Describe the bug

Using new IP-Adapter implementation from #10336
Works if FLUX.1 model is loaded without quantization, fails on load if transformers are quantized using BnB
Not sure why conversion would require gradients? Maybe as simple as including no_grad context during conversion?

Reproduction

import torch
import diffusers

dtype = torch.bfloat16
cache_dir = '/mnt/models/Diffusers'

quantization_config = diffusers.BitsAndBytesConfig()
print('load transformer')
transformer = diffusers.FluxTransformer2DModel.from_pretrained(
    'black-forest-labs/FLUX.1-schnell',
    subfolder="transformer",
    cache_dir=cache_dir,
    torch_dtype=dtype,
    quantization_config=quantization_config
)
print('load pipe')
pipe = diffusers.FluxPipeline.from_pretrained(
    'black-forest-labs/FLUX.1-schnell',
    transformer=transformer,
    cache_dir=cache_dir,
    torch_dtype=dtype,
)
print('load ipadapter')
pipe.load_ip_adapter(
    "XLabs-AI/flux-ip-adapter",
    weight_name="ip_adapter.safetensors",
    image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14"
)

Logs

Traceback (most recent call last):
  File "/home/vlado/dev/sdnext/tmp/flux.py", line 21, in <module>
    torch_dtype=dtype,
^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/loaders/ip_adapter.py", line 553, in load_ip_adapter
    self.transformer._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/loaders/transformer_flux.py", line 168, in _load_ip_adapter_weights
    attn_procs = self._convert_ip_adapter_attn_to_diffusers(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/loaders/transformer_flux.py", line 135, in _convert_ip_adapter_attn_to_diffusers
    attn_procs[name] = attn_processor_class(
                       ^^^^^^^^^^^^^^^^^^^^^
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/attention_processor.py", line 2683, in __init__
    nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype)
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 105, in __init__
    self.weight = Parameter(
                  ^^^^^^^^^^
  File "/home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/parameter.py", line 46, in __new__
    return torch.Tensor._make_subclass(cls, data, requires_grad)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Only Tensors of floating point and complex dtype can require gradients

System Info

diffusers==0.32.0.dev from main

Who can help?

@yiyixuxu @sayakpaul @DN6 @hlky

[hlky]

Try this #10338

[vladmandic]

no luck, error remains.
note that it fails during attn_procs[name] = attn_processor_class(...), so having eval() at the end is never reached.

[Caselles]

Sorry to ask, I have the same issue,

ValueError: Casting a quantized model to a new dtype is unsupported. To set the dtype of unquantized layers, please use the torch_dtype argument when loading the model using from_pretrained or from_single_file

Can you provide a sample code that works with flux dev quantized and ip-adapter if the problem is solved? Thanks.

[hlky]

@Caselles Can you share the traceback and a reproduction?

[Caselles]

Sure. Here's the reproduction :

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import FluxTransformer2DModel, FluxPipeline
from transformers import T5EncoderModel
from diffusers.utils import load_image


quant_config = TransformersBitsAndBytesConfig(load_in_4bit=True,)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True,)

transformer_4bit = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_4bit,
    text_encoder_2=text_encoder_2_4bit,
    torch_dtype=torch.float16)

image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg").resize((1024, 1024))

pipe.load_ip_adapter(
    "XLabs-AI/flux-ip-adapter",
    weight_name="ip_adapter.safetensors",
    image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14",
    torch_dtype=torch.float16,
)

pipe.set_ip_adapter_scale(1.0)

image = pipe(
    width=1024,
    height=1024,
    prompt="wearing sunglasses",
    negative_prompt="",
    true_cfg=4.0,
    generator=torch.Generator().manual_seed(4444),
    ip_adapter_image=image,
).images[0]

image.save('flux_ip_adapter_output.jpg')

And the traceback :

`low_cpu_mem_usage` was None, now default to True since model is quantized.
Downloading shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 6825.56it/s]
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:06<00:00,  3.07s/it]
Fetching 3 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 406.39it/s]
Loading pipeline components...:  71%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                     | 5/7 [00:00<00:00, 15.18it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 12.29it/s]
Traceback (most recent call last):
  File "/home/gohu/workspace/obvious/harmonization/ipadapter_flux.py", line 36, in <module>
    pipe.load_ip_adapter(
  File "/home/gohu/anaconda3/envs/flux/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
  File "/home/gohu/anaconda3/envs/flux/lib/python3.10/site-packages/diffusers/loaders/ip_adapter.py", line 553, in load_ip_adapter
    self.transformer._load_ip_adapter_weights(state_dicts, low_cpu_mem_usage=low_cpu_mem_usage)
  File "/home/gohu/anaconda3/envs/flux/lib/python3.10/site-packages/diffusers/loaders/transformer_flux.py", line 181, in _load_ip_adapter_weights
    self.to(dtype=self.dtype, device=self.device)
  File "/home/gohu/anaconda3/envs/flux/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 1199, in to
    raise ValueError(
ValueError: Casting a quantized model to a new `dtype` is unsupported. To set the dtype of unquantized layers, please use the `torch_dtype` argument when loading the model using `from_pretrained` or `from_single_file`

[hlky]

@Caselles Fix and example #10728 examples changes are pipe.enable_model_cpu_offload() (or pipe = pipe.to('cuda') could be used) and bnb_4bit_compute_dtype in BitsAndBytesConfig is needed to avoid slow inference, there is a warning about this if it's omitted.

[Caselles]

Thanks for your answer. However I got the same error when running your code in #10728.

If this helps here's my python environment :

accelerate==1.3.0
bitsandbytes==0.45.1
certifi==2024.12.14
charset-normalizer==3.4.1
diffusers @ git+https://github.com/huggingface/diffusers@aad69ac2f323734a083d66fa89197bf7d88e5a57
filelock==3.13.1
fsspec==2024.6.1
huggingface-hub==0.28.1
idna==3.10
importlib_metadata==8.6.1
Jinja2==3.1.4
MarkupSafe==2.1.5
mpmath==1.3.0
networkx==3.3
numpy==2.1.2
nvidia-cublas-cu11==11.11.3.6
nvidia-cuda-cupti-cu11==11.8.87
nvidia-cuda-nvrtc-cu11==11.8.89
nvidia-cuda-runtime-cu11==11.8.89
nvidia-cudnn-cu11==9.1.0.70
nvidia-cufft-cu11==10.9.0.58
nvidia-curand-cu11==10.3.0.86
nvidia-cusolver-cu11==11.4.1.48
nvidia-cusparse-cu11==11.7.5.86
nvidia-nccl-cu11==2.20.5
nvidia-nvtx-cu11==11.8.86
packaging==24.2
pillow==11.0.0
protobuf==5.29.3
psutil==6.1.1
PyYAML==6.0.2
regex==2024.11.6
requests==2.32.3
safetensors==0.5.2
sentencepiece==0.2.0
sympy==1.13.1
tokenizers==0.21.0
torch==2.4.1+cu118
torchaudio==2.4.1+cu118
torchvision==0.19.1+cu118
tqdm==4.67.1
transformers==4.48.1
triton==3.0.0
typing_extensions==4.12.2
urllib3==2.3.0
zipp==3.21.0

Quick related question. Is it possible to run the IP adapter in the FluxImage2Image pipeline as well ? Is it just a matter of having the prepare embeds function and then setting self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds ? Or is it just not supported at the moment and would require more work than that ? Thanks !

[hlky]

Doesn't look like you're using the PR branch.

FluxImage2Image

#10708

[Caselles]

My bad sorry, and your PR #10708 exactly answers my question as well. Thank you so much!

[JoshonSmith]

training code for ip-adapter with Flux.1-dev?