Try load gguf for FluxKontextPipeline but invalid input size

[dangph-alala]

Describe the bug

I try the gguf quantize method like FluxPipeline or FluxFillPipeline but fail, I use https://huggingface.co/bullerwins/FLUX.1-Kontext-dev-GGUF/blob/main/flux1-kontext-dev-Q4_K_M.gguf. I also tried resize image to width = (width // 64) * 64, height = (height // 64) * 64

Reproduction

transformer = FluxTransformer2DModel.from_single_file(
                                     "https://huggingface.co/bullerwins/FLUX.1-Kontext-dev-GGUF/blob/main/flux1-kontext-dev-Q4_K_M.gguf",
                                      quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),`
                                      torch_dtype=torch.bfloat16)
pipeline    = FluxKontextPipeline.from_pretrained(
                                "black-forest-labs/FLUX.1-Kontext-dev",
                                transformer=transformer,
                                torch_dtype=torch.bfloat16).to("cuda")

images = pipeline(
    prompt="restored photo",
    image=init_image,
    guidance_scale=2.5,
    num_inference_steps=28,
).images[0]

Logs

File "/root/miniconda3/envs/Flux_test/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_kontext.py", line 687, in prepare_latents
    image_latents = self._pack_latents(
  File "/root/miniconda3/envs/Flux_test/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_kontext.py", line 574, in _pack_latents
    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
RuntimeError: shape '[1, 32, 64, 2, 64, 2]' is invalid for input of size 262144

System Info

• Diffusers version: 0.35.0.dev0
• Platform: Linux-5.15.0-43-generic-x86_64-with-glibc2.35
• Python version: 3.10.18
• PyTorch version (GPU?): 2.6.0+cu124 (True)
• Huggingface_hub version: 0.33.1
• Transformers version: 4.53.0
• Accelerate version: 1.8.1
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.5.3
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 3090, 24576 MiB

Who can help?

No response

[DN6]

@dangph-alala What size image are you using? Could you try resizing so that height/width are divisible by 16?

[dangph-alala]

@dangph-alala What size image are you using? Could you try resizing so that height/width are divisible by 16?
@DN6 yeah, I tried keep ratio of image, and resized the image shape to be divisible by 64 so it is divisible by 16.
width = (width // 64) * 64, height = (height // 64) * 64

[asomoza]

@DN6 this code can test this:

import torch

from diffusers import FluxKontextPipeline, FluxTransformer2DModel, GGUFQuantizationConfig
from diffusers.utils import load_image


image = load_image("https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/resources/dog_plushie.png")


transformer = FluxTransformer2DModel.from_single_file(
    "https://huggingface.co/bullerwins/FLUX.1-Kontext-dev-GGUF/blob/main/flux1-kontext-dev-Q4_K_M.gguf",
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
)
pipeline = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
).to("cuda")

image = pipeline(
    prompt="change the plushie to a capybara plushie while maintaining the same style, position and background",
    image=image,
    guidance_scale=2.5,
    num_inference_steps=28,
).images[0]

image.save("dog_plushie_to_capybara_plushie.png")

without GGUF it works and with it doesn't, it seems that the num_channels_latents is 32 with the GGUF model instead of 16 with the non quantized transformer model.

[dangph-alala]

@asomoza yeah, I think so, too. And I haven’t figured out how to fix the diffusers to handle yet. Do you happen to have a suggestion for this issue? Idk why it still work in Comfy UI :(((

[asomoza]

yeah, I have an idea but it's this is something that probably @DN6 will fix a lot faster than me, since I will need to start looking for a way to differentiate the kontext model.

The problem lies here, since the shapes match, the function returns that is a depth model instead of a normal flux one.

The solution here is to come with a way to differentiate kontext from the depth model.

edit: Now that I look more into this, I can load the original model with from_single_file and it works, so this seems to be something wrong with the GGUF model, that shape in the original is 64 instead of 128 from the GGUF model.

[dangph-alala]

yeah, I have an idea but it's this is something that probably @DN6 will fix a lot faster than me, since I will need to start looking for a way to differentiate the kontext model.

The problem lies here, since the shapes match, the function returns that is a depth model instead of a normal flux one.

The solution here is to come with a way to differentiate kontext from the depth model.

Thank you about that!!!!!!

edit: Now that I look more into this, I can load the original model with from_single_file and it works, so this seems to be something wrong with the GGUF model, that shape in the original is 64 instead of 128 from the GGUF model.

Just to confirm - It means you used from_single_file without GGUF or still used from_single_file for both GGUF and original model for it working?

[asomoza]

I used from_single_file with the original unquantized model and it works, using the same from_single_file with the GGUF model doesn't work. What I mean with this is this is not an issue with the from_single_file but with the GGUF version of the model.

Sadly I don't use that much Flux and from_single file to have more knowledge as to what is causing this issue.