Add quantize_int4.py

quantize_int4.py

@@ -0,0 +1,49 @@
+"""Quantize Z-Image-Turbo OV FP16 -> INT4 using optimum (NNCF weight-only)."""
+from pathlib import Path
+from optimum.intel import OVDiffusionPipeline, OVQuantizer
+from optimum.intel.openvino.configuration import (
+    OVConfig,
+    OVWeightQuantizationConfig,
+    OVPipelineQuantizationConfig,
+)
+
+fp16_dir = "/home/user/app/z-image-turbo-ov-fp16"
+int4_dir = "/home/user/app/z-image-turbo-ov-int4"
+
+# INT4 config for transformer + text_encoder (main weights)
+# group_size=128, sym=False is NNCF default sweet spot for quality/size
+int4_config = OVWeightQuantizationConfig(
+    bits=4,
+    sym=False,
+    group_size=128,
+    group_size_fallback="adjust",
+    ratio=1.0,
+)
+
+# Keep VAE in FP16 (don't quantize) by using default FP32? Actually set bits=8 for others,
+# but we will explicitly only quantize transformer+text_encoder and copy rest as-is.
+# Using OVPipelineQuantizationConfig:
+pipeline_config = OVPipelineQuantizationConfig(
+    quantization_configs={
+        "transformer": int4_config,
+        "text_encoder": int4_config,
+    },
+    # default: 8-bit weight-only for any other submodel that gets quantized,
+    # but we skip vae by not including it? OVQuantizer quantizes only listed?
+    # To be safe, set default 8-bit; vae will be copied if not quantized?
+    default_config=OVWeightQuantizationConfig(bits=8, sym=True),
+)
+
+print(f"Loading FP16 OV pipeline from {fp16_dir} ...")
+pipe = OVDiffusionPipeline.from_pretrained(fp16_dir)
+print(f"Components: {list(pipe.components.keys())}")
+print(f"OV submodels: {list(pipe.ov_submodels) if hasattr(pipe, 'ov_submodels') else 'n/a'}")
+
+quantizer = OVQuantizer.from_pretrained(pipe)
+print("Quantizing transformer+text_encoder to INT4 (NNCF, optimum)...")
+ov_config = OVConfig(quantization_config=pipeline_config)
+quantizer.quantize(
+    save_directory=int4_dir,
+    ov_config=ov_config,
+)
+print(f"Saved INT4 model to {int4_dir}")