File size: 1,860 Bytes
d3fc334 eab6841 d3fc334 | 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 | """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 = "./z-image-turbo-ov-fp16"
int4_dir = "./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}")
|