OPQ INT4 Z-Image: 最优幂次量化扩散模型推理

使用 OPQ (Optimal Power Quantization) 将 Z-Image 6.02B 参数扩散模型量化到 INT4，GPU 峰值显存从 20.7 GB 降至 9.3 GB，图像质量几乎无损。

项目地址: www.modelscope.cn/dfytensor/OPQ-INT4-Z-Image

一键使用

# 1. 克隆项目 (含量化权重)
git clone https://www.modelscope.cn/dfytensor/OPQ-INT4-Z-Image.git
cd OPQ-INT4-Z-Image

# 2. 安装依赖
pip install -r requirements.txt

# 3. 运行推理 (自动加载量化权重 + 生成图像)
python opq_infer.py --load-saved

# 4. 查看结果
# → opq_output.png (768x1024, peak GPU 9.3 GB, ~57s)

显存要求: 仅需 12 GB GPU + 12 GB 内存即可运行。 速度取决于机器配置。

效果展示

bf16 原始 vs OPQ INT4

左：bf16 原始 (20.7 GB GPU) | 右：OPQ INT4 α=0.73 (9.3 GB GPU)

数值对比

量化前后的权重质量

对比图见 opq_full_compare.png。

OPQ 原理

传统均匀量化：q = round(w / Δ) — 均匀间距，不适应权重分布

OPQ 幂次量化：

归一化:  u = |w| / max(|w|)           ∈ [0, 1]
幂次变换: v = u^α                      (α=0.73 放大小值)
量化:    q = round(v × 15)             (INT4, 16级)
反变换:  w' = sign(w) × (q/15)^(1/α) × max(|w|)

α=0.73 是 Z-Image 权重分布的最优值，比 GPT-2 的 α=0.45 更大（扩散模型权重更集中）。

核心优势：

• SNR +7.2 dB（等效免费 +1 bit 精度）
• KLD 降低 691 倍（权重分布几乎无损）
• 零额外存储开销（同样 INT4 格式）

文件结构

OPQ_ZImage_INT4/
├── README.md                        ← 本文件
├── opq_quantize.py                  ← OPQ 量化模块 (OPQLinear + replace_with_opq)
├── opq_save.py                      ← 量化权重保存脚本
├── opq_infer.py                     ← 推理脚本 (量化 + 生成图像)
├── opq_full_compare.png             ← bf16 vs OPQ INT4 对比图
├── OPQ_Application_Report.md        ← 完整评估报告
├── requirements.txt                 ← 依赖列表
└── quantized_weights/               ← 已量化的权重文件 (可直接加载)
    ├── transformer_opq_int4.pt      ← Transformer 权重 (3.67 GB)
    ├── text_encoder_opq_int4.pt     ← T5 编码器权重 (2.17 GB)
    ├── vae_opq_int4.pt              ← VAE 权重 (~0 GB)
    └── config.pt                    ← 量化配置

环境要求

硬件（最低配置）

• GPU: 12 GB 显存 (如 RTX 3060 12GB / RTX 4070 / RTX 3090 / RTX 4090 等)
• 内存: 12 GB RAM
• 存储: 10 GB (克隆本项目含量化权重 ~6 GB + 模型缓存 ~4 GB)

推理速度参考

注意: 即使显存只有 12 GB, OPQ INT4 的 peak 仅 9.3 GB, 有足够余量。

软件

# 创建虚拟环境
python -m venv .venv

# Windows
.venv\Scripts\activate
# Linux
source .venv/bin/activate

# 安装依赖
pip install -r requirements.txt

requirements.txt 内容：

torch>=2.4.0
diffusers>=0.30.0
transformers>=4.45.0
accelerate>=0.34.0
Pillow>=9.0.0
numpy>=1.24.0
safetensors>=0.4.0

模型下载

Z-Image 原始模型需提前下载（推理脚本需要先加载原始模型结构，再覆盖为量化权重）：

# 方式1: modelscope
pip install modelscope
modelscope download --model Tongyi-MAI/Z-Image --local_dir /path/to/zimage

# 方式2: huggingface
huggingface-cli download Tongyi-MAI/Z-Image --local-dir /path/to/zimage

下载后修改 opq_infer.py 中的 MODEL_PATH：

MODEL_PATH = "/path/to/zimage"  # 改成你的路径

本项目已含量化权重 (quantized_weights/, ~6 GB), 无需单独下载量化模型。 原始模型仅用于加载模型结构, 推理时使用量化权重。

快速开始

方式 A：从量化权重加载（推荐，最快）

quantized_weights/ 已包含预计算的 OPQ INT4 权重，无需重新量化：

cd OPQ_ZImage_INT4
python opq_infer.py --load-saved

输出：

[1] Loading Z-Image...
      Original GPU: 20.7 GB
