| | --- |
| | language: |
| | - en |
| | license: other |
| | license_name: flux-1-dev-non-commercial-license |
| | license_link: LICENSE.md |
| | tags: |
| | - text-to-image |
| | - image-generation |
| | - flux |
| | --- |
| | |
| | `black-forest-labs/FLUX.1-dev` quantized the Transformer and the T5 Text Encoder to INT4 using Optimum Quanto with FP16 calculations. |
| |
|
| | ```shell |
| | pip install diffusers optimum-quanto |
| | ``` |
| |
|
| | ```python |
| | import json |
| | import torch |
| | import diffusers |
| | import transformers |
| | from optimum.quanto import requantize |
| | from safetensors.torch import load_file |
| | from huggingface_hub import hf_hub_download |
| | |
| | |
| | def load_quanto_transformer(repo_path): |
| | with open(hf_hub_download(repo_path, "transformer/quantization_map.json"), "r") as f: |
| | quantization_map = json.load(f) |
| | with torch.device("meta"): |
| | transformer = diffusers.FluxTransformer2DModel.from_config(hf_hub_download(repo_path, "transformer/config.json")).to(torch.float16) |
| | state_dict = load_file(hf_hub_download(repo_path, "transformer/diffusion_pytorch_model.safetensors")) |
| | requantize(transformer, state_dict, quantization_map, device=torch.device("cuda")) |
| | return transformer |
| | |
| | |
| | def load_quanto_text_encoder_2(repo_path): |
| | with open(hf_hub_download(repo_path, "text_encoder_2/quantization_map.json"), "r") as f: |
| | quantization_map = json.load(f) |
| | with open(hf_hub_download(repo_path, "text_encoder_2/config.json")) as f: |
| | t5_config = transformers.T5Config(**json.load(f)) |
| | with torch.device("meta"): |
| | text_encoder_2 = transformers.T5EncoderModel(t5_config).to(torch.float16) |
| | state_dict = load_file(hf_hub_download(repo_path, "text_encoder_2/model.safetensors")) |
| | requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cuda")) |
| | return text_encoder_2 |
| | |
| | |
| | pipe = diffusers.AutoPipelineForText2Image.from_pretrained("Disty0/FLUX.1-dev-qint4_tf-qint8_te", transformer=None, text_encoder_2=None, torch_dtype=torch.float16) |
| | pipe.transformer = load_quanto_transformer("Disty0/FLUX.1-dev-qint4_tf-qint8_te") |
| | pipe.text_encoder_2 = load_quanto_text_encoder_2("Disty0/FLUX.1-dev-qint4_tf-qint8_te") |
| | pipe = pipe.to("cuda", dtype=torch.float16) |
| | |
| | |
| | prompt = "A cat holding a sign that says hello world" |
| | image = pipe( |
| | prompt, |
| | height=1024, |
| | width=1024, |
| | guidance_scale=3.5, |
| | num_inference_steps=50, |
| | max_sequence_length=512, |
| | generator=torch.Generator("cpu").manual_seed(0) |
| | ).images[0] |
| | image.save("flux-dev.png") |
| | ``` |
| |
|