obj trace

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-01-25 11:41:34 -05:00
parent 6d08bf2558
commit 843011406f
2 changed files with 71 additions and 24 deletions
+47
View File
@@ -1,3 +1,5 @@
import re
import sys
import os
import psutil
import torch
@@ -56,3 +58,48 @@ def ram_stats():
return ram
except Exception:
return { 'used': 0, 'total': 0 }
class Object:
pattern = r"'(.*?)'"
def __init__(self, name, obj):
self.id = id(obj)
self.name = name
self.fn = sys._getframe(2).f_code.co_name
self.size = sys.getsizeof(obj)
self.refcount = sys.getrefcount(obj)
if torch.is_tensor(obj):
self.type = obj.dtype
self.size = obj.element_size() * obj.nelement()
else:
self.type = re.findall(self.pattern, str(type(obj)))[0]
self.size = sys.getsizeof(obj)
def __str__(self):
return f'{self.fn}.{self.name} type={self.type} size={self.size} ref={self.refcount}'
def get_objects(gcl={}, threshold:int=0):
objects = []
seen = []
for name, obj in gcl.items():
if id(obj) in seen:
continue
seen.append(id(obj))
if name == '__name__':
name = obj
elif name.startswith('__'):
continue
try:
o = Object(name, obj)
if o.size >= threshold:
objects.append(o)
except Exception:
pass
objects = sorted(objects, key=lambda x: x.size, reverse=True)
for obj in objects:
shared.log.trace(obj)
return objects
+24 -24
View File
@@ -1050,27 +1050,27 @@ def unload_model_weights(op='model'):
shared.compiled_model_state.compiled_cache.clear()
shared.compiled_model_state.req_cache.clear()
shared.compiled_model_state.partitioned_modules.clear()
if op == 'model' or op == 'dict':
if model_data.sd_model:
if not shared.native:
from modules import sd_hijack
move_model(model_data.sd_model, devices.cpu)
sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
elif not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
disable_offload(model_data.sd_model)
move_model(model_data.sd_model, 'meta')
model_data.sd_model = None
devices.torch_gc(force=True)
shared.log.debug(f'Unload weights {op}: {memory_stats()}')
elif op == 'refiner':
if model_data.sd_refiner:
if not shared.native:
from modules import sd_hijack
move_model(model_data.sd_refiner, devices.cpu)
sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
else:
disable_offload(model_data.sd_refiner)
move_model(model_data.sd_refiner, 'meta')
model_data.sd_refiner = None
devices.torch_gc(force=True)
shared.log.debug(f'Unload weights {op}: {memory_stats()}')
if (op == 'model' or op == 'dict') and model_data.sd_model:
shared.log.debug(f'Current {op}: {memory_stats()}')
if not shared.native:
from modules import sd_hijack
move_model(model_data.sd_model, devices.cpu)
sd_hijack.model_hijack.undo_hijack(model_data.sd_model)
elif not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
disable_offload(model_data.sd_model)
move_model(model_data.sd_model, 'meta')
model_data.sd_model = None
devices.torch_gc(force=True)
shared.log.debug(f'Unload {op}: {memory_stats()} after')
elif (op == 'refiner') and model_data.sd_refiner:
shared.log.debug(f'Current {op}: {memory_stats()}')
if not shared.native:
from modules import sd_hijack
move_model(model_data.sd_refiner, devices.cpu)
sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner)
else:
disable_offload(model_data.sd_refiner)
move_model(model_data.sd_refiner, 'meta')
model_data.sd_refiner = None
devices.torch_gc(force=True)
shared.log.debug(f'Unload {op}: {memory_stats()}')