Views
No views yet
black-forest-labs/FLUX.1-dev quantized the Transformer model to INT4 and the T5 Text Encoder to INT8 using Optimum Quanto.pip install diffusers optimum-quanto1import json
2import torch
3import diffusers
4import transformers
5from optimum.quanto import requantize
6from safetensors.torch import load_file
7from huggingface_hub import hf_hub_download
8
9
10def load_quanto_transformer(repo_path):
11 with open(hf_hub_download(repo_path, "transformer/quantization_map.json"), "r") as f:
12 quantization_map = json.load(f)
13 with torch.device("meta"):
14 transformer = diffusers.FluxTransformer2DModel.from_config(hf_hub_download(repo_path, "transformer/config.json")).to(torch.bfloat16)
15 state_dict = load_file(hf_hub_download(repo_path, "transformer/diffusion_pytorch_model.safetensors"))
16 requantize(transformer, state_dict, quantization_map, device=torch.device("cuda"))
17 return transformer
18
19
20def load_quanto_text_encoder_2(repo_path):
21 with open(hf_hub_download(repo_path, "text_encoder_2/quantization_map.json"), "r") as f:
22 quantization_map = json.load(f)
23 with open(hf_hub_download(repo_path, "text_encoder_2/config.json")) as f:
24 t5_config = transformers.T5Config(**json.load(f))
25 with torch.device("meta"):
26 text_encoder_2 = transformers.T5EncoderModel(t5_config).to(torch.bfloat16)
27 state_dict = load_file(hf_hub_download(repo_path, "text_encoder_2/model.safetensors"))
28 requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cuda"))
29 return text_encoder_2
30
31
32pipe = diffusers.AutoPipelineForText2Image.from_pretrained("Disty0/FLUX.1-dev-qint4_tf-qint8_te", transformer=None, text_encoder_2=None, torch_dtype=torch.bfloat16)
33pipe.transformer = load_quanto_transformer("Disty0/FLUX.1-dev-qint4_tf-qint8_te")
34pipe.text_encoder_2 = load_quanto_text_encoder_2("Disty0/FLUX.1-dev-qint4_tf-qint8_te")
35pipe = pipe.to("cuda", dtype=torch.bfloat16)
36
37
38prompt = "A cat holding a sign that says hello world"
39image = pipe(
40 prompt,
41 height=1024,
42 width=1024,
43 guidance_scale=3.5,
44 num_inference_steps=50,
45 max_sequence_length=512,
46 generator=torch.Generator("cpu").manual_seed(0)
47).images[0]
48image.save("flux-dev.png")