diff --git a/launch.py b/launch.py index e3b0b995e..bb07fd696 100755 --- a/launch.py +++ b/launch.py @@ -155,13 +155,12 @@ def run_extension_installer(ext_dir): # compatibility function def get_memory_stats(detailed:bool=False): - from modules.memstats import ram_stats, memory_stats + from modules.memstats import ram_stats, memory_stats, model_stats if not detailed: res = ram_stats() return f'{res["used"]}/{res["total"]}' else: - res = memory_stats() - return res + return { **memory_stats(), 'model': model_stats(as_gb=True) } # fresh dict: memory_stats returns a module global that the per-generation log also prints def clean_server(): diff --git a/modules/api/models.py b/modules/api/models.py index 7bb0f8f27..bbe5a5771 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -552,6 +552,7 @@ class ResEmbeddings(BaseModel): class ResMemory(BaseModel): ram: dict = Field(title="RAM", description="System memory stats") cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats") + model: dict = Field(default={}, title="Model", description="Loaded model bytes per component and device") class ResScripts(BaseModel): txt2img: list[str] = Field(title="Txt2img", description="Titles of scripts (txt2img)") diff --git a/modules/api/server.py b/modules/api/server.py index 3145ac61a..0e8e2b010 100644 --- a/modules/api/server.py +++ b/modules/api/server.py @@ -197,4 +197,5 @@ def get_memory(): cuda = { 'error': 'unavailable' } except Exception as err: cuda = { 'error': f'{err}' } - return models.ResMemory(ram = ram, cuda = cuda) + from modules import memstats + return models.ResMemory(ram = ram, cuda = cuda, model = memstats.model_stats()) diff --git a/modules/devices.py b/modules/devices.py index f46615165..5826dcf03 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -246,6 +246,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/lora/networks.py b/modules/lora/networks.py index bc3a896ae..35a412f0a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -13,6 +13,32 @@ native_active: bool = False default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_encoder_4', 'unet', 'transformer', 'transformer_2', 'llm_adapter'] +def group_will_mutate(module, network_layer_name: str, loaded) -> bool: + """True when the pass will write to this module: a loaded network covers its layer, a + tensor backup awaits restore, or an svd factor stash awaits removal.""" + if any(net.modules.get(network_layer_name, None) is not None for net in loaded): + return True + weights_backup = getattr(module, 'network_weights_backup', None) + if weights_backup is not None and not isinstance(weights_backup, bool): + return True + bias_backup = getattr(module, 'network_bias_backup', None) + if bias_backup is not None and not isinstance(bias_backup, bool): + return True + return getattr(module, 'sdnq_lora_svd_stash', None) is not None + + +def group_offload_strip(sd_model, component_name: str, stripped: dict): + """Group offload hooks come off before the first weight write in a component: a write + under live hooks either replaces a parameter out of the hook's group bookkeeping or is + lost on the next onload. With hooks removed the weights rest on cpu and the component + reports its truthful device, so writes land in place; the offload reapply at the end + of the pass snapshots the result into fresh groups.""" + component = getattr(sd_model, component_name, None) + sd_models.remove_group_offload_component(component) + stripped[component_name] = component.device + return stripped[component_name] + + def network_activate(include=None, exclude=None): if exclude is None: exclude = [] @@ -32,6 +58,8 @@ def network_activate(include=None, exclude=None): sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + group_offload = shared.opts.diffusers_offload_mode == "group" + group_stripped = {} device = None modules = {} components = include if len(include) > 0 else default_components @@ -67,6 +95,8 @@ def network_activate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks): + device = group_offload_strip(sd_model, component, group_stripped) backup_size += network_backup_weights(module, network_layer_name, component_wanted) if not component_wanted: weights_backup = getattr(module, "network_weights_backup", None) @@ -100,7 +130,7 @@ def network_activate(include=None, exclude=None): if l.debug and len(l.loaded_networks) > 0: log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') modules.clear() - if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential": + if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") @@ -121,6 +151,8 @@ def network_deactivate(include=None, exclude=None): sd_models.move_model(sd_model, device=devices.cpu) elif shared.opts.diffusers_offload_mode == "balanced": sd_model = sd_models.apply_balanced_offload(sd_model, force=True) # dispatched modules hold meta tensors backed by the offload map; rebuild them real on cpu with hooks intact before touching weights + group_offload = shared.opts.diffusers_offload_mode == "group" + group_stripped = {} modules = {} components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer', 'llm_adapter'] @@ -148,6 +180,8 @@ def network_deactivate(include=None, exclude=None): if task is not None: pbar.update(task, advance=1) continue + if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.previously_loaded_networks): + device = group_offload_strip(sd_model, component, group_stripped) batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit) if shared.opts.lora_fuse_native: network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) @@ -163,5 +197,5 @@ def network_deactivate(include=None, exclude=None): if l.debug and len(l.previously_loaded_networks) > 0: log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') modules.clear() - if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential": + if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") diff --git a/modules/memstats.py b/modules/memstats.py index 22023adde..63c95db81 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -115,6 +115,35 @@ def gpu_stats(): return gpu +def model_stats(as_gb: bool = False): + """Loaded-model bytes per component and device, so resident weights can be told from offloaded ones.""" + try: + from modules.modeldata import model_data + pipe = model_data.sd_model # raw slot: the shared.sd_model property can trigger a model load + if pipe is None: + return {} + components = getattr(pipe, 'components', None) or ({ 'model': pipe } if isinstance(pipe, torch.nn.Module) else {}) + placement = {} + seen = set() + for name, component in components.items(): + if not isinstance(component, torch.nn.Module): + continue + devmap = {} + for tensors in (component.parameters(), component.buffers()): + for t in tensors: + ptr = 0 if t.is_meta else t.untyped_storage().data_ptr() + if ptr: + if ptr in seen: # tied weights and offload rewraps share one storage across tensors + continue + seen.add(ptr) + devmap[t.device.type] = devmap.get(t.device.type, 0) + t.numel() * t.element_size() + if devmap: + placement[name] = { d: gb(v) for d, v in devmap.items() } if as_gb else devmap + return placement + except Exception as err: # walk can race a reload or an offload rewrap; every caller is a diagnostic that must not take its caller down + return { 'error': f'{err}' } + + def memory_stats(): mem['ram'] = ram_stats() mem['gpu'] = gpu_stats() diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index bd7ca9cff..d032d5bc9 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -477,6 +477,7 @@ def process_decode(p: processing.StableDiffusionProcessing, output): log.debug(f'Generated: frames={len(output.frames[0])}') output.images = output.frames[0] if output.images is not None and len(output.images) > 0 and isinstance(output.images[0], Image.Image): + sd_models.offload_ondemand(shared.sd_model) # in-pipe decode paths return materialized frames; the vae seam in processing_vae never runs return attach_audio(output.images, audio) model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner if not hasattr(model, 'vae'): diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 638eb8f86..6b721640c 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -189,6 +189,8 @@ def full_vae_encode(image, model): sd_models.move_model(model.unet, devices.cpu) if shared.opts.diffusers_offload_mode != "sequential" and hasattr(model, 'vae'): sd_models.move_model(model.vae, devices.device) + if getattr(model.vae, 'sdnext_ondemand', False): + model.vae.to(devices.device) # the image placement below derives from vae.device, and the entry bridge would onload the weights only after the input is already bound vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default" log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}') @@ -369,6 +371,7 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he if shared.cmd_opts.profile or debug: t1 = time.time() log.debug(f'Profile: VAE decode: {t1-t0:.2f}') + sd_models.offload_ondemand(model) devices.torch_gc() shared.state.end(jobid) return images @@ -393,6 +396,7 @@ def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable else: log.error('VAE not found in model') latents = [] + sd_models.offload_ondemand(model) devices.torch_gc() shared.state.end(jobid) return latents diff --git a/modules/sd_models.py b/modules/sd_models.py index c4a829509..0405f7002 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -16,7 +16,7 @@ from modules.memstats import memory_stats from modules.shared_helpers import walk_files from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoint_titles, get_closest_checkpoint_match, update_model_hashes, write_metadata, checkpoints_list # pylint: disable=unused-import -from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import +from modules.sd_offload import get_module_names, disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate, remove_group_offload_component, offload_ondemand # pylint: disable=unused-import from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, apply_function_to_model, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import @@ -231,13 +231,15 @@ def move_model(model, device=None, force=False): if model is None or device is None: return + if getattr(model, 'sdnext_ondemand', False) and device == devices.device: # on-demand components onload at their entry points instead of pre-moves + return if hasattr(model, 'pipe'): move_model(model.pipe, device, force) fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access if getattr(model, 'vae', None) is not None and get_diffusers_task(model) != DiffusersTaskType.TEXT_2_IMAGE: - if device == devices.device and model.vae.device.type != "meta": # force vae back to gpu if not in txt2img mode + if device == devices.device and model.vae.device.type != "meta" and not getattr(model.vae, 'sdnext_ondemand', False): # force vae back to gpu if not in txt2img mode; on-demand vaes onload at their entry point instead model.vae.to(device) if hasattr(model.vae, '_hf_hook'): debug_move(f'Model move: to={device} class={model.vae.__class__} fn={fn}') # pylint: disable=protected-access @@ -263,9 +265,14 @@ 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: + offload_ondemand(model) # a bulk move must not strand on-demand components on the accelerator; their entry points onload them when needed except Exception as e0: if 'Cannot copy out of meta tensor' in str(e0) or 'must be Tensor, not NoneType' in str(e0): if hasattr(model, "components"): diff --git a/modules/sd_offload.py b/modules/sd_offload.py index a0bfa6901..c7e608aac 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -19,6 +19,7 @@ debug_move = log.trace if debug else lambda *args, **kwargs: None offload_allow_none = ['sd', 'sdxl'] offload_post = ['h1'] offload_hook_instance = None +group_offload_vae_limit = 1.0 # GB; vae-class components above this rest on cpu and onload whole at encode/decode balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline'] no_split_module_classes = [ "Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "Embedding", @@ -44,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): @@ -76,30 +78,257 @@ def set_accelerate(sd_model): set_accelerate_to_module(sd_model.decoder_pipe) -def apply_group_offload(sd_model, op:str='model'): - offload_dct = { - 'onload_device': devices.device, - 'offload_device': devices.cpu, - 'offload_type': shared.opts.group_offload_type, - 'num_blocks_per_group': shared.opts.group_offload_blocks, +def group_offload_config(main: bool) -> dict: + """Effective group offload settings for one component. Components that run once per + generation take the leaf no-stream policy regardless of the main settings, so their + weights are never held in pinned host memory.""" + stream = shared.opts.group_offload_stream if main else False + blocks = max(1, int(shared.opts.group_offload_blocks)) + if stream and blocks != 1: + blocks = 1 # streamed prefetch supports one block per group; upstream clamps with a warning otherwise + return { + 'offload_type': shared.opts.group_offload_type if main else 'leaf_level', + 'num_blocks_per_group': blocks, 'non_blocking': shared.opts.diffusers_offload_nonblocking, - 'use_stream': shared.opts.group_offload_stream, - 'record_stream': shared.opts.group_offload_record, - 'low_cpu_mem_usage': False, + 'use_stream': stream, + 'record_stream': shared.opts.group_offload_record and stream, # record without streams is rejected upstream + 'low_cpu_mem_usage': stream and not shared.opts.group_offload_pin, } - if shared.opts.group_offload_type == 'block_level': - offload_dct['exclude_modules'] = ['vae'] - log.debug(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} options={offload_dct}') - if hasattr(sd_model, "enable_group_offload"): - sd_model.enable_group_offload(**offload_dct) - else: - log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported') + + +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 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) + registry.remove_hook(_GROUP_OFFLOADING, recurse=True) + 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 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) + if module is not None: + module.sdnext_ondemand = False + if hasattr(module, '_hf_hook'): + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + removed.append(f'{module_name}:ondemand') + if getattr(sd_model, 'sdnext_ondemand_modules', None): + sd_model.sdnext_ondemand_modules = [] + if removed: + log.debug(f'Offload: type=group op=remove modules={removed}') + + +def apply_group_offload_component(module, module_name: str, main: bool, op: str = 'model') -> bool: + """Apply group offload to one component. Re-application with unchanged settings is a no-op: + the hooks silently keep their original config when re-applied and raise before the first + forward, so a changed config must remove the old hooks first.""" + from diffusers.hooks import apply_group_offloading + cfg = group_offload_config(main) + if cfg['use_stream'] and not cfg['low_cpu_mem_usage']: + size_gb, _params = get_module_size(module) + pin_ok = getattr(module, 'sdnext_group_offload_pin', None) + if pin_ok is None: # decide once per module: a granted pin moves the weights into locked memory, so re-reading available on the next apply would see it lower by the pinned size and revoke its own grant + from modules import memstats + avail_gb = memstats.ram_stats().get('avail', 0) + reserve_gb = max(8.0, 0.25 * shared.cpu_memory) # pinned pages cannot be reclaimed or swapped, so a quarter of the machine, floored at 8 GB, stays pageable for the process and page cache + limit_gb = (avail_gb - reserve_gb) if avail_gb > 0 else (0.5 * shared.cpu_memory) # budget from memory free right now; total-derived ceiling only when psutil cannot say + pin_ok = size_gb <= limit_gb + module.sdnext_group_offload_pin = pin_ok + module.sdnext_group_offload_pin_limit = limit_gb + if not pin_ok: + # unpinned streaming degrades to per-transfer staging and leaf groups make that a per-module cost, + # so the whole leaf+stream shape goes with the pin: few large synchronous groups instead + cfg['low_cpu_mem_usage'] = True + cfg['use_stream'] = False + cfg['record_stream'] = False + cfg['offload_type'] = 'block_level' + cfg['num_blocks_per_group'] = max(4, int(shared.opts.group_offload_blocks)) + log.warning(f'Setting {op}: offload=group module={module_name} size={size_gb:.3f} limit={getattr(module, "sdnext_group_offload_pin_limit", 0):.3f} pin=denied type=block_level blocks={cfg["num_blocks_per_group"]} expect ~{size_gb:.0f} GB transferred per step') + sig = f'{devices.device}:{main}:' + ':'.join(str(v) for v in cfg.values()) + if getattr(module, 'sdnext_group_offload_sig', None) == sig: + return False + if hasattr(module, '_hf_hook'): # leftover accelerate hooks from a previous offload mode abort the group apply upstream + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + remove_group_offload_component(module) + module.requires_grad_(False) + apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg) + module.sdnext_group_offload_sig = sig + return True + + +def set_group_resident(module): + """VAE-class components never take group hooks: the hooks are forward-scoped, while + pipelines enter through encode/decode, and tiled calls re-enter per tile.""" + if hasattr(module, '_hf_hook'): + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + remove_group_offload_component(module) + module.requires_grad_(False) + module.to(devices.device) + + +def group_offload_role(module_name: str, module) -> str: + cls = module.__class__.__name__ + if 'vae' in module_name.lower() or cls.startswith(('Autoencoder', 'VQModel', 'AsymmetricAutoencoder', 'ConsistencyDecoder')): + return 'resident' + if module_name.startswith(('text_encoder', 'image_encoder', 'safety_checker')): + return 'aux' + return 'main' + + +def has_entry_bridge(module) -> bool: + """Entry points decorated with diffusers' apply_forward_hook fire _hf_hook.pre_forward, + which is what carries the on-demand onload for encode and decode calls that bypass forward.""" + for name in ('decode', 'encode'): + fn = getattr(module, name, None) + if fn is not None and getattr(fn, '__qualname__', '').startswith('apply_forward_hook'): + return True + return False + + +class OnDemandHook(accelerate.hooks.ModelHook): + """Whole-module onload for components entered through decode or encode rather than forward. + Tiled calls re-enter inside one entry point, so the module is on device before the first + tile; the return to cpu happens at the processing seams once outputs are materialized.""" + def pre_forward(self, module, *args, **kwargs): + param = next(module.parameters(), None) + if param is not None and not devices.same_device(param.device, devices.device): + t0 = time.time() + module.to(devices.device) + dt = time.time() - t0 + process_timer.add('onload', dt) + log.debug(f'Offload: type=ondemand op=onload module={module.__class__.__name__} time={dt:.3f}') + return args, kwargs + + +def set_group_vae(sd_model, module, module_name: str) -> str: + """Placement policy for vae-class components, which never take group hooks. Small + components stay resident; components above group_offload_vae_limit rest on cpu and + onload whole when their decode or encode entry point fires.""" + size_gb, _params = get_module_size(module) + if size_gb < group_offload_vae_limit or not has_entry_bridge(module): + set_group_resident(module) + module.sdnext_ondemand = False # a lingering stamp would let the seams offload a component with no onload hook + names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] + if module_name in names: + sd_model.sdnext_ondemand_modules = [n for n in names if n != module_name] + return 'resident' + if not getattr(module, 'sdnext_ondemand', False) or not hasattr(module, '_hf_hook'): + if hasattr(module, '_hf_hook'): + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) + remove_group_offload_component(module) + module.requires_grad_(False) + accelerate.hooks.add_hook_to_module(module, OnDemandHook(), append=False) + module.sdnext_ondemand = True + module.to(devices.cpu) + names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] + if module_name not in names: + sd_model.sdnext_ondemand_modules = names + [module_name] + return 'ondemand' + + +def offload_ondemand(sd_model): + """Return on-demand components to cpu once their outputs are materialized.""" + if sd_model is None: + return + names = getattr(sd_model, 'sdnext_ondemand_modules', None) + if not names and hasattr(sd_model, 'pipe'): + sd_model = sd_model.pipe + names = getattr(sd_model, 'sdnext_ondemand_modules', None) + for module_name in names or []: + module = getattr(sd_model, module_name, None) + param = next(module.parameters(), None) if module is not None else None + if param is not None and not devices.same_device(param.device, devices.cpu): + t0 = time.time() + module.to(devices.cpu) + dt = time.time() - t0 + process_timer.add('offload', dt) + log.debug(f'Offload: type=ondemand op=offload module={module_name} time={dt:.3f}') + + +def report_group_stats(sd_model, module_names): + """Per-component stats block once per loaded model; balanced mode prints its own from the hook map.""" + if getattr(sd_model, 'sdnext_group_stats_reported', False): + return + sd_model.sdnext_group_stats_reported = True + total = 0.0 + counted = [] + for module_name in module_names: + module = getattr(sd_model, module_name, None) + if isinstance(module, torch.nn.Module): + total += get_module_size(module)[0] + counted.append(module_name) + report_model_stats(module_name, module) + log.info(f'Model class={sd_model.__class__.__name__} modules={len(counted)} size={total:.3f}') + + +def apply_modular_group_offload(sd_model, op:str='model'): + """Per-component group offload for modular pipelines, which lack the pipeline-level + enable_*_offload entry points. The model and sequential modes also route here.""" + if shared.opts.diffusers_offload_mode != 'group' and not getattr(sd_model, 'sdnext_modular_offload_warned', False): + sd_model.sdnext_modular_offload_warned = True + log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} not supported on modular pipelines: using group offload') + applied = [] + for name in ('transformer', 'transformer_ref'): + transformer = getattr(sd_model, name, None) + if transformer is not None and apply_group_offload_component(transformer, name, main=True, op=op): + applied.append(name) + text_encoder = getattr(sd_model, 'text_encoder', None) + if text_encoder is not None: + # offload targets the inner model when present: conditioning may call it directly, + # and hooks on the wrapper forward would never fire + if apply_group_offload_component(getattr(text_encoder, 'model', text_encoder), 'text_encoder', main=False, op=op): + applied.append('text_encoder') + for name in ('vae', 'audio_vae'): + component = getattr(sd_model, name, None) + if component is not None: + placement = set_group_vae(sd_model, component, name) + applied.append(f'{name}:{placement}') + # has_accelerate stays unset: group hooks are not accelerate hooks, and the modular + # pipeline's own to() skips group-offloaded components when move_model runs + if any(':' not in name for name in applied): + log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied}') + report_group_stats(sd_model, ('transformer', 'transformer_ref', 'text_encoder', 'vae', 'audio_vae')) + + +def apply_group_offload(sd_model, op:str='model'): + applied, resident, ondemand = [], [], [] + for module_name in get_module_names(sd_model): + module = getattr(sd_model, module_name, None) + if not isinstance(module, torch.nn.Module): + continue + try: + role = group_offload_role(module_name, module) + if role == 'resident': + if set_group_vae(sd_model, module, module_name) == 'ondemand': + ondemand.append(module_name) + else: + resident.append(module_name) + elif apply_group_offload_component(module, module_name, main=role == 'main', op=op): + applied.append(module_name) + except Exception as e: + log.error(f'Setting {op}: offload=group module={module_name} {e}') set_accelerate(sd_model) + if applied: + log.info(f'Setting {op}: offload=group type={shared.opts.group_offload_type} modules={applied} resident={resident} ondemand={ondemand}') + report_group_stats(sd_model, get_module_names(sd_model)) return sd_model def apply_model_offload(sd_model, op:str='model', quiet:bool=False): try: + remove_group_offload(sd_model) log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False @@ -117,6 +346,7 @@ def apply_model_offload(sd_model, op:str='model', quiet:bool=False): def apply_sequential_offload(sd_model, op:str='model', quiet:bool=False): try: + remove_group_offload(sd_model) log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') if shared.opts.diffusers_move_base or shared.opts.diffusers_move_unet or shared.opts.diffusers_move_refiner: shared.opts.diffusers_move_base = False @@ -144,6 +374,7 @@ def apply_none_offload(sd_model, op:str='model', quiet:bool=False): log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') try: sd_model.has_accelerate = False + remove_group_offload(sd_model) if hasattr(sd_model, 'maybe_free_model_hooks'): sd_model.maybe_free_model_hooks() sd_model = accelerate.hooks.remove_hook_from_module(sd_model, recurse=True) @@ -224,6 +455,14 @@ class OffloadHook(accelerate.hooks.ModelHook): def model_size(self): return sum(self.offload_map.values()) + def matches(self, module, names: list, module_name: str | None = None) -> bool: + """Match against an always/never list by class name or by pipeline component name. + Component entries such as `text_encoder` cover every architecture without listing each encoder class.""" + if module.__class__.__name__ in names: + return True + module_name = module_name or getattr(module, 'module_name', None) + return module_name is not None and module_name in names + def init_hook(self, module): return module @@ -249,8 +488,7 @@ class OffloadHook(accelerate.hooks.ModelHook): 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 (module_instance is not None) and (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)): + if (module_instance is not None) and (_id != id(module_instance)) and (not self.matches(module_instance, self.offload_never, module_name)) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)): apply_balanced_offload_to_module(module_instance, op='pre') self.last_cls = module.__class__.__name__ process_timer.add('offload', time.time() - t0) @@ -440,9 +678,9 @@ def move_module_to_cpu(module, op='unk', force:bool=False): op = f'{op}:force' module = do_move(module) used_gpu -= module_size - elif module_cls in offload_hook_instance.offload_never: + elif offload_hook_instance.matches(module, offload_hook_instance.offload_never, module_name): op = f'{op}:never' - elif module_cls in offload_hook_instance.offload_always: + elif offload_hook_instance.matches(module, offload_hook_instance.offload_always, module_name): op = f'{op}:always' module = do_move(module) used_gpu -= module_size @@ -505,7 +743,9 @@ def get_logical_param_count(module: torch.nn.Module) -> int: def report_model_stats(module_name, module): try: - size = offload_hook_instance.offload_map.get(module_name, 0) + size = offload_hook_instance.offload_map.get(module_name, 0) if offload_hook_instance is not None else 0 + if size == 0: + size, _params = get_module_size(module) quant = getattr(module, "quantization_method", None) params = sum(p.numel() for p in module.parameters(recurse=True)) logical = get_logical_param_count(module) @@ -528,6 +768,7 @@ def apply_balanced_offload(sd_model=None, exclude: list[str] | None = None, forc exclude = [] if sd_model.__class__.__name__ in balanced_offload_exclude: return sd_model + remove_group_offload(sd_model) t0 = time.time() cached = True diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 5944dc426..4aa269a11 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -153,6 +153,7 @@ def create_settings(cmd_opts): "group_offload_type": OptionInfo("leaf_level", "Group offload type", gr.Radio, {"choices": ['leaf_level', 'block_level']}), "group_offload_stream": OptionInfo(False, "Use torch streams", gr.Checkbox), 'group_offload_record': OptionInfo(False, "Record torch streams", gr.Checkbox), + 'group_offload_pin': OptionInfo(True, "Pin offload memory", gr.Checkbox), 'group_offload_blocks': OptionInfo(1, "Offload blocks", gr.Number), })) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 91537762d..c281da8fd 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -620,7 +620,7 @@ {"id":"","label":"Generic","localized":"","hint":"","ui":"video"}, {"id":"","label":"Google GenAI","localized":"","hint":"","ui":"settings_model_options"}, {"id":"","label":"Group Offload","localized":"","hint":"","ui":"settings_offload"}, - {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings, and keeps the VAE resident

