scripts/quantize.py

"""Quantize Qwen-Image-2.1 linear weights to portable SDNQ UINT4."""
import argparse
import gc
import inspect
import json
import shutil
from collections import Counter
from pathlib import Path

import torch

from scripts.integrity import verify_files

REVISION = 'b3179ad355be050328e483a9dfdd9e60cd62adfa'
DIFFUSERS_COMMIT = '80c7ed262aeffbeb43ef13ae04baeb9b84515a69'
TRANSFORMER_SKIP = ['proj_out', 'img_in', 'txt_in', 'time_text_embed', 'modulation', 'norm_out']
TEXT_SKIP = ['lm_head', 'visual']
NOTICE = (
    'Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 '
    'Hangzhou Tongyi Laboratory Technology Co., Ltd. All Rights Reserved.\n\n'
    'Built with Qwen\n'
    'Independent derivative: Image21-INT4, by ixim / iximbox.\n'
    'Modified files: transformer and text_encoder weight shards, shard indexes and '
    'config.json files. Eligible linear weights converted to SDNQ UINT4 with SVD rank 32; '
    'floating-point exceptions are documented in component quantization reports.\n'
    'VAE, scheduler and processor retained from the pinned upstream snapshot.\n')


def sdnq_config(skip, device):
    import os
    os.environ.setdefault('DIFFUSERS_SDNQ_TRANSFORMERS', '1')
    from sdnq import SDNQConfig
    quant_device = torch.device(device)
    requested = dict(weights_dtype='uint4', use_svd=True, svd_rank=32,
                     modules_to_not_convert=list(skip), use_quantized_matmul=False,
                     quantization_device=quant_device, return_device=torch.device('cpu'))
    parameters = inspect.signature(SDNQConfig).parameters
    accepts_keywords = any(item.kind == inspect.Parameter.VAR_KEYWORD for item in parameters.values())
    kwargs = requested if accepts_keywords else {key: value for key, value in requested.items() if key in parameters}
    missing = [key for key in ('weights_dtype', 'use_svd') if key not in kwargs]
    if missing:
        raise RuntimeError(f'SDNQConfig cannot express {missing}')
    return SDNQConfig(**kwargs)


def _progress(component):
    import sdnq.quantizer as quantizer
    original = quantizer.sdnq_quantize_layer
    seen = {'count': 0}

    def wrapped(*args, **kwargs):
        seen['count'] += 1
        name = kwargs.get('param_name', args[3] if len(args) > 3 else '')
        if seen['count'] == 1 or seen['count'] % 10 == 0:
            print(f'  {component}: quantized {seen["count"]} layers, latest {name}', flush=True)
        return original(*args, **kwargs)

    quantizer.sdnq_quantize_layer = wrapped
    return original


def quantize_component(source, target, component, device):
    source, target = Path(source), Path(target)
    out = target / component
    if out.exists() and any(out.iterdir()):
        raise FileExistsError(f'Refusing to overwrite existing component: {out}')
    if component == 'transformer':
        from diffusers import QwenImage21Transformer2DModel
        cls, skip = QwenImage21Transformer2DModel, TRANSFORMER_SKIP
    elif component == 'text_encoder':
        from transformers import Qwen3VLForConditionalGeneration
        cls, skip = Qwen3VLForConditionalGeneration, TEXT_SKIP
    else:
        raise ValueError(component)
    config = sdnq_config(skip, device)
    print(f'Quantizing {component}; math on {device}, saved weights on CPU', flush=True)
    original_quantize = _progress(component)
    try:
        model = cls.from_pretrained(str(source / component), quantization_config=config,
                                   dtype=torch.bfloat16, device_map={'': 'cpu'},
                                   local_files_only=True)
    except Exception:
        if out.exists():
            shutil.rmtree(out)
        raise
    finally:
        import sdnq.quantizer as quantizer
        quantizer.sdnq_quantize_layer = original_quantize
    counts = Counter()
    for parameter in model.parameters():
        counts[str(parameter.dtype)] += parameter.numel()
    classes = Counter(type(module).__name__ for module in model.modules())
    quantized = [name for name, module in model.named_modules()
                 if 'SDNQ' in type(module).__name__ or 'Quant' in type(module).__name__]
    if not quantized:
        raise RuntimeError(f'No SDNQ modules produced for {component}')
    out.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(out), safe_serialization=True, max_shard_size='5GB')
    saved = json.loads((out / 'config.json').read_text(encoding='utf-8'))
    quant = saved.get('quantization_config') or {}
    if str(quant.get('quant_method', '')).lower() != 'sdnq' or str(quant.get('weights_dtype', '')).lower() != 'uint4':
        raise RuntimeError(f'Saved config is not SDNQ UINT4: {quant}')
    if quant.get('use_quantized_matmul'):
        raise RuntimeError('Refusing to save a CUDA-only quantized-matmul checkpoint')
    report = {'component': component, 'method': 'sdnq', 'weights_dtype': 'uint4',
              'use_svd': True, 'svd_rank': 32, 'use_quantized_matmul': False,
              'skip_modules': skip, 'quantized_modules': quantized,
              'module_types': dict(classes), 'parameter_dtypes': dict(counts),
              'saved_quantization_config': quant}
    (target / f'{component}-quantization.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
    del model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print(f'Saved {component}: {len(quantized)} SDNQ modules', flush=True)
    return report


def copy_support(source, target):
    source, target = Path(source), Path(target)
    target.mkdir(parents=True, exist_ok=True)
    for name in ('processor', 'scheduler', 'vae'):
        shutil.copytree(source / name, target / name, dirs_exist_ok=True,
                        ignore=shutil.ignore_patterns('.cache', '*.lock', '*.incomplete'))
    for name in ('model_index.json', 'LICENSE'):
        shutil.copy2(source / name, target / name)
    (target / 'Notice').write_text(NOTICE, encoding='utf-8')


def verify_source(source, manifest):
    source, manifest = Path(source), Path(manifest)
    payload = json.loads(manifest.read_text(encoding='utf-8'))
    revision = payload.get('revision') or payload.get('sha')
    if revision != REVISION:
        raise ValueError(f'Source manifest revision {revision} != {REVISION}')
    rows = payload.get('files')
    if rows is None:
        rows = [{'path': item['rfilename'], 'size': item['size'], 'sha256': item['lfs']['sha256']}
                for item in payload['siblings'] if item['rfilename'].endswith('.safetensors')]
    if len(rows) != 7:
        raise ValueError('Expected seven official weight files')
    verify_files(source, rows)
    return payload


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', default=str(Path('..') / 'Qwen-Image-2.1' / 'models' / 'bf16'))
    parser.add_argument('--output', default='models/int4')
    parser.add_argument('--manifest', default=str(Path('..') / 'Qwen-Image-2.1' / 'artifacts' / 'download-verification.json'))
    parser.add_argument('--component', choices=['transformer', 'text_encoder', 'all'], default='all')
    parser.add_argument('--device', default='cuda:0')
    args = parser.parse_args()
    verify_source(args.source, args.manifest)
    print(f'Source weights match {REVISION}', flush=True)
    copy_support(args.source, args.output)
    components = ['transformer', 'text_encoder'] if args.component == 'all' else [args.component]
    for component in components:
        quantize_component(args.source, args.output, component, args.device)
    print(f'Quantized pipeline written to {args.output}', flush=True)
    print(f'Diffusers commit used by this project: {DIFFUSERS_COMMIT}', flush=True)


if __name__ == '__main__':
    main()