diff --git a/modules/devices.py b/modules/devices.py index ea7a2c985..e1712219b 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -245,6 +245,8 @@ def torch_gc(force: bool = False, fast: bool = False, reason: str | None = None) torch.xpu.empty_cache() if hasattr(torch.xpu, "ipc_collect"): torch.xpu.ipc_collect() + if torch.cuda.is_available() and hasattr(torch._C, '_host_emptyCache'): # pylint: disable=protected-access + torch._C._host_emptyCache() # pylint: disable=protected-access # freed pinned host blocks otherwise stay cached in-process across model switches except Exception as e: log.error(f'Torch GC: {e}') else: diff --git a/modules/sd_models.py b/modules/sd_models.py index ff834aa6e..0405f7002 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -265,7 +265,10 @@ def move_model(model, device=None, force=False): if hasattr(model, 'device') and model.device == torch.device('meta'): set_execution_device(model, device) elif hasattr(model, 'to'): - model.to(device) + if device == devices.device and getattr(model, 'sdnext_ondemand_modules', None): + pass # the group engine already placed every component; a pipe-level move would only drag on-demand components to the accelerator for the trailing eviction to undo + else: + model.to(device) if hasattr(model, "prior_pipe"): model.prior_pipe.to(device) if device == devices.device: diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 1b0555e6d..229719615 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -45,6 +45,7 @@ def get_signature(cls): def disable_offload(sd_model): + remove_group_offload(sd_model) # group hooks block the meta move at unload, keeping component weights alive for as long as any reference to the pipe survives if not getattr(sd_model, 'has_accelerate', False): return for module_name in get_module_names(sd_model): @@ -95,11 +96,11 @@ def group_offload_config(main: bool) -> dict: } -def remove_group_offload_component(module): +def remove_group_offload_component(module) -> bool: if getattr(module, 'sdnext_group_offload_sig', None) is None: module = getattr(module, 'model', None) # wrapper components carry the hooks on the inner model if module is None or getattr(module, 'sdnext_group_offload_sig', None) is None: - return + return False from diffusers.hooks.group_offloading import _GROUP_OFFLOADING, _LAYER_EXECUTION_TRACKER, _LAZY_PREFETCH_GROUP_OFFLOADING from diffusers.hooks.hooks import HookRegistry registry = HookRegistry.check_if_exists_or_initialize(module) @@ -107,14 +108,14 @@ def remove_group_offload_component(module): registry.remove_hook(_LAYER_EXECUTION_TRACKER, recurse=True) registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=True) module.sdnext_group_offload_sig = None + return True def remove_group_offload(sd_model): removed = [] for module_name in get_module_names(sd_model): module = getattr(sd_model, module_name, None) - if isinstance(module, torch.nn.Module) and getattr(module, 'sdnext_group_offload_sig', None) is not None: - remove_group_offload_component(module) + if isinstance(module, torch.nn.Module) and remove_group_offload_component(module): removed.append(module_name) for module_name in getattr(sd_model, 'sdnext_ondemand_modules', None) or []: module = getattr(sd_model, module_name, None)