Add quantize_int4_flux.py

quantize_int4_flux.py

@@ -0,0 +1,88 @@
+"""NNCF weight-only INT4 quantization of the OpenVINO FP16 FLUX pipeline.
+
+transformer + text_encoder -> weight-only INT4 (asymmetric, group 128)
+all remaining components   -> default weight-only INT8
+
+Usage:
+    python quantize_int4_flux.py \
+        --model_path /home/user/app/flux-schnell-ov-fp16 \
+        --output_path /home/user/app/flux-schnell-ov-int4
+"""
+
+import argparse
+import shutil
+import time
+from pathlib import Path
+
+from optimum.intel import OVConfig, OVQuantizer
+from optimum.intel.openvino import (
+    OVPipelineQuantizationConfig,
+    OVWeightQuantizationConfig,
+)
+
+# For FLUX we need to import the specific pipeline class
+try:
+    from optimum.intel import OVFluxPipeline
+except ImportError:
+    # Fallback to generic diffusion pipeline
+    from optimum.intel import OVDiffusionPipeline as OVFluxPipeline
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="/home/user/app/flux-schnell-ov-fp16",
+        help="OpenVINO FP16 pipeline exported with optimum-cli",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="/home/user/app/flux-schnell-ov-int4",
+        help="Destination folder for the INT4 pipeline",
+    )
+    args = parser.parse_args()
+
+    model_path = Path(args.model_path)
+    output_path = Path(args.output_path)
+    if output_path.exists():
+        shutil.rmtree(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    int4 = dict(
+        bits=4,
+        sym=False,
+        group_size=128,
+        group_size_fallback="adjust",
+        ratio=1.0,
+    )
+    quantization_configs = {
+        # FLUX uses "transformer" instead of "unet"
+        "transformer": OVWeightQuantizationConfig(**int4),
+        "text_encoder": OVWeightQuantizationConfig(**int4),
+        # Also quantize the second text encoder (T5)
+        "text_encoder_2": OVWeightQuantizationConfig(**int4),
+    }
+    default_config = OVWeightQuantizationConfig(bits=8)
+
+    quantization_config = OVPipelineQuantizationConfig(
+        quantization_configs=quantization_configs,
+        default_config=default_config,
+    )
+    ov_config = OVConfig(quantization_config=quantization_config)
+
+    print(f"loading FP16 pipeline from {model_path} ...", flush=True)
+    t0 = time.perf_counter()
+    model = OVFluxPipeline.from_pretrained(str(model_path), device="CPU")
+    print(f"loaded in {time.perf_counter() - t0:.1f}s", flush=True)
+
+    quantizer = OVQuantizer(model=model)
+    t0 = time.perf_counter()
+    quantizer.quantize(ov_config=ov_config, save_directory=str(output_path))
+    print(f"quantization took {time.perf_counter() - t0:.1f}s", flush=True)
+    print(f"INT4 pipeline saved to {output_path}")
+
+
+if __name__ == "__main__":
+    main()