[wip] Higgs quantization

[sayakpaul]

What does this PR do?

Testing code (doesn't work):

from diffusers import DiffusionPipeline, AutoModel
from diffusers.quantizers.quantization_config import HiggsConfig
import torch 

ckpt_id = "black-forest-labs/FLUX.1-dev"
transformer = AutoModel.from_pretrained(
    ckpt_id, 
    subfolder="transformer", 
    quantization_config=HiggsConfig(modules_to_not_convert=["x_embedder"]), 
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)
pipe = DiffusionPipeline.from_pretrained(
    ckpt_id, 
    transformer=transformer,
    torch_dtype=torch.bfloat16
).to("cuda")

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 28,
    "max_sequence_length": 512,
    "generator": torch.manual_seed(0),
}
image = pipe(**pipe_kwargs).images[0]
image.save("higgs.png")

Errors out with:

in_features=4096, out_features=3072
torch.Size([768, 4096])
Loading checkpoint shards:   0%|                                                                                     | 0/3 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/home/sayak/diffusers/check_higgs.py", line 6, in <module>
    transformer = AutoModel.from_pretrained(
  File "/home/sayak/.pyenv/versions/higgs/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
  File "/home/sayak/diffusers/src/diffusers/models/auto_model.py", line 206, in from_pretrained
    return model_cls.from_pretrained(pretrained_model_or_path, **kwargs)
  File "/home/sayak/.pyenv/versions/higgs/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
  File "/home/sayak/diffusers/src/diffusers/models/modeling_utils.py", line 1248, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/sayak/diffusers/src/diffusers/models/modeling_utils.py", line 1507, in _load_pretrained_model
    offload_index, state_dict_index = load_model_dict_into_meta(
  File "/home/sayak/diffusers/src/diffusers/models/model_loading_utils.py", line 305, in load_model_dict_into_meta
    set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)
  File "/home/sayak/.pyenv/versions/higgs/lib/python3.10/site-packages/accelerate/utils/modeling.py", line 292, in set_module_tensor_to_device
    raise ValueError(
ValueError: Trying to set a tensor of shape torch.Size([3072, 4096]) in "weight" (which has shape torch.Size([768, 4096])), this looks incorrect.

Cc: @SunMarc @MekkCyber

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.