[2] Loading saved quantized weights...
      transformer: 276 layers, 0 missing buffers
      text_encoder: 252 layers, 0 missing buffers
      vae: 8 layers, 0 missing buffers
      Total: 20.7 -> 7.3 GB (saved 13.4 GB)
[3] Generating image...
  Result:
    Image:       opq_output.png
    Size:        768x1024
    Mean:        118
    Std:         86.3
    CorrH:       0.9956
    Peak GPU:    9.3 GB
    Time:        57s

方式 B：在线量化（用于换模型或调参数）

# 1. 量化原始权重并保存（首次运行，约 10 秒）
python opq_save.py

# 2. 之后用 --load-saved 直接加载
python opq_infer.py --load-saved

输出：

  Quantizing transformer...
    276 layers quantized (2s), saving...
    Saved: quantized_weights/transformer_opq_int4.pt (3.94 GB, 5s)
  Quantizing text_encoder...
    252 layers quantized (1s), saving...
    Saved: quantized_weights/text_encoder_opq_int4.pt (2.33 GB, 2s)
  Quantizing vae...
    8 layers quantized (0s), saving...
    Saved: quantized_weights/vae_opq_int4.pt (0.00 GB, 0s)
  Done! Files in quantized_weights/

方式 C：在线量化 + 推理（一步到位，不保存）

python opq_infer.py

推理参数配置

编辑 opq_infer.py 顶部的参数：

# 模型路径（必须修改为你的本地路径）
MODEL_PATH = r"F:\zimage_cache\models\Tongyi-MAI--Z-Image\snapshots\master"

# 量化参数
OPQ_BITS = 4       # 量化位数: 3, 4, 5, 6
OPQ_ALPHA = 0.73   # 幂次参数: 0.3-1.0 (0.73=最优)
OPQ_GS = 128       # 分组大小: 32, 64, 128, 256

# 生成参数
PROMPT = "你的提示词"
NEGATIVE = "负面提示词"
HEIGHT = 1024
WIDTH = 768
STEPS = 30
SEED = 42

从量化权重加载的流程说明

--load-saved 模式的工作原理：

1. 加载原始 Z-Image bf16 模型到 GPU        (20.7 GB)
2. 对每个组件:
   a. 运行 replace_with_opq() 创建 OPQLinear 结构
   b. 用 quantized_weights/*.pt 覆盖 packed/sign_packed/mx buffers
3. 释放原始权重，仅保留量化数据             (7.3 GB)
4. 正常推理生成图像                         (peak 9.3 GB)

注意：即使使用 --load-saved，仍需下载原始 Z-Image 模型。因为 diffusers 需要先加载模型结构，然后用量化权重覆盖。

不同位宽的效果

权重文件说明

quantized_weights/ 目录包含预计算好的 OPQ INT4 量化权重：

合计 5.84 GB（原始 bf16 20.5 GB，压缩 3.5×）

量化哪些组件

Z-Image pipeline 包含三个模型组件：

默认量化全部组件。如需只量化 Transformer（保留 T5/VAE 为 bf16），修改 opq_infer.py：

# 只量化 transformer
components = ['transformer']

在你自己的代码中使用

from opq_quantize import replace_with_opq

# 加载任意 diffusers 模型
pipe = DiffusionPipeline.from_pretrained("your-model", torch_dtype=torch.bfloat16)
pipe.to("cuda")

# OPQ 量化 (一行搞定)
replace_with_opq(pipe.transformer, bits=4, alpha=0.73)
replace_with_opq(pipe.text_encoder, bits=4, alpha=0.73)

# 正常推理
image = pipe(prompt="hello").images[0]

技术细节

存储格式

每个权重存储 3 部分：

1. Magnitude (4-bit): round((|w|/max)^0.73 × 15) → 8 个打包进 1 个 uint32
2. Sign (1-bit): 正/负 → 32 个打包进 1 个 uint32
3. Group max (float16): 每 128 个权重共享一个

有效 BPW = (4 + 0.125 + 0.125) = 4.25 BPW

推理解码

# 从存储格式重建权重:
qn = q / 15                    # [0, 1]
inv = qn ** (1.0 / 0.73)       # 逆幂次
w = sign * inv * group_max     # 正负号 × 幅度 × 组最大值

为什么 α=0.73？

• Z-Image 权重 per-group 归一化后分布较集中（σ/μ ≈ 0.39）
• 比均匀分布需要更弱的非线性（α 接近 1.0）
• GPT-2 的 α=0.45 更小因为 LLM 权重分布尾部更长
• α=0.73 是 Z-Image 上的平坦最优区 [0.65, 0.80] 的中心

引用

@misc{opq2026,
  title={Optimal Power Quantization for Diffusion Model Weight Compression},
  year={2026},
}

License

仅供学术研究使用。