Fix model offload by not focring the model to GPU

This commit is contained in:
Kubuxu
2023-07-29 20:13:36 +01:00
parent 3258b27523
commit a4a68f33e3
+4 -3
View File
@@ -702,7 +702,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:
if op == 'refiner' and not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
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
@@ -721,7 +721,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:
elif not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
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}")
@@ -752,7 +752,8 @@ 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:
elif not shared.opts.diffusers_seq_cpu_offload and not shared.opts.diffusers_model_cpu_offload:
# In offload modes, accelerate will move models around.
sd_model.to(devices.device)
if op == 'refiner' and base_sent_to_cpu:
shared.log.debug('Moving base model back to GPU')