File size: 4,212 Bytes
ba48d54 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Serialize standard bitsandbytes INT8 components for a Diffusers pipeline."""
import argparse
import gc
import json
import shutil
from collections import Counter
from pathlib import Path
import torch
from diffusers import BitsAndBytesConfig as DBitsAndBytesConfig
from diffusers import QwenImage21Transformer2DModel
from transformers import BitsAndBytesConfig as TBitsAndBytesConfig
from transformers import Qwen3VLForConditionalGeneration
TRANSFORMER_SKIP = ['proj_out', 'img_in', 'txt_in', 'time_text_embed', 'modulation', 'norm_out']
TEXT_SKIP = ['lm_head', 'visual']
def quantize_component(source, target, component):
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':
cls, config_cls, skip = QwenImage21Transformer2DModel, DBitsAndBytesConfig, TRANSFORMER_SKIP
elif component == 'text_encoder':
cls, config_cls, skip = Qwen3VLForConditionalGeneration, TBitsAndBytesConfig, TEXT_SKIP
else:
raise ValueError(component)
config = config_cls(load_in_8bit=True, llm_int8_threshold=6.0,
llm_int8_skip_modules=skip)
model = cls.from_pretrained(str(source / component), quantization_config=config,
dtype=torch.bfloat16, device_map={'': 0},
local_files_only=True)
import bitsandbytes as bnb
modules = [name for name, mod in model.named_modules() if isinstance(mod, bnb.nn.Linear8bitLt)]
counts = Counter()
for p in model.parameters():
counts[str(p.dtype)] += p.numel()
report = {'component': component, 'method': 'bitsandbytes LLM.int8',
'threshold': 6.0, 'skip_modules': skip, 'linear8bit_modules': modules,
'parameter_dtypes': dict(counts), 'int8_parameters': counts['torch.int8']}
if not modules or not report['int8_parameters']:
raise RuntimeError(f'No integer weights produced for {component}')
out.mkdir(parents=True, exist_ok=True)
model.save_pretrained(str(out), safe_serialization=True, max_shard_size='5GB')
(target / f'{component}-quantization.json').write_text(
json.dumps(report, indent=2), encoding='utf-8')
del model
gc.collect()
torch.cuda.empty_cache()
return report
def copy_support(source, 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(
'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-INT8, by ixim / iximbox.\n'
'Modified files: transformer and text_encoder weight shards, shard indexes and '
'config.json files. Eligible linear weights converted to bitsandbytes INT8; '
'floating-point exceptions are documented in component quantization reports.\n'
'VAE, scheduler and processor retained from the pinned upstream snapshot.\n',
encoding='utf-8')
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--source', default='models/bf16')
ap.add_argument('--output', default='models/int8')
ap.add_argument('--component', choices=['transformer', 'text_encoder', 'all'], default='all')
args = ap.parse_args()
components = ['transformer', 'text_encoder'] if args.component == 'all' else [args.component]
for component in components:
print(f'Quantizing {component}', flush=True)
report = quantize_component(args.source, args.output, component)
print(f'Saved {component}: {report["int8_parameters"]:,} INT8 parameters', flush=True)
copy_support(Path(args.source), Path(args.output))
if __name__ == '__main__':
main()
|