From 018040256320163363e27103600a20e3632bd4a6 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 30 Jul 2023 11:36:47 +0300 Subject: [PATCH] Better move and accelerate handling --- modules/processing_diffusers.py | 12 ++++++------ modules/sd_models.py | 24 ++++++++---------------- modules/sd_vae.py | 4 ++-- 3 files changed, 16 insertions(+), 24 deletions(-) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index ce1457967..ec2683b0f 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -36,7 +36,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro def vae_decode(latents, model, output_type='np'): if hasattr(model, 'vae') and torch.is_tensor(latents): shared.log.debug(f'Diffusers VAE decode: name={model.vae.config.get("_name_or_path", "default")} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}') - if shared.opts.diffusers_move_unet: + if shared.opts.diffusers_move_unet and not model.has_accelerate: shared.log.debug('Diffusers: Moving UNet to CPU') unet_device = model.unet.device model.unet.to(devices.cpu) @@ -44,7 +44,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro latents.to(model.vae.device) decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0] imgs = model.image_processor.postprocess(decoded, output_type=output_type) - if shared.opts.diffusers_move_unet: + if shared.opts.diffusers_move_unet and not model.has_accelerate: model.unet.to(unet_device) return imgs else: @@ -134,7 +134,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro if shared.state.interrupted or shared.state.skipped: return results - if shared.opts.diffusers_move_base: + if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate: shared.sd_model.to(devices.device) refiner_enabled = shared.sd_refiner is not None and p.enable_hr @@ -168,7 +168,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro for i in range(len(decoded)): images.save_image(decoded[i], path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-refiner") - if shared.opts.diffusers_move_base: + if shared.opts.diffusers_move_base and not shared.sd_model.has_accelerate: shared.log.debug('Diffusers: Moving base model to CPU') shared.sd_model.to('cpu') devices.torch_gc() @@ -182,7 +182,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro if shared.state.interrupted or shared.state.skipped: return results - if shared.opts.diffusers_move_refiner: + if shared.opts.diffusers_move_refiner and not shared.sd_refiner.has_accelerate: shared.sd_refiner.to(devices.device) p.ops.append('refine') for i in range(len(output.images)): @@ -205,7 +205,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro refiner_images = vae_decode(refiner_output.images, shared.sd_refiner) results.append(refiner_images[0]) - if shared.opts.diffusers_move_refiner: + if shared.opts.diffusers_move_refiner and not shared.sd_refiner.has_accelerate: shared.log.debug('Diffusers: Moving refiner model to CPU') shared.sd_refiner.to('cpu') else: diff --git a/modules/sd_models.py b/modules/sd_models.py index fe39ea5a9..6c3ea782e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -659,18 +659,10 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No 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: + if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload: + shared.log.debug(f'Diffusers {op}: enable sequential 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 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') @@ -753,7 +745,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model.sd_model_hash = checkpoint_info.hash # pylint: disable=attribute-defined-outside-init 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') - if op == 'refiner' and shared.opts.diffusers_move_refiner: + if op == 'refiner' and shared.opts.diffusers_move_refiner and not sd_model.has_accelerate: shared.log.debug('Moving refiner model to CPU') sd_model.to("cpu") elif not sd_model.has_accelerate: @@ -943,12 +935,12 @@ 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 sd_model.has_accelerate: + if not sd_model.has_accelerate: if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.send_everything_to_cpu() else: sd_model.to(devices.cpu) - if reuse_dict or (shared.opts.model_reuse_dict and sd_model is not None): + if (reuse_dict or (shared.opts.model_reuse_dict and sd_model is not None)) and not sd_model.has_accelerate: shared.log.info('Reusing previous model dictionary') sd_hijack.model_hijack.undo_hijack(sd_model) else: @@ -980,7 +972,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 sd_model.has_accelerate): + if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not sd_model.has_accelerate: sd_model.to(devices.device) timer.record("device") shared.log.info(f"Weights loaded in {timer.summary()}") @@ -990,7 +982,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 model_data.sd_model.has_accelerate: + if 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) @@ -998,7 +990,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 model_data.sd_refiner.has_accelerate: + if 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 3843d6496..acab7cfda 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 sd_model.has_accelerate: + if 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 sd_model.has_accelerate): + if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not sd_model.has_accelerate: sd_model.to(devices.device) shared.log.info(f"VAE weights loaded: {vae_file}") return sd_model