#!/usr/bin/env python3 """Convert stabilityai/stable-diffusion-3.5-medium to OpenVINO IR (FP16), then compress weights to INT4 with NNCF. Method mirrors the conversion documented in https://huggingface.co/HelloSun/FLUX.2-klein-4B-OpenVINO-INT4 INT4 : transformer + text_encoder_3 (T5-XXL) OVWeightQuantizationConfig(bits=4, sym=False, group_size=128, group_size_fallback="adjust", ratio=1.0) rest : default INT8 OVWeightQuantizationConfig(bits=8) via OVQuantizer + ov_config=OVConfig( quantization_config=OVPipelineQuantizationConfig(...)) SD3.5-medium has three text encoders. Only text_encoder_3 (T5-XXL, 4.7B params, ~9.5 GB FP16) is pushed to INT4; the two small CLIP encoders (CLIP-L / CLIP-G) stay at INT8 because quantizing them costs little size and risks prompt fidelity. Inference parameters are NOT changed by this script: they stay at the values in the original model card (1024x1024, 40 steps, guidance 4.5). Usage ----- export HF_TOKEN=... # SD3.5-medium is gated python quantize_int4.py # or split the two phases python quantize_int4.py --export-only python quantize_int4.py --quantize-only """ from __future__ import annotations import argparse import json import os import subprocess import sys import tempfile import time from pathlib import Path MODEL_ID = "stabilityai/stable-diffusion-3.5-medium" REPO_ID = "HelloSun/SD3.5-medium-OpenVINO-INT4" # INT4 keeps the two large components. group_size_fallback="adjust" downshifts to # the largest compatible group when a tensor's shape is not divisible by 128, # which is common for the small projection matrices inside MMDiT blocks. INT4_COMPONENTS = ("transformer", "text_encoder_3") INT4_KWARGS = { "bits": 4, "sym": False, "group_size": 128, "group_size_fallback": "adjust", "ratio": 1.0, } def log(msg: str) -> None: print(f"[quantize_int4] {msg}", flush=True) def dir_size_gb(path: Path) -> float: total = 0 for root, _dirs, files in os.walk(path): for name in files: fp = Path(root) / name if fp.is_file() and not fp.is_symlink(): total += fp.stat().st_size return total / 1024**3 def component_sizes(path: Path) -> dict[str, float]: """Per-subfolder size in GB, for the components we actually quantize.""" out: dict[str, float] = {} if not path.is_dir(): return out for child in sorted(p for p in path.iterdir() if p.is_dir()): out[child.name] = dir_size_gb(child) return out def require_token() -> None: """SD3.5-medium is gated (license must be accepted on the Hub).""" if os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"): return log("No HF_TOKEN in environment. If the repo is not already accepted for") log("your account, run: hf auth login") def export_fp16(out_dir: Path) -> None: """FP16 export via optimum-cli, matching the reference card's export step.""" if (out_dir / "openvino_config.json").is_file(): log(f"FP16 export already present at {out_dir}, skipping.") return cmd = [ "optimum-cli", "export", "openvino", "--model", MODEL_ID, "--task", "text-to-image", "--library", "diffusers", "--weight-format", "fp16", str(out_dir), ] log("running: " + " ".join(cmd)) log("(downloads ~16 GB of FP16 weights; this is the long part)") subprocess.run(cmd, check=True) def quantize_int4(fp16_dir: Path, out_dir: Path) -> None: import shutil # immediate_save=True overwrites files in save_directory in place, so the # FP16 IR must already be there. Seed it automatically on first run. out_exists = out_dir.exists() and any(out_dir.iterdir()) if not out_exists: if (fp16_dir / "model_index.json").is_file(): log(f"seeding {out_dir} from {fp16_dir} (required for immediate_save=True)") shutil.copytree(fp16_dir, out_dir, dirs_exist_ok=True) # TemporaryDirectory (used internally for each sub-model) must live on the # same filesystem as out_dir, otherwise rename() fails with EXDEV # (Invalid cross-device link). E.g.: export TMPDIR=/data-nvme/work/tmp try: if out_dir.exists() and Path(tempfile.gettempdir()).stat().st_dev != Path(out_dir).stat().st_dev: log("WARNING: TMPDIR is on a different filesystem than out_dir; " "set TMPDIR inside the same filesystem to avoid EXDEV rename errors.") except OSError: pass from optimum.intel import ( OVConfig, OVDiffusionPipeline, OVPipelineQuantizationConfig, OVQuantizer, OVWeightQuantizationConfig, ) log(f"loading FP16 OpenVINO pipeline from {fp16_dir}") model = OVDiffusionPipeline.from_pretrained(str(fp16_dir)) available = sorted(c for c in ("transformer", "text_encoder_3") if getattr(model, c, None) is not None) if available != list(INT4_COMPONENTS): log(f"WARNING: expected components {list(INT4_COMPONENTS)}, pipeline exposes {available}") def int4_cfg(): return OVWeightQuantizationConfig(**INT4_KWARGS) quantization_config = OVPipelineQuantizationConfig( quantization_configs={name: int4_cfg() for name in INT4_COMPONENTS}, # Everything else (CLIP encoders, VAE decoder/encoder) at INT8. default_config=OVWeightQuantizationConfig(bits=8), ) quantizer = OVQuantizer.from_pretrained(model) log("quantizing -> INT4 (transformer, text_encoder_3) / INT8 (rest)") quantizer.quantize( ov_config=OVConfig(quantization_config=quantization_config), save_directory=str(out_dir), # Saves each sub-model as soon as it is compressed, so the FP16 copy of a # component can be released. Cuts peak RSS substantially. The in-memory # pipeline is unusable afterwards, which is fine here. immediate_save=True, ) def main() -> int: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--fp16-dir", default="sd3.5-medium-ov-fp16", type=Path) parser.add_argument("--out-dir", default=REPO_ID.split("/")[-1], type=Path) parser.add_argument("--export-only", action="store_true") parser.add_argument("--quantize-only", action="store_true") args = parser.parse_args() if args.export_only and args.quantize_only: parser.error("--export-only and --quantize-only are mutually exclusive") require_token() report: dict[str, object] = {"model_id": MODEL_ID, "repo_id": REPO_ID} if not args.quantize_only: t0 = time.perf_counter() export_fp16(args.fp16_dir) report["export_seconds"] = round(time.perf_counter() - t0, 2) report["fp16_size_gb"] = round(dir_size_gb(args.fp16_dir), 2) report["fp16_components_gb"] = {k: round(v, 2) for k, v in component_sizes(args.fp16_dir).items()} log(f"FP16 export done: {report['fp16_size_gb']} GB") if args.export_only: (args.out_dir).mkdir(parents=True, exist_ok=True) (args.out_dir / "benchmark_quantization.json").write_text( json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8" ) return 0 if args.out_dir.exists() and args.out_dir.is_dir() and any(args.out_dir.iterdir()): log(f"WARNING: {args.out_dir} is not empty; quantize will overwrite its contents.") t0 = time.perf_counter() quantize_int4(args.fp16_dir, args.out_dir) report["quantize_seconds"] = round(time.perf_counter() - t0, 2) report["int4_size_gb"] = round(dir_size_gb(args.out_dir), 2) report["int4_components_gb"] = {k: round(v, 2) for k, v in component_sizes(args.out_dir).items()} if isinstance(report.get("fp16_size_gb"), float) and report["fp16_size_gb"]: report["compression_ratio"] = round(report["fp16_size_gb"] / report["int4_size_gb"], 2) report["int4_config"] = {"components": list(INT4_COMPONENTS), **INT4_KWARGS, "rest": "int8"} args.out_dir.mkdir(parents=True, exist_ok=True) (args.out_dir / "benchmark_quantization.json").write_text( json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8" ) log(json.dumps(report, indent=2, ensure_ascii=False)) log("done. copy README.md / REPORT.md / *.py into the output dir, then:") log(f" huggingface-cli upload {REPO_ID} {args.out_dir} --repo-type model") return 0 if __name__ == "__main__": sys.exit(main())