diff --git a/CHANGELOG.md b/CHANGELOG.md index e5d5833ef..e23ea25d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,16 +10,17 @@ if you have a compatible nVidia GPU, Nunchaku is the fastest quantization engine, currently available for Flux.1, SANA and Qwen-Image models *note*: release version of `nunchaku==0.3.2` does NOT include support, so you need to build [nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html) from source - updated [SD.Next Model Samples Gallery](https://vladmandic.github.io/sd-samples/compare.html) -- **Core** - - enable offload during pre-forward by default - - improve offloading of very large models - - update `requirements` - **UI** - improved image scaling in img2img and control interfaces - add base model type to networks display, thanks @Artheriax - additional hints to ui, thanks @Artheriax - add video support to gallery, thanks @CalamitousFelicitousness - additional artwork for reference models in networks, thanks @liutyi +- **Offloading** + - enable offload during pre-forward by default + - improve offloading of models with multiple dits + - improve offloading of models with impliciy vae processing + - improve offloading of models with controlnet - **Fixes** - normalize path hanlding when deleting images - fix hidden model tags in networks display diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index 3d1913165..416dfa6b5 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -313,6 +313,7 @@ class ControlNet(): errors.display(e, 'Control') if self.model is None: return + self.model.offload_never = True if self.dtype is not None: self.model.to(self.dtype) if "Control" in opts.sdnq_quantize_weights: @@ -441,7 +442,7 @@ class ControlNetPipeline(): tokenizer=pipeline.tokenizer, transformer=pipeline.transformer, scheduler=pipeline.scheduler, - controlnet=controlnets, # can be a list + controlnet=controlnets[0] if isinstance(controlnets, list) else controlnets, # can be a list ) elif len(loras) > 0: self.pipeline = pipeline @@ -463,11 +464,13 @@ class ControlNetPipeline(): if dtype is not None: self.pipeline = self.pipeline.to(dtype) + controlnet = None # free up memory + controlnets = None sd_models.copy_diffuser_options(self.pipeline, pipeline) if opts.diffusers_offload_mode == 'none': sd_models.move_model(self.pipeline, devices.device) - from modules.sd_models import set_diffuser_offload - set_diffuser_offload(self.pipeline, 'model') + sd_models.clear_caches() + sd_models.set_diffuser_offload(self.pipeline, 'model') t1 = time.time() debug_log(f'Control {what} pipeline: class={self.pipeline.__class__.__name__} time={t1-t0:.2f}') diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index 9c46ea7c2..82ca4ec5d 100644 --- a/modules/sd_hijack_te.py +++ b/modules/sd_hijack_te.py @@ -19,15 +19,14 @@ def hijack_encode_prompt(*args, **kwargs): res = None t1 = time.time() timer.process.add('te', t1-t0) - if hasattr(shared.sd_model, "maybe_free_model_hooks"): - shared.sd_model.maybe_free_model_hooks() + # if hasattr(shared.sd_model, "maybe_free_model_hooks"): + # shared.sd_model.maybe_free_model_hooks() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) shared.state.end() return res def init_hijack(pipe): - if shared.opts.te_hijack and pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'): - # shared.log.debug(f'Model: cls={pipe.__class__.__name__} hijack encode') + if pipe is not None and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'): pipe.orig_encode_prompt = pipe.encode_prompt pipe.encode_prompt = hijack_encode_prompt diff --git a/modules/sd_hijack_vae.py b/modules/sd_hijack_vae.py new file mode 100644 index 000000000..26c36b347 --- /dev/null +++ b/modules/sd_hijack_vae.py @@ -0,0 +1,64 @@ +import os +import time +import torch +from modules import shared, sd_models, devices, timer, errors + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def hijack_vae_decode(*args, **kwargs): + shared.state.begin('VAE') + t0 = time.time() + res = None + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + try: + sd_models.move_model(shared.sd_model.vae, devices.device) + if torch.is_tensor(args[0]): + latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype + res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs) + t1 = time.time() + shared.log.debug(f'Decode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') + else: + res = shared.sd_model.vae.orig_decode(*args, **kwargs) + except Exception as e: + shared.log.error(f'Decode: vae={shared.sd_model.vae.__class__.__name__} {e}') + errors.display(e, 'vae') + res = None + t1 = time.time() + timer.process.add('vae', t1-t0) + shared.state.end() + return res + + +def hijack_vae_encode(*args, **kwargs): + shared.state.begin('VAE') + t0 = time.time() + res = None + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + try: + sd_models.move_model(shared.sd_model.vae, devices.device) + if torch.is_tensor(args[0]): + latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype + res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs) + t1 = time.time() + shared.log.debug(f'Encode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') + else: + res = shared.sd_model.vae.orig_encode(*args, **kwargs) + except Exception as e: + shared.log.error(f'Encode: vae={shared.sd_model.vae.__class__.__name__} {e}') + errors.display(e, 'vae') + res = None + t1 = time.time() + timer.process.add('vae', t1-t0) + shared.state.end() + return res + + +def init_hijack(pipe): + if pipe is not None and hasattr(pipe, 'vae') and hasattr(pipe.vae, 'decode') and not hasattr(pipe.vae, 'orig_decode'): + pipe.vae.orig_decode = pipe.vae.decode + pipe.vae.decode = hijack_vae_decode + if pipe is not None and hasattr(pipe, 'vae') and hasattr(pipe.vae, 'encode') and not hasattr(pipe.vae, 'orig_encode'): + pipe.vae.orig_encode = pipe.vae.encode + pipe.vae.encode = hijack_vae_encode diff --git a/modules/sd_models.py b/modules/sd_models.py index fef82c570..85682ccc8 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -57,22 +57,6 @@ i2i_pipes = [ ] -def copy_diffuser_options(new_pipe, orig_pipe): - new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None) - new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None) - new_pipe.embedding_db = getattr(orig_pipe, 'embedding_db', None) - new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None) - new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False) - new_pipe.current_attn_name = getattr(orig_pipe, 'current_attn_name', None) - new_pipe.default_scheduler = getattr(orig_pipe, 'default_scheduler', None) - new_pipe.is_sdxl = getattr(orig_pipe, 'is_sdxl', False) # a1111 compatibility item - new_pipe.is_sd2 = getattr(orig_pipe, 'is_sd2', False) - new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True) - add_noise_pred_to_diffusers_callback(new_pipe) - if new_pipe.has_accelerate: - set_accelerate(new_pipe) - - def set_huggingface_options(): if shared.opts.diffusers_to_gpu: # and model_type.startswith('Stable Diffusion'): sd_hijack_accelerate.hijack_accelerate() @@ -872,6 +856,28 @@ def clean_diffuser_pipe(pipe): pipe.register_to_config(**internal_dict) +def copy_diffuser_options(new_pipe, orig_pipe): + new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None) + new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None) + new_pipe.embedding_db = getattr(orig_pipe, 'embedding_db', None) + new_pipe.loaded_loras = getattr(orig_pipe, 'loaded_loras', {}) + new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None) + new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False) + new_pipe.current_attn_name = getattr(orig_pipe, 'current_attn_name', None) + new_pipe.default_scheduler = getattr(orig_pipe, 'default_scheduler', None) + new_pipe.image_encoder = getattr(orig_pipe, 'image_encoder', None) + new_pipe.feature_extractor = getattr(orig_pipe, 'feature_extractor', None) + new_pipe.mask_processor = getattr(orig_pipe, 'mask_processor', None) + new_pipe.restore_pipeline = getattr(orig_pipe, 'restore_pipeline', None) + new_pipe.task_args = getattr(orig_pipe, 'task_args', None) + new_pipe.is_sdxl = getattr(orig_pipe, 'is_sdxl', False) # a1111 compatibility item + new_pipe.is_sd2 = getattr(orig_pipe, 'is_sd2', False) + new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True) + add_noise_pred_to_diffusers_callback(new_pipe) + if new_pipe.has_accelerate: + set_accelerate(new_pipe) + + def backup_pipe_components(pipe): if pipe is None: return {} diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 8678d8002..0a35e3950 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -203,18 +203,19 @@ class OffloadHook(accelerate.hooks.ModelHook): return module def pre_forward(self, module, *args, **kwargs): - if self.last_pre != id(module): # offload every other module first time when new module starts pre-forward - self.last_pre = id(module) + _id = id(module) + if self.last_pre != _id and not hasattr(module, "offload_never"): # offload every other module first time when new module starts pre-forward + self.last_pre = _id if shared.opts.diffusers_offload_pre: debug_move(f'Offload: type=balanced op=pre module={module.__class__.__name__}') for pipe in get_pipe_variants(): for module_name in get_module_names(pipe): module_instance = getattr(pipe, module_name, None) module_cls = module_instance.__class__.__name__ - if (id(module) != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)): + if (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)): apply_balanced_offload_to_module(module_instance, op='pre') - if not devices.same_device(module.device, devices.device): + if not devices.same_device(module.device, devices.device): # move-to-device device_index = torch.device(devices.device).index if device_index is None: device_index = 0 @@ -233,6 +234,13 @@ class OffloadHook(accelerate.hooks.ModelHook): module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory + + if debug: + for pipe in get_pipe_variants(): + for module_name in get_module_names(pipe): + module_instance = getattr(pipe, module_name, None) + shared.log.trace(f'Offload: type=balanced op=pre check module={module_instance.__class__.__name__} device={module_instance.device} dtype={module_instance.dtype}') + return args, kwargs def post_forward(self, module, output): @@ -292,7 +300,7 @@ def get_module_sizes(pipe=None, exclude=[]): return modules -def move_module_to_cpu(module, op='unk'): +def move_module_to_cpu(module, op='unk', force:bool=False): try: module_name = getattr(module, "module_name", module.__class__.__name__) module_size = offload_hook_instance.offload_map.get(module_name, offload_hook_instance.model_size()) @@ -301,7 +309,11 @@ def move_module_to_cpu(module, op='unk'): prev_gpu = used_gpu module_cls = module.__class__.__name__ op = f'{op}:skip' - if module_cls in offload_hook_instance.offload_never: + if force: + op = f'{op}:force' + module = module.to(devices.cpu) + used_gpu -= module_size + elif module_cls in offload_hook_instance.offload_never: op = f'{op}:never' elif module_cls in offload_hook_instance.offload_always: op = f'{op}:always' @@ -313,7 +325,7 @@ def move_module_to_cpu(module, op='unk'): used_gpu -= module_size if debug: quant = getattr(module, "quantization_method", None) - debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}') + debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f}:{shared.opts.diffusers_offload_min_gpu_memory} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}') except Exception as e: if 'out of memory' in str(e): devices.torch_gc(fast=True, force=True, reason='oom') @@ -325,7 +337,7 @@ def move_module_to_cpu(module, op='unk'): errors.display(e, f'Offload: type=balanced op=apply module={getattr(module, "__name__", None)}') -def apply_balanced_offload_to_module(module, op="apply"): +def apply_balanced_offload_to_module(module, op="apply", force:bool=False): module_name = getattr(module, "module_name", module.__class__.__name__) network_layer_name = getattr(module, "network_layer_name", None) device_map = getattr(module, "balanced_offload_device_map", None) @@ -334,7 +346,7 @@ def apply_balanced_offload_to_module(module, op="apply"): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) except Exception as e: shared.log.warning(f'Offload remove hook: module={module_name} {e}') - move_module_to_cpu(module, op=op) + move_module_to_cpu(module, op=op, force=force) try: module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) except Exception as e: @@ -345,7 +357,7 @@ def apply_balanced_offload_to_module(module, op="apply"): if device_map and max_memory: module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory - module.offload_post = shared.sd_model_type in offload_post and shared.opts.te_hijack and module_name.startswith("text_encoder") + module.offload_post = shared.sd_model_type in offload_post and module_name.startswith("text_encoder") if shared.opts.layerwise_quantization or getattr(module, 'quantization_method', None) == 'LayerWise': model_quant.apply_layerwise(module, quiet=True) # need to reapply since hooks were removed/readded devices.torch_gc(fast=True, force=True, reason='offload') diff --git a/modules/shared.py b/modules/shared.py index 407d15fa1..f389d0e0d 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -251,7 +251,6 @@ options_templates.update(options_section(('text_encoder', "Text Encoder"), { "sd_textencoder_cache_size": OptionInfo(4, "Text encoder cache size", gr.Slider, {"minimum": 0, "maximum": 16, "step": 1}), "sd_textencder_linebreak": OptionInfo(True, "Use line break as prompt segment marker", gr.Checkbox), "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox), - "te_hijack": OptionInfo(True, "Offload after prompt encode", gr.Checkbox), "te_optional_sep": OptionInfo("