fix(offload): hand back the loaded cpu tensors on no-stream offload

Group hooks on the no-stream path and the on-demand hook return a
component to cpu through a device copy, so a memory-mapped text encoder
sits in memory twice, as the mapped file its never-run vision tower
keeps alive and as the copies, and every encode pays a device-to-host
transfer of unchanged weights. The engine now records the cpu tensors
at onload and hands them back at offload; a component moved by any
other path still takes the copy.

modules/sd_offload_group.py

@@ -1,4 +1,5 @@
 import time
+import types
 import itertools
 import torch
 import accelerate.hooks
@@ -28,6 +29,59 @@ def group_offload_config(main: bool) -> dict:
     }
 
 
+def group_tensors(group) -> list:
+    """Every parameter and buffer a diffusers group moves, read at call time so a tensor replaced while offloaded is seen."""
+    tensors = []
+    for module in group.modules:
+        tensors.extend(module.parameters())
+        tensors.extend(module.buffers())
+    tensors.extend(group.parameters)
+    tensors.extend(group.buffers)
+    return tensors
+
+
+def loaded_tensors(tensors) -> dict:
+    """The cpu tensors held right now, keyed by parameter, to hand back on offload in place of a fresh copy."""
+    return {t: t.data for t in tensors if t.data.device.type == 'cpu'}
+
+
+def restore_tensors(tensors, loaded: dict | None, non_blocking: bool = False):
+    """Return tensors to cpu: the tensor each was onloaded from where known, a copy otherwise."""
+    for t in tensors:
+        if t.data.device.type == 'cpu':
+            continue
+        source = loaded.get(t) if loaded else None
+        t.data = source if source is not None else t.data.to(devices.cpu, non_blocking=non_blocking)
+
+
+def onload_remember(group):
+    group.sdnext_loaded = loaded_tensors(group_tensors(group))
+    group.sdnext_onload()
+
+
+def offload_restore(group):
+    restore_tensors(group_tensors(group), getattr(group, 'sdnext_loaded', None))
+    group.sdnext_loaded = None # a record lives from one onload to its offload
+
+
+def keep_loaded_tensors(module) -> int:
+    """Groups on the no-stream path copy their weights to fresh cpu memory on every offload; record the cpu
+    tensors at onload and hand them back at offload instead. Returns the number of groups patched."""
+    from diffusers.hooks.group_offloading import _GROUP_OFFLOADING
+    count = 0
+    for sub in module.modules():
+        registry = getattr(sub, '_diffusers_hook', None)
+        hook = registry.get_hook(_GROUP_OFFLOADING) if registry is not None else None
+        group = getattr(hook, 'group', None)
+        if group is None or group.stream is not None or getattr(group, 'offload_to_disk_path', None) or hasattr(group, 'sdnext_onload'):
+            continue
+        group.sdnext_onload = group._onload_from_memory # pylint: disable=protected-access
+        group._onload_from_memory = types.MethodType(onload_remember, group) # pylint: disable=protected-access
+        group._offload_to_memory = types.MethodType(offload_restore, group) # pylint: disable=protected-access
+        count += 1
+    return count
+
+
 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
@@ -99,6 +153,8 @@ def apply_group_offload_component(module, module_name: str, main: bool) -> bool:
     s.debug_move(f'Offload: type=group op=apply type={shared.opts.group_offload_type} module={module_name} pin={cfg["use_stream"] and not cfg["low_cpu_mem_usage"]}') # before the apply: pinning large components takes a while and would otherwise run silently
     module.sdnext_group_offload_sig = 'partial' # a raise below leaves hooks that only a non-empty signature will remove
     apply_group_offloading(module, onload_device=devices.device, offload_device=devices.cpu, **cfg)
+    if not cfg['use_stream']:
+        s.debug_move(f'Offload: type=group op=keep module={module_name} groups={keep_loaded_tensors(module)}')
     module.sdnext_group_offload_sig = sig
     return True
 
@@ -153,6 +209,7 @@ class OnDemandHook(accelerate.hooks.ModelHook):
         param = next(module.parameters(), None)
         if param is not None and not devices.same_device(param.device, devices.device):
             t0 = time.time()
