diff --git a/launch.py b/launch.py index e3b0b995e..bb07fd696 100755 --- a/launch.py +++ b/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(): diff --git a/modules/api/server.py b/modules/api/server.py index aee51ae4a..0e8e2b010 100644 --- a/modules/api/server.py +++ b/modules/api/server.py @@ -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()) diff --git a/modules/memstats.py b/modules/memstats.py index 22023adde..63c95db81 100644 --- a/modules/memstats.py +++ b/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()