HighCWu/FLUX.1-Kontext-dev-bnb-hqq-4bit

Flux-Kontext-dev model outputs using BnB&Hqq 4bit quantization

Original Input
BnB 4-bit (DiT) & Hqq 4-bit (T5)

Usage with Diffusers

To use this quantized FLUX.1 [dev] checkpoint, you need to install the 🧨 diffusers, transformers, bitsandbytes and hqq library:

pip install git+https://github.com/huggingface/diffusers.git@main # add support for `FluxKontextPipeline`
pip install transformers>=4.53.1 # add support for hqq quantized model in diffusers pipeline
pip install -U bitsandbytes
pip install -U hqq

After installing the required library, you can run the following script:

import torch
from diffusers import FluxKontextPipeline
from diffusers.utils import load_image

pipe = FluxKontextPipeline.from_pretrained("HighCWu/FLUX.1-Kontext-dev-bnb-hqq-4bit", torch_dtype=torch.bfloat16)
pipe.to("cuda")

input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")

image = pipe(
  image=input_image,
  prompt="Add a hat to the cat",
  guidance_scale=2.5,
  num_inference_steps=28,
  generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save(f"kontext.1-dev.png")

How to generate this quantized checkpoint ?

This checkpoint was created with the following script using "black-forest-labs/FLUX.1-Kontext-dev" checkpoint:

import torch

assert torch.cuda.is_available() # force initialization of cuda

from diffusers import FluxKontextPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import HqqConfig as TransformersHqqConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16),
        "text_encoder_2": TransformersHqqConfig(nbits=4, group_size=64),
    }
)

pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16
)

pipe.save_pretrained("FLUX.1-Kontext-dev-bnb-hqq-4bit")