+            module.sdnext_loaded = loaded_tensors(list(module.parameters()) + list(module.buffers()))
             module.to(devices.device, non_blocking=shared.opts.diffusers_offload_nonblocking)
             t1 = time.time()
             process_timer.add('onload', t1 - t0)
@@ -194,7 +251,12 @@ def offload_ondemand(sd_model, include=[], exclude=[], reason='', force=False):
                 continue
             try:
                 t0 = time.time()
-                module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking)
+                loaded = getattr(module, 'sdnext_loaded', None)
+                if loaded:
+                    restore_tensors(list(module.parameters()) + list(module.buffers()), loaded, non_blocking=shared.opts.diffusers_offload_nonblocking)
+                    module.sdnext_loaded = None # a record lives from one onload to its offload
+                else:
+                    module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking)
                 dt = time.time() - t0
                 process_timer.add('offload', dt)
                 moved.append(module_name)

test/test-offload-roles.py

@@ -466,6 +466,59 @@ def test_resident_placement_clears_the_ondemand_stamp():
     assert not hasattr(module, '_hf_hook'), 'the on-demand hook must be removed'
 
 
+def accelerator():
+    """The restore round trip needs a device that swaps tensor data with cpu both ways; meta cannot, so these two tests take a real accelerator or skip."""
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    log.warning('  SKIP: no accelerator for the round trip')
+    return None
+
+
+def test_ondemand_offload_hands_back_the_loaded_tensors():
+    device = accelerator()
+    if device is None:
+        return True
+    module = BridgeModule()
+    sd_offload_group.apply_group_offload_ondemand(module)
+    loaded = {name: p.data for name, p in module.named_parameters()}
+    pipe = FakePipe({'vae': module})
+    orig_device = sd_offload_group.devices.device
+    sd_offload_group.devices.device = device
+    try:
+        module._hf_hook.pre_forward(module, torch.zeros(1, 4)) # pylint: disable=protected-access
+        assert next(module.parameters()).device.type == device.type, 'the entry hook must onload the whole module'
+        sd_offload_group.offload_ondemand(pipe, force=True)
+    finally:
+        sd_offload_group.devices.device = orig_device
+    for name, param in module.named_parameters():
+        assert param.data.data_ptr() == loaded[name].data_ptr(), f'{name} came back as a copy rather than the loaded tensor'
+    return True
+
+
+def test_group_offload_hands_back_the_loaded_tensors():
+    from diffusers.hooks.group_offloading import _GROUP_OFFLOADING
+    device = accelerator()
+    if device is None:
+        return True
+    module = PlainModule()
+    orig_device = sd_offload_group.devices.device
+    sd_offload_group.devices.device = device
+    try:
+        assert sd_offload_group.apply_group_offload_component(module, 'text_encoder', main=False) is True
+        loaded = {name: p.data for name, p in module.named_parameters()}
+        group = module.proj._diffusers_hook.get_hook(_GROUP_OFFLOADING).group # pylint: disable=protected-access
+        assert group.stream is None and hasattr(group, 'sdnext_onload'), 'aux components take the no-stream path and must carry the restore patch'
+        group.onload_()
+        assert next(module.parameters()).device.type == device.type, 'onload must still move the group'
+        group.offload_()
+    finally:
+        sd_offload_group.devices.device = orig_device
+    for name, param in module.named_parameters():
+        assert param.data.data_ptr() == loaded[name].data_ptr(), f'{name} came back as a copy rather than the loaded tensor'
+    assert sd_offload_group.keep_loaded_tensors(module) == 0, 'a second pass must not patch the same groups again'
+    return True
+
+
 # ============================================================
 # get_module_names
 # ============================================================
@@ -706,6 +759,8 @@ def run_all():
         test_ondemand_apply_returns_bool_and_is_idempotent,
         test_ondemand_apply_leaves_weights_on_cpu,
         test_resident_placement_clears_the_ondemand_stamp,
+        test_ondemand_offload_hands_back_the_loaded_tensors,
+        test_group_offload_hands_back_the_loaded_tensors,
     ]:
         run_test(cat, fn)