Applies only when Model offload mode is group.

Default is leaf_level.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by group offload.
- leaf_level: offloads at the smallest module level; maximum memory savings, slower
- block_level: offloads groups of transformer blocks (size set by Offload blocks); faster with less savings
Text encoders always offload at leaf level. Small VAEs stay resident on the GPU; VAEs above 1GB rest in system memory and load whole for each encode or decode.

Applies only when Model offload mode is group.

Default is leaf_level.","reload":"model","ui":"settings_offload"}, {"id":"","label":"Grid Options","localized":"","hint":"","ui":"settings_saving-images"}, {"id":"","label":"Grids","localized":"","hint":"","ui":"settings_saving-paths"}, {"id":"","label":"Guider","localized":"","hint":"","ui":"txt2img"}, @@ -943,8 +943,8 @@ {"id":"","label":"Model load model direct to GPU","localized":"","hint":"","ui":"settings_sd"}, {"id":"","label":"Model offload mode","localized":"","hint":"Controls how model components move between VRAM and system RAM to fit larger models on less VRAM.
- none: keeps everything on the GPU; fastest, but only works if the whole model fits in VRAM
- balanced: the recommended default; offloads only when VRAM use crosses a threshold, so it suits almost any GPU (tuned by the watermarks below)
- group: offloads groups of layers via diffusers group offloading; an alternative middle ground with optional stream prefetch
- model: offloads whole components such as the VAE or text encoder when idle; a more compatible fallback when balanced or group are unsupported, with smaller savings
- sequential: offloads layer by layer; the most memory saving but slowest, for when even balanced runs out of memory

