mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
@@ -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
@@ -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()}')
|
||||
|
||||
Reference in New Issue
Block a user