From 7e1030d499d7364512f155db7924156106e2dea2 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Sat, 29 Jul 2023 20:38:44 +0100 Subject: [PATCH] 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 96cc58c1f..8ac7dc9c4 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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) 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