File size: 4,768 Bytes
8da9d20 5cf1cfd 8da9d20 5cf1cfd 8da9d20 5cf1cfd 8da9d20 5cf1cfd 8da9d20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """FLUX.2-klein-4B -> OpenVINO FP16 export + NNCF weight-only INT4 quantization.
Environment (pinned, do not change):
diffusers==0.37.1 transformers==4.57.6 tokenizers==0.22.0
huggingface-hub==0.35.1 optimum-intel==2.2.0 optimum==2.3.0
openvino==2026.4.0 nncf==3.4.0 torch pillow psutil
FP16 export CLI equivalent (reference, same as task example pattern):
optimum-cli export openvino -m black-forest-labs/FLUX.2-klein-4B \
--task text-to-image --library diffusers --weight-format fp16 \
./flux-2-klein-4b-ov-fp16
# same syntax, only -m and output dir differ.
INT4 quantization:
OVQuantizer + OVPipelineQuantizationConfig (pipeline variant of OVConfig)
transformer + text_encoder: OVWeightQuantizationConfig(
bits=4, sym=False, group_size=128, group_size_fallback="adjust", ratio=1.0)
rest: default INT8 weight-only (OVWeightQuantizationConfig() defaults bits=8)
passed as ov_config=... to quantizer.quantize()
"""
import os
import time
import json
from pathlib import Path
FP16_DIR = "./flux-2-klein-4b-ov-fp16"
INT4_DIR = "./flux-2-klein-4b-ov-int4"
MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
def export_fp16():
from optimum.intel import OVDiffusionPipeline
if os.path.exists(FP16_DIR) and os.listdir(FP16_DIR):
print(f"FP16 model already exists at {FP16_DIR}, skipping export...")
print("CLI equivalent: optimum-cli export openvino -m "
f"{MODEL_ID} --task text-to-image --library diffusers "
f"--weight-format fp16 {FP16_DIR}")
return FP16_DIR, 0.0
print(f"Exporting {MODEL_ID} to FP16 OpenVINO...")
print("CLI equivalent: optimum-cli export openvino -m "
f"{MODEL_ID} --task text-to-image --library diffusers "
f"--weight-format fp16 {FP16_DIR}")
start = time.time()
pipeline = OVDiffusionPipeline.from_pretrained(
MODEL_ID,
export=True,
compile=False,
weight_format="fp16",
token=os.environ.get("HF_TOKEN"),
)
pipeline.save_pretrained(FP16_DIR)
export_time = time.time() - start
print(f"FP16 export completed in {export_time:.2f}s")
return FP16_DIR, export_time
def quantize_int4(fp16_dir):
from optimum.intel import OVDiffusionPipeline
from optimum.intel.openvino import (
OVQuantizer,
OVConfig,
OVPipelineQuantizationConfig,
OVWeightQuantizationConfig,
)
print(f"Quantizing {fp16_dir} to INT4...")
start = time.time()
pipeline = OVDiffusionPipeline.from_pretrained(fp16_dir, compile=False)
quantizer = OVQuantizer.from_pretrained(pipeline)
int4_cfg = OVWeightQuantizationConfig(
bits=4,
sym=False,
group_size=128,
group_size_fallback="adjust",
ratio=1.0,
)
int8_default = OVWeightQuantizationConfig() # bits=8 default INT8
# OVPipelineQuantizationConfig is the pipeline variant of OVConfig;
# passed as ov_config=OVConfig(quantization_config=...) to quantize().
# transformer/text_encoder INT4, everything else defaults to INT8.
ov_config = OVConfig(
quantization_config=OVPipelineQuantizationConfig(
quantization_configs={
"transformer": int4_cfg,
"text_encoder": int4_cfg,
},
default_config=int8_default,
)
)
quantizer.quantize(
save_directory=INT4_DIR,
ov_config=ov_config,
)
quantize_time = time.time() - start
print(f"INT4 quantization completed in {quantize_time:.2f}s")
return INT4_DIR, quantize_time
def main():
print("=" * 60)
print("FLUX.2-klein-4B OpenVINO INT4 Quantization")
print("=" * 60)
os.makedirs(OUTPUT_DIR, exist_ok=True)
fp16_dir, export_time = export_fp16()
int4_dir, quantize_time = quantize_int4(fp16_dir)
data = {
"model_id": MODEL_ID,
"fp16_export_dir": fp16_dir,
"int4_dir": int4_dir,
"export_time_seconds": export_time,
"quantize_time_seconds": quantize_time,
"total_time_seconds": export_time + quantize_time,
"quant_config": {
"transformer": "OVWeightQuantizationConfig(bits=4,sym=False,group_size=128,group_size_fallback=adjust,ratio=1.0)",
"text_encoder": "OVWeightQuantizationConfig(bits=4,sym=False,group_size=128,group_size_fallback=adjust,ratio=1.0)",
"others": "default INT8 (OVWeightQuantizationConfig bits=8)",
"ov_config_class": "OVPipelineQuantizationConfig (pipeline variant of OVConfig)",
},
}
with open(os.path.join(OUTPUT_DIR, "benchmark_quantization.json"), "w") as f:
json.dump(data, f, indent=2)
print(f"\nDone. FP16: {fp16_dir} INT4: {int4_dir}")
if __name__ == "__main__":
main()
|