From a4a68f33e37da74f16beca791963e01c07d25a85 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Sat, 29 Jul 2023 20:13:36 +0100 Subject: [PATCH 1/3] Fix model offload by not focring the model to GPU --- modules/sd_models.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 341083477..b0fe7d1a2 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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') From 3c9eb2cd65c3c6606b732daae2909af00b81c455 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Sat, 29 Jul 2023 20:38:44 +0100 Subject: [PATCH 2/3] Introduce sd_model.has_accelerate --- modules/sd_models.py | 35 ++++++++++++++++++----------------- modules/sd_vae.py | 4 ++-- 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index b0fe7d1a2..ec0cbee3c 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -657,23 +657,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') @@ -702,7 +704,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 @@ -721,7 +723,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}") @@ -752,7 +754,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: @@ -910,7 +912,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 @@ -937,7 +938,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: @@ -974,7 +975,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()}") @@ -984,7 +985,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) @@ -992,7 +993,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) diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 53b61ae01..3843d6496 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -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 From b01ada06d5822fcc8c2d145298f033bdc849f708 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Sat, 29 Jul 2023 21:12:43 +0100 Subject: [PATCH 3/3] Fix pipeline switching --- modules/sd_models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index ec0cbee3c..b89df92fe 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -790,6 +790,7 @@ def set_diffuser_pipe(pipe, new_pipe_type): sd_checkpoint_info = pipe.sd_checkpoint_info sd_model_checkpoint = pipe.sd_model_checkpoint sd_model_hash = pipe.sd_model_hash + has_accelerate = pipe.has_accelerate if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE: new_pipe = diffusers.AutoPipelineForText2Image.from_pipe(pipe) @@ -804,6 +805,7 @@ def set_diffuser_pipe(pipe, new_pipe_type): new_pipe.sd_checkpoint_info = sd_checkpoint_info new_pipe.sd_model_checkpoint = sd_model_checkpoint new_pipe.sd_model_hash = sd_model_hash + new_pipe.has_accelerate = has_accelerate model_data.sd_model = new_pipe shared.log.info(f"Pipeline class changed from {pipe.__class__.__name__} to {new_pipe.__class__.__name__}") @@ -930,7 +932,7 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model') else: model_data.sd_dict = 'None' shared.log.debug(f'Load model weights: existing={sd_model is not None} target={checkpoint_info.filename} info={info}') - if not sd_model: + if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner if sd_model is None: # previous model load failed current_checkpoint_info = None