Command-line flags override the automatic choice:
- --lowvram: forces sequential
- --medvram: forces balanced with low watermark 0

With no flag, balanced is the automatic default on any GPU, with watermarks set by GPU memory (low / high):
- 12 GB or less: 0 / 0.6
- 12-24 GB: 0.2 / 0.6
- 24 GB or more: 0.2 / 0.8
(or none if no GPU is detected)","reload":"model","ui":"settings_offload"}, {"id":"","label":"Model types not to offload","localized":"","hint":"Model architectures to skip when offloading, separated by spaces or commas.
Useful for model types that misbehave when offloaded.

Applies only to balanced offload.

Default is empty.","ui":"settings_offload"}, - {"id":"","label":"Modules to always offload","localized":"","hint":"Module names that are always offloaded in balanced mode, separated by spaces, commas, or semicolons, regardless of the watermarks.

Applies only to balanced offload.

Default by GPU memory: the large text encoders (T5, UMT5) are added at roughly 4-12 GB and at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"}, - {"id":"","label":"Modules to never offload","localized":"","hint":"Module names that are never offloaded in balanced mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.

Applies only to balanced offload.

Default by GPU memory: the CLIP text encoders and the VAE are kept resident at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Modules to always offload","localized":"","hint":"Modules that are always offloaded in balanced mode, separated by spaces, commas, or semicolons, regardless of the watermarks.
Entries match either a class name (T5EncoderModel) or a pipeline component name (text_encoder, text_encoder_2).
A component entry covers every model architecture without naming each encoder class.

