feat(monitor): report model placement in the memory monitor

Move the per-component placement walk from the api server into memstats
next to ram and gpu stats, and add it to the --monitor tick so offload
placement is visible in a log rather than only over the api.

- guard the walk internally: the supervisor loop logs its monitor line
  unguarded, and the walk can race a reload or an offload rewrap
- monitor reports gb and merges into a fresh dict, since memory_stats
  returns a module global that the per-generation log also prints
- endpoint keeps reporting raw bytes
This commit is contained in:
CalamitousFelicitousness
2026-07-21 23:55:57 +01:00
parent b12cbcb523
commit d38f2c97b2
3 changed files with 33 additions and 34 deletions
+2 -31
View File
@@ -165,32 +165,6 @@ def post_skip():
shared.state.skip()
return Response(status_code=204)
def get_model_placement():
import torch
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:
continue
seen.add(ptr)
devmap[t.device.type] = devmap.get(t.device.type, 0) + t.numel() * t.element_size()
if devmap:
placement[name] = devmap
return placement
def get_memory():
try:
import psutil
@@ -223,8 +197,5 @@ def get_memory():
cuda = { 'error': 'unavailable' }
except Exception as err:
cuda = { 'error': f'{err}' }
try:
model = get_model_placement() # walk can race a model reload or offload rewrap; report the error and keep the endpoint alive
except Exception as err:
model = { 'error': f'{err}' }
return models.ResMemory(ram = ram, cuda = cuda, model = model)
from modules import memstats
return models.ResMemory(ram = ram, cuda = cuda, model = memstats.model_stats())