# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import torch
from PIL.Image import Image
from collections import OrderedDict
from scepter.modules.utils.distribute import we
from scepter.modules.utils.config import Config
from scepter.modules.utils.logger import get_logger
from scepter.studio.utils.env import get_available_memory
from scepter.modules.model.registry import MODELS, BACKBONES, EMBEDDERS
from scepter.modules.utils.registry import Registry, build_from_config
def get_model(model_tuple):
assert 'model' in model_tuple
return model_tuple['model']
class BaseInference():
'''
support to load the components dynamicly.
create and load model when run this model at the first time.
'''
def __init__(self, cfg, logger=None):
if logger is None:
logger = get_logger(name='scepter')
self.logger = logger
self.name = cfg.NAME
def init_from_modules(self, modules):
for k, v in modules.items():
self.__setattr__(k, v)
def infer_model(self, cfg, module_paras=None):
module = {
'model': None,
'cfg': cfg,
'device': 'offline',
'name': cfg.NAME,
'function_info': {},
'paras': {}
}
if module_paras is None:
return module
function_info = {}
paras = {
k.lower(): v
for k, v in module_paras.get('PARAS', {}).items()
}
for function in module_paras.get('FUNCTION', []):
input_dict = {}
for inp in function.get('INPUT', []):
if inp.lower() in self.input:
input_dict[inp.lower()] = self.input[inp.lower()]
function_info[function.NAME] = {
'dtype': function.get('DTYPE', 'float32'),
'input': input_dict
}
module['paras'] = paras
module['function_info'] = function_info
return module
def init_from_ckpt(self, path, model, ignore_keys=list()):
if path.endswith('safetensors'):
from safetensors.torch import load_file as load_safetensors
sd = load_safetensors(path)
else:
sd = torch.load(path, map_location='cpu', weights_only=True)
new_sd = OrderedDict()
for k, v in sd.items():
ignored = False
for ik in ignore_keys:
if ik in k:
if we.rank == 0:
self.logger.info(
'Ignore key {} from state_dict.'.format(k))
ignored = True
break
if not ignored:
new_sd[k] = v
missing, unexpected = model.load_state_dict(new_sd, strict=False)
if we.rank == 0:
self.logger.info(
f'Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys'
)
if len(missing) > 0:
self.logger.info(f'Missing Keys:\n {missing}')
if len(unexpected) > 0:
self.logger.info(f'\nUnexpected Keys:\n {unexpected}')
def load(self, module):
if module['device'] == 'offline':
from scepter.modules.utils.import_utils import LazyImportModule
if (LazyImportModule.get_module_type(('MODELS', module['cfg'].NAME)) or
module['cfg'].NAME in MODELS.class_map):
model = MODELS.build(module['cfg'], logger=self.logger).eval()
elif (LazyImportModule.get_module_type(('BACKBONES', module['cfg'].NAME)) or
module['cfg'].NAME in BACKBONES.class_map):
model = BACKBONES.build(module['cfg'],
logger=self.logger).eval()
elif (LazyImportModule.get_module_type(('EMBEDDERS', module['cfg'].NAME)) or
module['cfg'].NAME in EMBEDDERS.class_map):
model = EMBEDDERS.build(module['cfg'],
logger=self.logger).eval()
else:
raise NotImplementedError
if 'DTYPE' in module['cfg'] and module['cfg']['DTYPE'] is not None:
model = model.to(getattr(torch, module['cfg'].DTYPE))
if module['cfg'].get('RELOAD_MODEL', None):
self.init_from_ckpt(module['cfg'].RELOAD_MODEL, model)
module['model'] = model
module['device'] = 'cpu'
if module['device'] == 'cpu':
module['device'] = we.device_id
module['model'] = module['model'].to(we.device_id)
return module
def unload(self, module):
if module is None:
return module
mem = get_available_memory()
free_mem = int(mem['available'] / (1024**2))
total_mem = int(mem['total'] / (1024**2))
if free_mem < 0.5 * total_mem:
if module['model'] is not None:
module['model'] = module['model'].to('cpu')
del module['model']
module['model'] = None
module['device'] = 'offline'
print('delete module')
else:
if module['model'] is not None:
module['model'] = module['model'].to('cpu')
module['device'] = 'cpu'
else:
module['device'] = 'offline'
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
return module
def dynamic_load(self, module=None, name=''):
self.logger.info('Loading {} model'.format(name))
if name == 'all':
for subname in self.loaded_model_name:
self.loaded_model[subname] = self.dynamic_load(
getattr(self, subname), subname)
elif name in self.loaded_model_name:
if name in self.loaded_model:
if module['cfg'] != self.loaded_model[name]['cfg']:
self.unload(self.loaded_model[name])
module = self.load(module)
self.loaded_model[name] = module
return module
elif module['device'] == 'cpu' or module['device'] == 'offline':
module = self.load(module)
return module
else:
return module
else:
module = self.load(module)
self.loaded_model[name] = module
return module
else:
return self.load(module)
def dynamic_unload(self, module=None, name='', skip_loaded=False):
self.logger.info('Unloading {} model'.format(name))
if name == 'all':
for name, module in self.loaded_model.items():
module = self.unload(self.loaded_model[name])
self.loaded_model[name] = module
elif name in self.loaded_model_name:
if name in self.loaded_model:
if not skip_loaded:
module = self.unload(self.loaded_model[name])
self.loaded_model[name] = module
else:
self.unload(module)
else:
self.unload(module)
def load_default(self, cfg):
module_paras = {}
if cfg is not None:
self.paras = cfg.PARAS
self.input_cfg = {k.lower(): v for k, v in cfg.INPUT.items()}
self.input = {k.lower(): dict(v).get('DEFAULT', None) if isinstance(v, (dict, OrderedDict, Config)) else v for k, v in cfg.INPUT.items()}
self.output = {k.lower(): v for k, v in cfg.OUTPUT.items()}
module_paras = cfg.MODULES_PARAS
return module_paras
def load_image(self, image, num_samples=1):
if isinstance(image, torch.Tensor):
pass
elif isinstance(image, Image):
pass
elif isinstance(image, Image):
pass
def get_function_info(self, module, function_name=None):
all_function = module['function_info']
if function_name in all_function:
return function_name, all_function[function_name]['dtype']
if function_name is None and len(all_function) == 1:
for k, v in all_function.items():
return k, v['dtype']
@torch.no_grad()
def __call__(self,
input,
**kwargs):
return
def build_inference(cfg, registry, logger=None, *args, **kwargs):
""" After build model, load pretrained model if exists key `pretrain`.
pretrain (str, dict): Describes how to load pretrained model.
str, treat pretrain as model path;
dict: should contains key `path`, and other parameters token by function load_pretrained();
"""
if not isinstance(cfg, Config):
raise TypeError(f'Config must be type dict, got {type(cfg)}')
model = build_from_config(cfg, registry, logger=logger, *args, **kwargs)
return model
# reigister cls for diffusion.
INFERENCES = Registry('INFERENCE', build_func=build_inference)