Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Better move and accelerate handling

Browse Source
This commit is contained in:
Disty0
2023-07-30 11:36:47 +03:00
parent acc8233f52
commit 0180402563
3 changed files with 16 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/processing_diffusers.py
+6 -6
View File
@@ -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:
modules/sd_models.py
+8 -16
View File
@@ -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)
modules/sd_vae.py
+2 -2
View File
@@ -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