Merge pull request #5016 from vladmandic/feat/offload-engine

feat(offload): per-component group offload engine and memory observability

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():

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)")

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())

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:

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")

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()

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'):

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

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"):

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

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),
     }))
 

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 <b>group</b> offload.<br>- <b>leaf_level</b>: offloads at the smallest module level; maximum memory savings, slower<br>- <b>block_level</b>: offloads groups of transformer blocks (size set by <b><i>Offload blocks</i></b>); faster with less savings, and keeps the VAE resident<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","reload":"model","ui":"settings_offload"},
+    {"id":"","label":"Group offload type","localized":"","hint":"Granularity used by <b>group</b> offload.<br>- <b>leaf_level</b>: offloads at the smallest module level; maximum memory savings, slower<br>- <b>block_level</b>: offloads groups of transformer blocks (size set by <b><i>Offload blocks</i></b>); faster with less savings<br>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.<br><br>Applies only when <b><i>Model offload mode</i></b> is <b>group</b>.<br><br>Default is <b>leaf_level</b>.","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.<br>- <b>none</b>: keeps everything on the GPU; fastest, but only works if the whole model fits in VRAM<br>- <b>balanced</b>: the recommended default; offloads only when VRAM use crosses a threshold, so it suits almost any GPU (tuned by the watermarks below)<br>- <b>group</b>: offloads groups of layers via diffusers group offloading; an alternative middle ground with optional stream prefetch<br>- <b>model</b>: 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<br>- <b>sequential</b>: offloads layer by layer; the most memory saving but slowest, for when even balanced runs out of memory<br><br>Command-line flags override the automatic choice:<br>- <code>--lowvram</code>: forces <b>sequential</b><br>- <code>--medvram</code>: forces <b>balanced</b> with low watermark <b>0</b><br><br>With no flag, <b>balanced</b> is the automatic default on any GPU, with watermarks set by GPU memory (low / high):<br>- 12 GB or less: <b>0</b> / <b>0.6</b><br>- 12-24 GB: <b>0.2</b> / <b>0.6</b><br>- 24 GB or more: <b>0.2</b> / <b>0.8</b><br>(or <b>none</b> 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.<br>Useful for model types that misbehave when offloaded.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default is empty.","ui":"settings_offload"},
-    {"id":"","label":"Modules to always offload","localized":"","hint":"Module names that are always offloaded in <b>balanced</b> mode, separated by spaces, commas, or semicolons, regardless of the watermarks.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the large text encoders (<i>T5</i>, <i>UMT5</i>) 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 <b>balanced</b> mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.<br><br>Applies only to <b>balanced</b> offload.<br><br>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 <b>balanced</b> mode, separated by spaces, commas, or semicolons, regardless of the watermarks.<br>Entries match either a class name (<i>T5EncoderModel</i>) or a pipeline component name (<i>text_encoder</i>, <i>text_encoder_2</i>).<br>A component entry covers every model architecture without naming each encoder class.<br><br>Applies only to <b>balanced</b> offload.<br><br>Default by GPU memory: the large text encoders (<i>T5</i>, <i>UMT5</i>) 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 <b>balanced</b> mode, separated by spaces, commas, or semicolons, keeping them resident in VRAM.<br>Entries match either a class name (<i>CLIPTextModel</i>) or a pipeline component name (<i>vae</i>).<br>This list takes precedence, so a class entry here exempts one model from a component entry in <b><i>Modules to always offload</i></b>.<br><br>Applies only to <b>balanced</b> offload.<br><br>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.<br><br>Codes are the short family names shown in the load log, such as <code>sd</code>, <code>sdxl</code>, <code>sd3</code>, <code>f1</code>.<br><br>Example: <code>sd sdxl</code> leaves <i>SD</i> and <i>SDXL</i> checkpoints in full precision while other families are still quantized.<br><br>Applies to all quantization backends.<br><br>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.<br>Useful for layers that are sensitive to quantization, such as gate or projection layers. Example: <code>proj_out, x_embedder</code>.<br><br>Some models already exclude sensitive modules by default; entries here extend that built-in list rather than replacing it.<br><br>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 <b><i>Quantization type</i></b>. Example: <code>{\"uint4\": [\"proj_out\"]}</code>.<br><br>Some models already assign certain modules a specific type by default; entries here merge with those built-in mappings rather than replacing them.<br><br>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.<br>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.<br>Modules larger than half of system memory fall back to per-group pinning automatically.<br><br>Applies only to <b>group</b> offload with <b><i>Use torch streams</i></b> enabled.<br><br>Enabled by default.","reload":"model","ui":"settings_offload"},
     {"id":"extras_nav","label":"Process","localized":"","hint":"Process existing image<br>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 <b>group</b> offload so reused buffers stay correct when <b><i>Use torch streams</i></b> is enabled.<br><br>Applies only to <b>group</b> offload with streams.<br><br>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.<br>Slightly faster at the cost of slightly higher VRAM use; correctness is maintained either way.<br>Has no effect unless <b><i>Use torch streams</i></b> is enabled.<br><br>Applies only to <b>group</b> offload with streams.<br><br>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:<br>- <b>Ideogram 4</b>: the unconditional transformer<br>- <b>Wan</b> combined stage: the second expert (transformer_2)<br><br>Shown only when the loaded model has a second transformer.<br>Default loads the component from the base model.","ui":"settings_sd"},
-    {"id":"","label":"Use torch streams","localized":"","hint":"In <b>group</b> offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.<br>Faster, but uses more VRAM.<br><br>Applies only to <b>group</b> offload.<br><br>Disabled by default.","reload":"model","ui":"settings_offload"},
+    {"id":"","label":"Use torch streams","localized":"","hint":"In <b>group</b> offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.<br>Faster, but offloaded weights are then staged in pinned non-pageable host memory: the whole module when <b><i>Pin offload memory</i></b> is enabled, one group at a time otherwise.<br>Text encoders are exempt and always offload without streams.<br><br>Applies only to <b>group</b> offload.<br><br>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.<br>Tuned by <b><i>SVD rank size</i></b> and <b><i>SVD steps</i></b>.<br><br>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 <b><i>Quantization type</i></b> (the minimum) and steps up to higher precision until its error meets the <b><i>Dynamic loss threshold</i></b>.<br>Protects error-sensitive layers at the cost of a larger model and slower load.<br><br>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.<br>Group size is set by <b><i>Hadamard group size</i></b>.<br><br>Disabled by default.","reload":"model","ui":"settings_quantization"},