File size: 6,380 Bytes
9116984 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Load Image21-INT4 or the BF16 source on CUDA, Apple Silicon, or CPU.
SDNQ stores UINT4 weights and dequantizes them with ordinary PyTorch operators
unless a CUDA/XPU matmul kernel is requested. This loader does not request that
kernel, so CUDA and MPS execute the same eager implementation. Bitsandbytes is
not used.
"""
import json
import os
from pathlib import Path
os.environ.setdefault('DIFFUSERS_SDNQ_TRANSFORMERS', '1')
import torch
from scripts.device import choose_runtime
def configure_caches():
root = Path(__file__).resolve().parents[1]
os.environ.setdefault('HF_HOME', str(root / '.cache' / 'huggingface'))
os.environ.setdefault('NUMBA_CACHE_DIR', str(root / '.cache' / 'numba'))
def detect_device(requested=None):
if requested:
device = torch.device(requested)
if device.type == 'cuda' and not torch.cuda.is_available():
raise RuntimeError('CUDA was requested but is not available')
if device.type == 'mps' and not torch.backends.mps.is_available():
raise RuntimeError('MPS was requested but is not available')
return device
if torch.cuda.is_available():
return torch.device('cuda')
if torch.backends.mps.is_available():
return torch.device('mps')
return torch.device('cpu')
def accelerator_bytes(device):
if device.type == 'cuda':
return torch.cuda.get_device_properties(device).total_memory
import psutil
return psutil.virtual_memory().total
def read_quantization(model):
path = Path(model) / 'transformer' / 'config.json'
if not path.is_file():
return {}
return json.loads(path.read_text(encoding='utf-8')).get('quantization_config') or {}
def quantization_dict(module):
config = getattr(module, 'config', None)
raw = getattr(config, 'quantization_config', None)
if raw is None:
return {}
if hasattr(raw, 'to_dict'):
return raw.to_dict()
if isinstance(raw, dict):
return raw
return dict(raw)
def _config_value(value):
value = getattr(value, 'value', value)
return str(value).split('.')[-1].lower()
def _require_int4(module, name):
quant = quantization_dict(module)
method = _config_value(quant.get('quant_method', ''))
dtype = _config_value(quant.get('weights_dtype', ''))
if method != 'sdnq' or dtype != 'uint4':
raise ValueError(f'{name} is not an SDNQ UINT4 checkpoint: {method} {dtype}')
if quant.get('use_quantized_matmul'):
raise ValueError(f'{name} enables quantized matmul, which is not part of the portable INT4 runtime')
def _dtype(name):
override = os.environ.get('IMAGE21_DTYPE')
return getattr(torch, override or name)
def _enable_group(module, device, offload_type):
from diffusers.hooks.group_offloading import apply_group_offloading
kwargs = dict(onload_device=device, offload_device=torch.device('cpu'),
offload_type=offload_type, use_stream=False, non_blocking=False)
if offload_type == 'block_level':
kwargs['num_blocks_per_group'] = 1
if hasattr(module, 'enable_group_offload'):
module.enable_group_offload(**kwargs)
else:
apply_group_offloading(module, **kwargs)
def apply_offload(pipe, spec, device):
mode = spec['offload']
if mode == 'resident':
pipe.to(device)
elif mode == 'model':
if device.type != 'cuda':
raise ValueError('Model CPU offload is the CUDA path; Apple Silicon uses resident MPS')
pipe.enable_model_cpu_offload(gpu_id=device.index or 0)
elif mode == 'group':
if device.type != 'cuda':
raise ValueError('Group offload in this loader targets CUDA; Apple Silicon uses resident MPS')
# The transformer is the activation-heavy component, so it moves one block at a time.
# The text encoder moves one linear or embedding at a time: a whole decoder block's
# dequantized MLP does not leave enough room under an 8GB cap.
_enable_group(pipe.transformer, device, 'block_level')
_enable_group(pipe.text_encoder, device, 'leaf_level')
# A 1024 decode keeps several GiB live beside the denoiser cache. On an 8GB
# board that peak does not fit, and tiling changes pixels. Decode on CPU.
pipe.vae.to('cpu')
decode = pipe.vae.decode
encode_image = pipe._encode_vae_image
def _decode_on_cpu(latents, *args, **kwargs):
if torch.is_tensor(latents):
latents = latents.detach().to('cpu')
return decode(latents, *args, **kwargs)
def _encode_on_cpu(image, generator):
device = image.device
encoded = encode_image(image.detach().to('cpu'), generator)
return encoded.to(device=device, dtype=image.dtype)
pipe.vae.decode = _decode_on_cpu
pipe._encode_vae_image = _encode_on_cpu
else:
raise ValueError(f'Unknown offload mode: {mode}')
pipe.image21_runtime = dict(spec, device=str(device))
return pipe
def load_pipeline(model, offload='auto', device=None, local_files_only=False):
"""Load a local directory or a Hub repository id."""
configure_caches()
quant = read_quantization(model) if Path(model).exists() else {'quant_method': 'sdnq', 'weights_dtype': 'uint4'}
method = str(quant.get('quant_method', '')).lower()
if method == 'sdnq':
import sdnq # noqa: F401 (registers the quantizer with diffusers and transformers)
elif quant:
raise ValueError(f'Unsupported quantization config: {quant}')
from diffusers import QwenImage21Pipeline
device = detect_device(device)
spec = choose_runtime(device.type, accelerator_bytes(device))
if offload != 'auto':
if offload not in ('model', 'group', 'resident'):
raise ValueError('offload must be auto, model, group, or resident')
spec = dict(spec, offload=offload)
dtype = _dtype(spec['dtype'])
pipe = QwenImage21Pipeline.from_pretrained(str(model), dtype=dtype, local_files_only=local_files_only)
if method == 'sdnq':
_require_int4(pipe.transformer, 'transformer')
_require_int4(pipe.text_encoder, 'text_encoder')
if quantization_dict(pipe.vae):
raise ValueError('VAE is expected to stay in floating point')
apply_offload(pipe, spec, device)
return pipe
|