diffusers/FLUX.1-dev-torchao-int4

1---2base_model: black-forest-labs/FLUX.1-dev3library_name: diffusers4base_model_relation: quantized5tags:6- quantization7---8 9# Visual comparison of Flux-dev model outputs using BF16 and torchao int4_weight_only quantization10 11<td style="text-align: center;">12  BF16<br>13  <medium-zoom background="rgba(0,0,0,.7)"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/quantization-backends-diffusers/combined_flux-dev_bf16_combined.png" alt="Flux-dev output with BF16: Baroque, Futurist, Noir styles"></medium-zoom>14</td>15<td style="text-align: center;">16  torchao int4_weight_only<br>17  <medium-zoom background="rgba(0,0,0,.7)"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/quantization-backends-diffusers/combined_flux-dev_torchao_4bit_combined.png" alt="torchao int4_weight_only Output"></medium-zoom>18</td>19 20# Usage with Diffusers21 22To use this quantized FLUX.1 [dev] checkpoint, you need to install the 🧨 diffusers and torchao library:23 24```25pip install -U torchao26```27 28For now, we require this specific branch in diffusers library to fix an error when loading the model 29 30```31pip install git+https://github.com/huggingface/diffusers.git@torchao-int4-serialization32```33 34After installing the required library, you can run the following script: 35 36```python37from diffusers import FluxPipeline38 39pipe = FluxPipeline.from_pretrained(40    "diffusers/FLUX.1-dev-torchao-int4",41    torch_dtype=torch.bfloat16,42    use_safetensors=False,43    device_map="balanced"44)45 46prompt = "Baroque style, a lavish palace interior with ornate gilded ceilings, intricate tapestries, and dramatic lighting over a grand staircase."47 48pipe_kwargs = {49    "prompt": prompt,50    "height": 1024,51    "width": 1024,52    "guidance_scale": 3.5,53    "num_inference_steps": 50,54    "max_sequence_length": 512,55}56 57image = pipe(58    **pipe_kwargs, generator=torch.manual_seed(0),59).images[0]60 61image.save("flux.png")62```63 64# How to generate this quantized checkpoint ? 65 66This checkpoint was created with the following script using "black-forest-labs/FLUX.1-dev" checkpoint:67 68```python69 70import torch71from diffusers import FluxPipeline72from diffusers.quantizers import PipelineQuantizationConfig73from diffusers import TorchAoConfig as DiffusersTorchAoConfig74from transformers import TorchAoConfig as TransformersTorchAoConfig75 76pipeline_quant_config = PipelineQuantizationConfig(77    quant_mapping={78        "transformer": DiffusersTorchAoConfig("int4_weight_only"),79        "text_encoder_2": TransformersTorchAoConfig("int4_weight_only"),80    }81)82 83pipe = FluxPipeline.from_pretrained(84    "black-forest-labs/FLUX.1-dev",85    quantization_config=pipeline_quant_config,86    torch_dtype=torch.bfloat16,87    device_map="balanced"88)89 90# safe_serialization set to `False` as we can't save torchao quantized model to safetensors format91pipe.save_pretrained("FLUX.1-dev-torchao-int4", safe_serialization=False)92```