From 290970e536d837ab8d63a1ea811c671cdc4a7b9c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 9 Nov 2023 12:57:46 -0500 Subject: [PATCH] safe move offloads --- modules/processing_diffusers.py | 12 ++++++------ modules/sd_models.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 54564f822..f401a7b0a 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -89,12 +89,12 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro def full_vae_decode(latents, model): t0 = time.time() - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): shared.log.debug('Moving to CPU: model=UNet') unet_device = model.unet.device model.unet.to(devices.cpu) devices.torch_gc() - if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload: + if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'): model.vae.to(devices.device) latents.to(model.vae.device) @@ -104,7 +104,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0] - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): model.unet.to(unet_device) t1 = time.time() shared.log.debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={upcast} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}') @@ -112,15 +112,15 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro def full_vae_encode(image, model): shared.log.debug(f'VAE encode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}') - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): shared.log.debug('Moving to CPU: model=UNet') unet_device = model.unet.device model.unet.to(devices.cpu) devices.torch_gc() - if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload: + if not shared.cmd_opts.lowvram and not shared.opts.diffusers_seq_cpu_offload and hasattr(model, 'vae'): model.vae.to(devices.device) encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() - if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False): + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): model.unet.to(unet_device) return encoded diff --git a/modules/sd_models.py b/modules/sd_models.py index ddea17a01..b9d673e5b 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -777,7 +777,7 @@ def set_diffuser_options(sd_model, vae, op: str): shared.log.debug(f'Setting {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}') if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'): sd_model.enable_xformers_memory_efficient_attention() - if shared.opts.opt_channelslast: + if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'): shared.log.debug(f'Setting {op}: enable channels last') sd_model.unet.to(memory_format=torch.channels_last)