File size: 5,028 Bytes
a5e8975 | 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 91 92 93 94 95 96 97 | """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
TRANSFORMER_SKIP = ['proj_out', 'img_in', 'txt_in', 'time_text_embed', 'modulation', 'norm_out']
TEXT_SKIP = ['lm_head', 'visual', 'model.visual']
def quantize_component(source, target, component):
import torch
from diffusers import BitsAndBytesConfig as DBitsAndBytesConfig, QwenImage21Transformer2DModel
from transformers import BitsAndBytesConfig as TBitsAndBytesConfig, Qwen3VLForConditionalGeneration
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)
from scripts.integrity import sha256
source_files = [{'path': p.relative_to(source).as_posix(), 'size': p.stat().st_size,
'sha256': sha256(p)} for p in sorted((source / component).glob('*'))
if p.is_file() and (p.suffix == '.json' or p.suffix == '.safetensors')]
import bitsandbytes as bnb
modules = [name for name, mod in model.named_modules() if isinstance(mod, bnb.nn.Linear8bitLt)]
if component == 'text_encoder' and any(name.startswith(('visual.', 'model.visual.')) for name in modules):
raise RuntimeError('The vision tower must remain in floating point')
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'],
'source_files': source_files}
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', newline='\n')
del model
gc.collect()
torch.cuda.empty_cache()
return report
def copy_support(source, target, derivative='Image21-INT8'):
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)
if (source / 'upstream.json').exists():
shutil.copy2(source / 'upstream.json', target / 'upstream.json')
(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'
f'Independent derivative: {derivative}, 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', newline='\n')
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')
ap.add_argument('--name', default='Image21-INT8', help='Independent derivative name for attribution')
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), args.name)
if __name__ == '__main__':
main()
|