| """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() |
|
|
| |
| |
| |
| 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() |
|
|