Applies only to balanced offload.

Default by GPU memory: the large text encoders (T5, UMT5) are added at roughly 4-12 GB and at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Modules to never offload","localized":"","hint":"Modules that are never offloaded in balanced mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.
Entries match either a class name (CLIPTextModel) or a pipeline component name (vae).
This list takes precedence, so a class entry here exempts one model from a component entry in Modules to always offload.

Applies only to balanced offload.

Default by GPU memory: the CLIP text encoders and the VAE are kept resident at 24 GB or more; empty otherwise.","reload":"model","ui":"settings_offload"}, {"id":"","label":"Model types not to quantize","localized":"","hint":"Model families that quantization always skips, even when it is otherwise enabled. Space or comma separated list of model type codes; when the loaded model matches, none of its components are quantized.

Codes are the short family names shown in the load log, such as sd, sdxl, sd3, f1.

Example: sd sdxl leaves SD and SDXL checkpoints in full precision while other families are still quantized.

Applies to all quantization backends.

Default is empty.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Modules to not convert","localized":"","hint":"Names of modules to leave unquantized (kept in original precision), separated by spaces, commas, or semicolons.
Useful for layers that are sensitive to quantization, such as gate or projection layers. Example: proj_out, x_embedder.

Some models already exclude sensitive modules by default; entries here extend that built-in list rather than replacing it.

