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

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

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

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,