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

Introduce sd_model.has_accelerate

Browse Source
This commit is contained in:
Kubuxu
2023-07-29 20:38:44 +01:00
committed by Disty0
parent f945cf14b0
commit 7e1030d499
2 changed files with 20 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_models.py
+18 -17
View File
@@ -652,23 +652,25 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
if hasattr(sd_model, "watermark"):
sd_model.watermark = NoWatermark()
sd_model.has_accelerate = False
if hasattr(sd_model, "enable_model_cpu_offload"):
if (shared.cmd_opts.medvram and devices.backend != "directml") or shared.opts.diffusers_model_cpu_offload:
shared.log.debug(f'Diffusers {op}: enable model CPU offload')
sd_model.enable_model_cpu_offload()
sd_model.has_accelerate = True
if hasattr(sd_model, "enable_sequential_cpu_offload"):
if shared.opts.diffusers_seq_cpu_offload:
sd_model.enable_sequential_cpu_offload(device=devices.device)
sd_model.has_accelerate = True
shared.log.debug(f'Diffusers {op}: enable sequential CPU offload')
if shared.opts.diffusers_move_base or shared.opts.diffusers_move_refiner or shared.opts.diffusers_move_unet:
shared.log.warning("Moving models to CPU is not compatible with sequential CPU offload")
shared.log.debug('Disabled moving base model to CPU')
shared.log.debug('Disabled moving refiner model to CPU')
shared.log.debug('Disabled moving UNet to CPU')
shared.opts.diffusers_move_base=False
shared.opts.diffusers_move_refiner=False
shared.opts.diffusers_move_unet=False
if sd_model.has_accelerate and (shared.opts.diffusers_move_base or shared.opts.diffusers_move_refiner or shared.opts.diffusers_move_unet):
shared.log.warning("Moving models to CPU is not compatible with sequential CPU offload")
shared.log.debug('Disabled moving base model to CPU')
shared.log.debug('Disabled moving refiner model to CPU')
shared.log.debug('Disabled moving UNet to CPU')
shared.opts.diffusers_move_base=False
shared.opts.diffusers_move_refiner=False
shared.opts.diffusers_move_unet=False
if hasattr(sd_model, "enable_vae_slicing"):
if shared.cmd_opts.lowvram or shared.opts.diffusers_vae_slicing:
shared.log.debug(f'Diffusers {op}: enable VAE slicing')
@@ -704,7 +706,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
base_sent_to_cpu=False
if shared.opts.cuda_compile and torch.cuda.is_available():
if op == 'refiner' and not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
if op == 'refiner' and not sd_model.has_accelerate:
gpu_vram = memory_stats().get('gpu', {})
free_vram = gpu_vram.get('total', 0) - gpu_vram.get('used', 0)
refiner_enough_vram = free_vram >= 7 if "StableDiffusionXL" in sd_model.__class__.__name__ else 3
@@ -723,7 +725,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
devices.torch_gc(force=True)
sd_model.to(devices.device)
base_sent_to_cpu=True
elif not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
elif not sd_model.has_accelerate:
sd_model.to(devices.device)
try:
shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_mode}")
@@ -754,7 +756,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
if op == 'refiner' and shared.opts.diffusers_move_refiner:
shared.log.debug('Moving refiner model to CPU')
sd_model.to("cpu")
elif not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
elif not sd_model.has_accelerate:
# In offload modes, accelerate will move models around.
sd_model.to(devices.device)
if op == 'refiner' and base_sent_to_cpu:
@@ -912,7 +914,6 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, timer=None,
devices.torch_gc(force=True)
shared.log.info(f'Model load finished: {memory_stats()} cached={len(checkpoints_loaded.keys())}')
def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model'):
load_dict = shared.opts.sd_model_dict != model_data.sd_dict
global skip_next_load # pylint: disable=global-statement
@@ -939,7 +940,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
current_checkpoint_info = getattr(sd_model, 'sd_checkpoint_info', None)
if current_checkpoint_info is not None and checkpoint_info is not None and current_checkpoint_info.filename == checkpoint_info.filename:
return
if shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload:
if shared.backend == shared.Backend.ORIGINAL or not sd_model.has_accelerate:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
else:
@@ -976,7 +977,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
timer.record("hijack")
script_callbacks.model_loaded_callback(sd_model)
timer.record("callbacks")
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and (shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload):
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and (shared.backend == shared.Backend.ORIGINAL or not sd_model.has_accelerate):
sd_model.to(devices.device)
timer.record("device")
shared.log.info(f"Weights loaded in {timer.summary()}")
@@ -986,7 +987,7 @@ def unload_model_weights(op='model'):
from modules import sd_hijack
if op == 'model' or op == 'dict':
if model_data.sd_model:
if shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload:
if shared.backend == shared.Backend.ORIGINAL or not model_data.sd_model.has_accelerate:
model_data.sd_model.to(devices.cpu)
if shared.backend == shared.Backend.ORIGINAL:
sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
@@ -994,7 +995,7 @@ def unload_model_weights(op='model'):
shared.log.debug(f'Weights unloaded {op}: {memory_stats()}')
else:
if model_data.sd_refiner:
if shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload:
if shared.backend == shared.Backend.ORIGINAL or not model_data.sd_refiner.has_accelerate:
model_data.sd_refiner.to(devices.cpu)
if shared.backend == shared.Backend.ORIGINAL:
sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
modules/sd_vae.py
+2 -2
View File
@@ -232,7 +232,7 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
vae_source = "from function argument"
if loaded_vae_file == vae_file:
return
if shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload:
if shared.backend == shared.Backend.ORIGINAL or not sd_model.has_accelerate:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
else:
@@ -246,7 +246,7 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
sd_hijack.model_hijack.hijack(sd_model)
script_callbacks.model_loaded_callback(sd_model)
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and (shared.backend == shared.Backend.ORIGINAL or not shared.opts.diffusers_seq_cpu_offload):
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and (shared.backend == shared.Backend.ORIGINAL or not sd_model.has_accelerate):
sd_model.to(devices.device)
shared.log.info(f"VAE weights loaded: {vae_file}")
return sd_model