hf-internal-testing/flux.1-dev-int8-pkg

Running Flux.1-dev under 12GBs

This repository contains the mixed int8 params for the T5 and transformer of Flux.1-Dev.

This is how the checkpoints were obtained:

import torch
from diffusers import BitsAndBytesConfig, FluxTransformer2DModel
from transformers import T5EncoderModel, BitsAndBytesConfig
from huggingface_hub import create_repo, upload_folder
import tempfile

model_id = "black-forest-labs/FLUX.1-dev"

config = BitsAndBytesConfig(load_in_8bit=True)
transformer = FluxTransformer2DModel.from_pretrained(
    model_id, subfolder="transformer", quantization_config=config, torch_dtype=torch.float16
)

config = BitsAndBytesConfig(load_in_8bit=True)
t5 = T5EncoderModel.from_pretrained(
    model_id, subfolder="text_encoder_2", quantization_config=config, torch_dtype=torch.float16
)

repo_id = create_repo("sayakpaul/flux.1-dev-int8-pkg", exist_ok=True).repo_id

with tempfile.TemporaryDirectory() as tmpdirname:
    transformer.save_pretrained(tmpdirname)
    upload_folder(repo_id=repo_id, folder_path=tmpdirname, path_in_repo="transformer")

with tempfile.TemporaryDirectory() as tmpdirname:
    t5.save_pretrained(tmpdirname)
    upload_folder(repo_id=repo_id, folder_path=tmpdirname, path_in_repo="text_encoder_2")

Respective diffusers PR: https://github.com/huggingface/diffusers/pull/9213/.

The checkpoints of this repository were optimized to run on a T4 notebook. More specifically, the compute datatype of the quantized checkpoints was kept to FP16. In practice, if you have a GPU card that supports BF16, you should change the compute datatype to BF16 (bnb_4bit_compute_dtype).