Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

VRAM efficient loading and compile

Browse Source
This commit is contained in:
Disty0
2024-03-14 01:42:36 +03:00
parent e956f8b43b
commit 164ada5805
2 changed files with 44 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_models.py
+13 -30
View File
@@ -1058,32 +1058,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
elif "Kandinsky" in sd_model.__class__.__name__:
sd_model.scheduler.name = 'DDIM'
base_sent_to_cpu=False
if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none') or shared.opts.ipex_optimize or shared.opts.nncf_compress_weights:
if op == 'refiner' and not getattr(sd_model, 'has_accelerate', False):
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
if not shared.opts.diffusers_move_base and refiner_enough_vram:
move_model(sd_model, devices.device)
base_sent_to_cpu=False
else:
if not refiner_enough_vram and not (shared.opts.diffusers_move_base and shared.opts.diffusers_move_refiner):
shared.log.warning(f"Insufficient GPU memory, using system memory as fallback: free={free_vram} GB")
if not shared.opts.shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
shared.log.debug('Enabled moving base model to CPU')
shared.log.debug('Enabled moving refiner model to CPU')
shared.opts.diffusers_move_base=True
shared.opts.diffusers_move_refiner=True
shared.log.debug('Moving base model to CPU')
move_model(model_data.sd_model, devices.cpu)
devices.torch_gc(force=True)
move_model(sd_model, devices.device)
base_sent_to_cpu=True
else:
move_model(sd_model, devices.device)
sd_models_compile.compile_diffusers(sd_model)
if sd_model is None:
shared.log.error('Diffuser model not loaded')
return
@@ -1097,15 +1071,24 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
if hasattr(sd_model, "set_progress_bar_config"):
sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba')
set_diffuser_options(sd_model, vae, op)
if op == 'refiner' and shared.opts.diffusers_move_refiner:
shared.log.debug('Moving refiner model to CPU')
move_model(sd_model, devices.cpu)
else:
move_model(sd_model, devices.device)
if op == 'refiner' and base_sent_to_cpu:
shared.log.debug('Moving base model back to GPU')
move_model(model_data.sd_model, devices.device)
set_diffuser_options(sd_model, vae, op) # offloading should enabled after all
if shared.opts.ipex_optimize:
sd_model = sd_models_compile.ipex_optimize(sd_model)
if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
sd_model = sd_models_compile.nncf_compress_weights(sd_model)
if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none'):
sd_model = sd_models_compile.compile_diffusers(sd_model)
except Exception as e:
shared.log.error("Failed to load diffusers model")
errors.display(e, "loading Diffusers model")
modules/sd_models_compile.py
+31 -7
View File
@@ -35,7 +35,20 @@ def ipex_optimize(sd_model):
import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
model.eval()
model.training = False
model = ipex.optimize(model, dtype=devices.dtype, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init
if model.device.type != "meta":
return_device = model.device
model = ipex.optimize(model.to(devices.device),
dtype=devices.dtype,
inplace=True,
weights_prepack=False
).to(return_device) # pylint: disable=attribute-defined-outside-init
else:
model = ipex.optimize(model,
dtype=devices.dtype,
inplace=True,
weights_prepack=False
) # pylint: disable=attribute-defined-outside-init
devices.torch_gc()
return model
if "Model" in shared.opts.ipex_optimize:
@@ -79,9 +92,10 @@ def nncf_compress_weights(sd_model):
def nncf_compress_model(model):
return_device = model.device
model.eval()
if hasattr(model, "get_input_embeddings"):
backup_embeddings = copy.deepcopy(model.get_input_embeddings())
model = nncf.compress_weights(model.eval().to(devices.device)).to(return_device)
model = nncf.compress_weights(model.to(devices.device)).to(return_device)
if hasattr(model, "set_input_embeddings"):
model.set_input_embeddings(backup_embeddings)
devices.torch_gc(force=True)
@@ -188,7 +202,21 @@ def compile_torch(sd_model):
shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access
def torch_compile_model(model):
return torch.compile(model, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
if model.device.type != "meta":
return_device = model.device
model = torch.compile(model.to(devices.device),
mode=shared.opts.cuda_compile_mode,
backend=shared.opts.cuda_compile_backend,
fullgraph=shared.opts.cuda_compile_fullgraph
).to(return_device)
else:
model = torch.compile(model,
mode=shared.opts.cuda_compile_mode,
backend=shared.opts.cuda_compile_backend,
fullgraph=shared.opts.cuda_compile_fullgraph
)
devices.torch_gc()
return model
if shared.opts.cuda_compile_backend == "openvino_fx":
sd_model = optimize_openvino(sd_model)
@@ -273,10 +301,6 @@ def compile_deepcache(sd_model):
def compile_diffusers(sd_model):
if shared.opts.ipex_optimize:
sd_model = ipex_optimize(sd_model)
if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
sd_model = nncf_compress_weights(sd_model)
if not shared.opts.cuda_compile:
return sd_model
if shared.opts.cuda_compile_backend == 'none':