48 lines
1.9 KiB
Python
48 lines
1.9 KiB
Python
import logging
|
|
import torch
|
|
from comfy import model_management
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
class MemSafeWrapper(ModelPatcher):
|
|
def __init__(self, model:torch.nn.Module, device:torch.device=None, offload_device:torch.device=None, dtype:torch.dtype=None):
|
|
if dtype is None:
|
|
device = model_management.get_torch_device()
|
|
if offload_device is None:
|
|
offload_device = model_management.intermediate_device()
|
|
if dtype is None:
|
|
dtype = model_management.unet_dtype()
|
|
|
|
self.model = model.to(device=offload_device, dtype=dtype)
|
|
self.model.eval()
|
|
super().__init__(self.model, device, offload_device)
|
|
|
|
overlap_attr = set(dir(super())) & set(dir(self.model)) - set(dir(object))
|
|
if len(overlap_attr) != 0:
|
|
attr_list = ", ".join(sorted(overlap_attr))
|
|
logging.warning(
|
|
f"WARNING[MemSafeWrapper]: The following attributes are present in both "
|
|
f"'{self.__class__.__name__}' and '{self.model.__class__.__name__}':\n"
|
|
f"{attr_list}\n"
|
|
f"Attributes from '{self.model.__class__.__name__}' will be ignored for these."
|
|
)
|
|
|
|
def to_gpu(self, func):
|
|
def wrapper(*args, **kwargs):
|
|
if self.current_device == self.offload_device:
|
|
model_management.load_model_gpu(self)
|
|
return func(*args, **kwargs)
|
|
return wrapper
|
|
|
|
def __getattr__(self, name:str):
|
|
if hasattr(self.model, name):
|
|
attr = getattr(self.model, name)
|
|
if callable(attr):
|
|
return self.to_gpu(attr)
|
|
else:
|
|
return attr
|
|
else:
|
|
raise AttributeError(f"{self.__class__.__name__} does not has {name} attribute.")
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self.model(*args, **kwargs)
|