diff --git a/modules/memstats.py b/modules/memstats.py index d43e5bbfa..cb43f246c 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -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 diff --git a/modules/sd_models.py b/modules/sd_models.py index df49b4918..2de127347 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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()}')