From 619a25eea689a2fecd176e346130fc2389af8e74 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 7 Sep 2026 22:35:38 +0100 Subject: [PATCH] fix(offload): key the group stats report by component A checkpoint-name key survives unload, so a reload of the same checkpoint never printed the per-module stats block again, even with a different quantization. Each component now carries its own reported stamp: a task switch rebuilds the pipe around the same modules and stays quiet, while a reload or a component swap brings new modules and reports them. --- modules/sd_offload_group.py | 23 ++++++++++------------- modules/sd_offload_state.py | 1 - test/test-offload-roles.py | 23 +++++++++++++++++++++++ 3 files changed, 33 insertions(+), 14 deletions(-) diff --git a/modules/sd_offload_group.py b/modules/sd_offload_group.py index 07776596b..014456be7 100644 --- a/modules/sd_offload_group.py +++ b/modules/sd_offload_group.py @@ -206,20 +206,17 @@ def offload_ondemand(sd_model, include=[], exclude=[], reason='', force=False): 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.""" - checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else sd_model.__class__.__name__ - if checkpoint_name in s.group_stats_reported: # keyed by checkpoint since a task switch rebuilds the pipe object + """Per-component stats block once per loaded component; balanced mode prints its own from the hook map.""" + modules = {name: getattr(sd_model, name, None) for name in module_names} + modules = {name: module for name, module in modules.items() if isinstance(module, torch.nn.Module)} + pending = {name: module for name, module in modules.items() if not getattr(module, 'sdnext_stats_reported', False)} # a task switch reuses the modules, a reload brings new ones + if not pending: return - s.group_stats_reported.add(checkpoint_name) - 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}') + for module_name, module in pending.items(): + module.sdnext_stats_reported = True + report_model_stats(module_name, module) + total = sum(get_module_size(module)[0] for module in modules.values()) + log.info(f'Model class={sd_model.__class__.__name__} modules={len(modules)} size={total:.3f}') def apply_group_offload(sd_model): diff --git a/modules/sd_offload_state.py b/modules/sd_offload_state.py index 2c1297675..375340a02 100644 --- a/modules/sd_offload_state.py +++ b/modules/sd_offload_state.py @@ -30,5 +30,4 @@ no_split_module_classes = [ ] accelerate_dtype_byte_size = None # monkey-patch accelerate.utils.modeling.dtype_byte_size -group_stats_reported = set() move_stream = None diff --git a/test/test-offload-roles.py b/test/test-offload-roles.py index 8931caf53..533e4a8d7 100644 --- a/test/test-offload-roles.py +++ b/test/test-offload-roles.py @@ -364,6 +364,28 @@ def test_dispatch_skips_non_modules(): assert list(calls) == ['transformer'], f'dispatched {list(calls)}' +def test_stats_report_once_per_component(): + seen = [] + orig_stats = sd_offload_group.report_model_stats + sd_offload_group.report_model_stats = lambda module_name, module: seen.append(module_name) + try: + transformer, vae = PlainModule(), BridgeModule() + pipe = FakePipe({'transformer': transformer, 'vae': vae}) + names = sd_offload_utils.get_module_names(pipe) + sd_offload_group.report_group_stats(pipe, names) + assert sorted(seen) == ['transformer', 'vae'], f'first report covered {seen}' + sd_offload_group.report_group_stats(pipe, names) + assert len(seen) == 2, f'a reapply reported again: {seen}' + switched = FakePipe({'transformer': transformer, 'vae': vae}) # a task switch rebuilds the pipe around the same components + sd_offload_group.report_group_stats(switched, names) + assert len(seen) == 2, f'a task switch reported again: {seen}' + reloaded = FakePipe({'transformer': PlainModule(), 'vae': vae}) # a reload or a component swap brings a new module + sd_offload_group.report_group_stats(reloaded, names) + assert seen[2:] == ['transformer'], f'a new component was not reported on its own: {seen}' + finally: + sd_offload_group.report_model_stats = orig_stats + + def test_force_sweep_moves_only_stamped_components(): stamped = SweepModule() stamped.sdnext_ondemand = True @@ -671,6 +693,7 @@ def run_all(): for fn in [ test_dispatch_is_one_arm_per_component, test_dispatch_skips_non_modules, + test_stats_report_once_per_component, test_ondemand_list_tracks_the_stamps, test_force_sweep_moves_only_stamped_components, test_reapply_after_clearing_the_never_list_restores_hooks,