Blame for quantize_int4.py

#!/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())