Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

add diffusers_offload_nonblocking setting

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-08-01 16:38:30 -04:00
parent 52bef2871b
commit 4be093b80f
7 changed files with 22 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_quant.py
+9 -10
View File
@@ -302,7 +302,7 @@ def load_fp8_model_layerwise(checkpoint_info, load_model_func, diffusers_load_co
repo_path = checkpoint_info.path
try:
import torch
from modules import devices
from modules import devices, shared
from diffusers.quantizers import quantization_config
if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'):
setattr(quantization_config.QuantizationMethod, 'LAYERWISE', 'layerwise') # noqa: B010
@@ -315,7 +315,7 @@ def load_fp8_model_layerwise(checkpoint_info, load_model_func, diffusers_load_co
model = load_model_func(repo_path, **load_args)
model = upcast_non_layerwise_modules(model, devices.dtype)
model._skip_layerwise_casting_patterns = None # pylint: disable=protected-access
model.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=False, skip_modules_pattern=[])
model.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=shared.opts.diffusers_offload_nonblocking, skip_modules_pattern=[])
model.layerwise_storage_dtype = storage_dtype
model.quantization_method = 'LayerWise'
except Exception as e:
@@ -336,7 +336,6 @@ def apply_layerwise(sd_model, quiet:bool=False):
storage_dtype = None
log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
return
non_blocking = False
if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'):
setattr(quantization_config.QuantizationMethod, 'LAYERWISE', 'layerwise') # noqa: B010
for module in sd_models.get_signature(sd_model).keys():
@@ -347,25 +346,25 @@ def apply_layerwise(sd_model, quiet:bool=False):
m = getattr(sd_model, module)
if getattr(m, "quantization_method", None) in {'LayerWise', quantization_config.QuantizationMethod.LAYERWISE}: # pylint: disable=no-member
storage_dtype = getattr(m, "layerwise_storage_dtype", storage_dtype)
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=shared.opts.diffusers_offload_nonblocking)
elif module.startswith('unet') and ('Model' in shared.opts.layerwise_quantization):
if hasattr(m, 'enable_layerwise_casting'):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=shared.opts.diffusers_offload_nonblocking)
m.layerwise_storage_dtype = storage_dtype
m.quantization_method = 'LayerWise'
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not shared.opts.diffusers_offload_nonblocking}')
elif module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization):
if hasattr(m, 'enable_layerwise_casting'):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=shared.opts.diffusers_offload_nonblocking)
m.layerwise_storage_dtype = storage_dtype
m.quantization_method = 'LayerWise'
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not shared.opts.diffusers_offload_nonblocking}')
elif module.startswith('text_encoder') and ('TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
if hasattr(m, 'enable_layerwise_casting'):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=shared.opts.diffusers_offload_nonblocking)
m.layerwise_storage_dtype = storage_dtype
m.quantization_method = quantization_config.QuantizationMethod.LAYERWISE # pylint: disable=no-member
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not shared.opts.diffusers_offload_nonblocking}')
except Exception as e:
if 'Hook with name' not in str(e):
log.error(f'Quantization: type=layerwise {e}')