From ee1a4c607d46e23a329c239478a48754566ba84a Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 15 Apr 2025 10:06:34 -0400 Subject: [PATCH] offload cleanup Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 4 ++++ installer.py | 2 +- modules/model_hidream.py | 10 +++++++--- modules/sd_models.py | 2 ++ modules/sd_offload.py | 30 +++++++++++++++++------------- 5 files changed, 31 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ff351dd89..e2c4fe086 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Change Log for SD.Next +## Update for 2025-04-15 + +- **HiDream** optimized offloading, now works in 12GB VRAM / 26GB RAM + ## Update for 2025-04-14 - [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) new guidance method optimized for flow-matching models diff --git a/installer.py b/installer.py index 005e81dcf..7d9336c5c 100644 --- a/installer.py +++ b/installer.py @@ -538,7 +538,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git or args.experimental: return - sha = 'a8f5134c113da402a93580ef7a021557e816c98d' # diffusers commit hash + sha = 'b6156aafe998eb57902efd3b8cce9a6fde35c1ea' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' diff --git a/modules/model_hidream.py b/modules/model_hidream.py index c9702e7cf..9f37d0f98 100644 --- a/modules/model_hidream.py +++ b/modules/model_hidream.py @@ -24,7 +24,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): **quant_args, ) if shared.opts.diffusers_offload_mode != 'none': - transformer = transformer.to(devices.cpu) + sd_models.move_model(transformer, devices.cpu) load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') @@ -36,7 +36,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): **quant_args, ) if shared.opts.diffusers_offload_mode != 'none': - text_encoder_3 = text_encoder_3.to(devices.cpu) + sd_models.move_model(text_encoder_3, devices.cpu) load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True) shared.log.debug(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') @@ -55,7 +55,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): **load_args, ) if shared.opts.diffusers_offload_mode != 'none': - text_encoder_4 = text_encoder_4.to(devices.cpu) + sd_models.move_model(text_encoder_4, devices.cpu) load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') @@ -69,6 +69,10 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): **load_args, ) sd_hijack_te.init_hijack(pipe) + del text_encoder_3 + del text_encoder_4 + del tokenizer_4 + del transformer devices.torch_gc() return pipe diff --git a/modules/sd_models.py b/modules/sd_models.py index 25e0c7005..bed80bdb5 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -225,6 +225,8 @@ def move_model(model, device=None, force=False): pass # ignore model move if sequential offload is enabled elif 'Params4bit' in str(e0) or 'Params8bit' in str(e0): pass # ignore model move if quantization is enabled + elif 'already been set to the correct devices' in str(e0): + pass # ignore errors on pre-quant models else: raise e0 t1 = time.time() diff --git a/modules/sd_offload.py b/modules/sd_offload.py index a9273d38d..9352d4f38 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -9,8 +9,10 @@ from modules import shared, devices, errors, model_quant from modules.timer import process as process_timer -debug_move = log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None -should_offload = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4'] +debug = os.environ.get('SD_MOVE_DEBUG', None) is not None +debug_move = log.trace if debug else lambda *args, **kwargs: None +offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4'] +offload_post = ['h1'] offload_hook_instance = None balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline'] @@ -66,7 +68,7 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False): if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if shared.opts.diffusers_offload_mode == "none": - if shared.sd_model_type in should_offload or 'video' in shared.sd_model_type: + if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type: shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') else: shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') @@ -175,19 +177,20 @@ class OffloadHook(accelerate.hooks.ModelHook): return args, kwargs def post_forward(self, module, output): - if getattr(module, "do_offload", False) and shared.opts.te_hijack and module.device != devices.cpu: + if getattr(module, "offload_post", False) and module.device != devices.cpu: used_gpu, used_ram = devices.torch_gc(fast=True) perc_gpu = used_gpu / shared.gpu_memory try: module_size = self.model_size() prev_gpu = used_gpu - do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) - if do_offload: + offload_now = perc_gpu > shared.opts.diffusers_offload_min_gpu_memory + if offload_now: module = module.to(devices.cpu) used_gpu -= module_size - cls = module.__class__.__name__ - quant = getattr(module, "quantization_method", None) - debug_move(f'Offload: type=balanced op={"move post forward" if do_offload else "skip post forward"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}') + if debug: + cls = module.__class__.__name__ + quant = getattr(module, "quantization_method", None) + debug_move(f'Offload: type=balanced op={"post" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} 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') @@ -269,15 +272,16 @@ def apply_balanced_offload(sd_model=None, exclude=[]): perc_gpu = used_gpu / shared.gpu_memory try: prev_gpu = used_gpu - do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu) - if do_offload: + offload_now = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu) + if offload_now: module = module.to(devices.cpu) used_gpu -= module_size cls = module.__class__.__name__ quant = getattr(module, "quantization_method", None) if not cached: shared.log.debug(f'Model module={module_name} type={cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}') - debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}') + if debug: + debug_move(f'Offload: type=balanced op={"move" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} 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') @@ -295,7 +299,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]): if device_map and max_memory: module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory - module.do_offload = bool("HiDreamImage" in sd_model.__class__.__name__ and module_name.startswith("text_encoder")) + module.offload_post = shared.sd_model_type in [offload_post] and shared.opts.te_hijack and module_name.startswith("text_encoder") devices.torch_gc(fast=True, force=True, reason='offload') apply_balanced_offload_to_module(sd_model)