Default is empty.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Modules dtype dict","localized":"","hint":"Advanced: JSON mapping a quantization type to a list of module names, to quantize specific modules at a different type than the global Quantization type. Example: {\"uint4\": [\"proj_out\"]}.

Some models already assign certain modules a specific type by default; entries here merge with those built-in mappings rather than replacing them.

Default is empty.","reload":"model","ui":"settings_quantization"}, @@ -1067,6 +1067,7 @@ {"id":"","label":"OpenBody","localized":"","hint":"","ui":"control"} ], "p": [ + {"id":"","label":"Pin offload memory","localized":"","hint":"Keeps the CPU copy of every stream-offloaded weight in pinned non-pageable memory for the fastest transfers, at a host memory cost equal to the full module size.
When disabled, memory is pinned one group at a time during transfer: slower, but the weights stay pageable and use no extra memory at rest.
Modules larger than half of system memory fall back to per-group pinning automatically.

Applies only to group offload with Use torch streams enabled.

Enabled by default.","reload":"model","ui":"settings_offload"}, {"id":"extras_nav","label":"Process","localized":"","hint":"Process existing image
Can be used to upscale images, remove backgrounds, obfuscate NSFW content, apply various filters and effects"}, {"id":"txt2img_prompts","label":"Prompts","localized":"","hint":"Image prompt and negative prompt","ui":"txt2img"}, {"id":"txt2img_pause","label":"Pause","localized":"","hint":"Pause processing","ui":"txt2img"}, @@ -1237,7 +1238,7 @@ {"id":"","label":"Rebase","localized":"","hint":"","ui":"tab_update"}, {"id":"","label":"Repos","localized":"","hint":"","ui":"component-8779"}, {"id":"","label":"Refiner model","localized":"","hint":"Refiner model used for second-pass operations","ui":"settings_sd"}, - {"id":"","label":"Record torch streams","localized":"","hint":"Records CUDA stream usage during group offload so reused buffers stay correct when Use torch streams is enabled.

Applies only to group offload with streams.

Disabled by default.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Record torch streams","localized":"","hint":"Skips a stream synchronization each time a group offloads, letting transfers and compute overlap more tightly.
Slightly faster at the cost of slightly higher VRAM use; correctness is maintained either way.
Has no effect unless Use torch streams is enabled.

Applies only to group offload with streams.

Disabled by default.","reload":"model","ui":"settings_offload"}, {"id":"","label":"Remote VAE image type","localized":"","hint":"","ui":"settings_vae_encoder"}, {"id":"","label":"Remote VAE for encode","localized":"","hint":"","ui":"settings_vae_encoder"}, {"id":"","label":"RAS enabled","localized":"","hint":"","ui":"settings_advanced"}, @@ -1562,7 +1563,7 @@ {"id":"","label":"Username","localized":"","hint":"","ui":"component-8823"}, {"id":"","label":"UNET model","localized":"","hint":"","ui":"settings_sd"}, {"id":"","label":"UNET model secondary","localized":"","hint":"Override for the second transformer of dual-transformer architectures:
- Ideogram 4: the unconditional transformer
- Wan combined stage: the second expert (transformer_2)

Shown only when the loaded model has a second transformer.
Default loads the component from the base model.","ui":"settings_sd"}, - {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but uses more VRAM.

Applies only to group offload.

Disabled by default.","reload":"model","ui":"settings_offload"}, + {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but offloaded weights are then staged in pinned non-pageable host memory: the whole module when Pin offload memory is enabled, one group at a time otherwise.
Text encoders are exempt and always offload without streams.

Applies only to group offload.

Disabled by default.","reload":"model","ui":"settings_offload"}, {"id":"","label":"Use SVD quantization","localized":"","hint":"Adds a low-rank (SVDQuant) correction on top of SDNQ to recover accuracy lost at low bit widths, at the cost of extra size and compute.
Tuned by SVD rank size and SVD steps.

Disabled by default.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Use Dynamic quantization","localized":"","hint":"Picks a per-layer weight type automatically instead of one type everywhere. Each layer starts at the Quantization type (the minimum) and steps up to higher precision until its error meets the Dynamic loss threshold.
Protects error-sensitive layers at the cost of a larger model and slower load.

Disabled by default.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Use Hadamard rotations","localized":"","hint":"Applies a Hadamard rotation before quantizing to spread out weight outliers, which can improve accuracy at low bit widths.
Group size is set by Hadamard group size.

Disabled by default.","reload":"model","ui":"settings_quantization"},