608 lines
26 KiB
Python
608 lines
26 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import copy
|
|
import os.path
|
|
import random
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torchvision.transforms.functional as TF
|
|
from PIL.Image import Image
|
|
|
|
from scepter.modules.model.network.diffusion.diffusion import GaussianDiffusion
|
|
from scepter.modules.model.network.diffusion.schedules import noise_schedule
|
|
from scepter.modules.model.registry import (BACKBONES, EMBEDDERS, MODELS,
|
|
TOKENIZERS)
|
|
from scepter.modules.model.utils.data_utils import crop_back
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import FS
|
|
from scepter.studio.utils.env import get_available_memory
|
|
|
|
|
|
def get_model(model_tuple):
|
|
assert 'model' in model_tuple
|
|
return model_tuple['model']
|
|
|
|
|
|
class LargenInference():
|
|
'''
|
|
define vae, unet, text-encoder, tuner, refiner components
|
|
support to load the components dynamicly.
|
|
create and load model when run this model at the first time.
|
|
'''
|
|
def __init__(self, logger=None):
|
|
self.logger = logger
|
|
self.loaded_model = {}
|
|
self.loaded_model_name = [
|
|
'diffusion_model', 'first_stage_model', 'cond_stage_model'
|
|
]
|
|
|
|
def init_from_cfg(self, cfg):
|
|
self.name = cfg.NAME
|
|
self.is_default = cfg.get('IS_DEFAULT', False)
|
|
module_paras = self.load_default(cfg.get('DEFAULT_PARAS', None))
|
|
assert cfg.have('MODEL')
|
|
cfg.MODEL = self.redefine_paras(cfg.MODEL)
|
|
self.diffusion = self.load_schedule(cfg.MODEL.SCHEDULE)
|
|
self.diffusion_model = self.infer_model(
|
|
cfg.MODEL.DIFFUSION_MODEL, module_paras.get(
|
|
'DIFFUSION_MODEL',
|
|
None)) if cfg.MODEL.have('DIFFUSION_MODEL') else None
|
|
self.first_stage_model = self.infer_model(
|
|
cfg.MODEL.FIRST_STAGE_MODEL,
|
|
module_paras.get(
|
|
'FIRST_STAGE_MODEL',
|
|
None)) if cfg.MODEL.have('FIRST_STAGE_MODEL') else None
|
|
self.cond_stage_model = self.infer_model(
|
|
cfg.MODEL.COND_STAGE_MODEL,
|
|
module_paras.get(
|
|
'COND_STAGE_MODEL',
|
|
None)) if cfg.MODEL.have('COND_STAGE_MODEL') else None
|
|
self.refiner_cond_model = self.infer_model(
|
|
cfg.MODEL.REFINER_COND_MODEL,
|
|
module_paras.get(
|
|
'REFINER_COND_MODEL',
|
|
None)) if cfg.MODEL.have('REFINER_COND_MODEL') else None
|
|
self.refiner_diffusion_model = self.infer_model(
|
|
cfg.MODEL.REFINER_MODEL, module_paras.get(
|
|
'REFINER_MODEL',
|
|
None)) if cfg.MODEL.have('REFINER_MODEL') else None
|
|
self.tokenizer = TOKENIZERS.build(
|
|
cfg.MODEL.TOKENIZER,
|
|
logger=self.logger) if cfg.MODEL.have('TOKENIZER') else None
|
|
|
|
if self.tokenizer is not None:
|
|
self.cond_stage_model['cfg'].KWARGS = {
|
|
'vocab_size': self.tokenizer.vocab_size
|
|
}
|
|
|
|
def redefine_paras(self, cfg):
|
|
if cfg.get('PRETRAINED_MODEL', None):
|
|
assert FS.isfile(cfg.PRETRAINED_MODEL)
|
|
with FS.get_from(cfg.PRETRAINED_MODEL,
|
|
wait_finish=True) as local_path:
|
|
if local_path.endswith('safetensors'):
|
|
from safetensors.torch import load_file as load_safetensors
|
|
sd = load_safetensors(local_path)
|
|
else:
|
|
sd = torch.load(local_path, map_location='cpu')
|
|
|
|
if 'model' in sd:
|
|
sd = sd['model']
|
|
|
|
first_stage_model_path = os.path.join(
|
|
os.path.dirname(local_path), 'first_stage_model.pth')
|
|
cond_stage_model_path = os.path.join(
|
|
os.path.dirname(local_path), 'cond_stage_model.pth')
|
|
diffusion_model_path = os.path.join(
|
|
os.path.dirname(local_path), 'diffusion_model.pth')
|
|
if (not os.path.exists(first_stage_model_path)
|
|
or not os.path.exists(cond_stage_model_path)
|
|
or not os.path.exists(diffusion_model_path)):
|
|
self.logger.info(
|
|
'Now read the whole model and rearrange the modules, it may take several mins.'
|
|
)
|
|
first_stage_model = OrderedDict()
|
|
cond_stage_model = OrderedDict()
|
|
diffusion_model = OrderedDict()
|
|
for k, v in sd.items():
|
|
if k.startswith('first_stage_model.'):
|
|
first_stage_model[k.replace(
|
|
'first_stage_model.', '')] = v
|
|
elif k.startswith('conditioner.'):
|
|
cond_stage_model[k.replace('conditioner.', '')] = v
|
|
elif k.startswith('cond_stage_model.'):
|
|
if k.startswith('cond_stage_model.model.'):
|
|
cond_stage_model[k.replace(
|
|
'cond_stage_model.model.', '')] = v
|
|
else:
|
|
cond_stage_model[k.replace(
|
|
'cond_stage_model.', '')] = v
|
|
elif k.startswith('model.diffusion_model.'):
|
|
diffusion_model[k.replace('model.diffusion_model.',
|
|
'')] = v
|
|
elif k.startswith('model.'):
|
|
diffusion_model[k.replace('model.', '')] = v
|
|
else:
|
|
continue
|
|
if cfg.have('FIRST_STAGE_MODEL'):
|
|
with open(first_stage_model_path + 'cache', 'wb') as f:
|
|
torch.save(first_stage_model, f)
|
|
os.rename(first_stage_model_path + 'cache',
|
|
first_stage_model_path)
|
|
self.logger.info(
|
|
'First stage model has been processed.')
|
|
if cfg.have('COND_STAGE_MODEL'):
|
|
with open(cond_stage_model_path + 'cache', 'wb') as f:
|
|
torch.save(cond_stage_model, f)
|
|
os.rename(cond_stage_model_path + 'cache',
|
|
cond_stage_model_path)
|
|
self.logger.info(
|
|
'Cond stage model has been processed.')
|
|
if cfg.have('DIFFUSION_MODEL'):
|
|
with open(diffusion_model_path + 'cache', 'wb') as f:
|
|
torch.save(diffusion_model, f)
|
|
os.rename(diffusion_model_path + 'cache',
|
|
diffusion_model_path)
|
|
self.logger.info('Diffusion model has been processed.')
|
|
if not cfg.FIRST_STAGE_MODEL.get('PRETRAINED_MODEL', None):
|
|
cfg.FIRST_STAGE_MODEL.PRETRAINED_MODEL = first_stage_model_path
|
|
else:
|
|
cfg.FIRST_STAGE_MODEL.RELOAD_MODEL = first_stage_model_path
|
|
if not cfg.COND_STAGE_MODEL.get('PRETRAINED_MODEL', None):
|
|
cfg.COND_STAGE_MODEL.PRETRAINED_MODEL = cond_stage_model_path
|
|
else:
|
|
cfg.COND_STAGE_MODEL.RELOAD_MODEL = cond_stage_model_path
|
|
if not cfg.DIFFUSION_MODEL.get('PRETRAINED_MODEL', None):
|
|
cfg.DIFFUSION_MODEL.PRETRAINED_MODEL = diffusion_model_path
|
|
else:
|
|
cfg.DIFFUSION_MODEL.RELOAD_MODEL = diffusion_model_path
|
|
return cfg
|
|
|
|
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')
|
|
|
|
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':
|
|
if module['cfg'].NAME in MODELS.class_map:
|
|
model = MODELS.build(module['cfg'], logger=self.logger).eval()
|
|
elif module['cfg'].NAME in BACKBONES.class_map:
|
|
model = BACKBONES.build(module['cfg'],
|
|
logger=self.logger).eval()
|
|
elif module['cfg'].NAME in EMBEDDERS.class_map:
|
|
model = EMBEDDERS.build(module['cfg'],
|
|
logger=self.logger).eval()
|
|
else:
|
|
raise NotImplementedError
|
|
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:
|
|
module['model'] = module['model'].to('cpu')
|
|
module['device'] = 'cpu'
|
|
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 = {k.lower(): 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_schedule(self, cfg):
|
|
parameterization = cfg.get('PARAMETERIZATION', 'eps')
|
|
assert parameterization in [
|
|
'eps', 'x0', 'v'
|
|
], 'currently only supporting "eps" and "x0" and "v"'
|
|
num_timesteps = cfg.get('TIMESTEPS', 1000)
|
|
|
|
schedule_args = {
|
|
k.lower(): v
|
|
for k, v in cfg.get('SCHEDULE_ARGS', {
|
|
'NAME': 'logsnr_cosine_interp',
|
|
'SCALE_MIN': 2.0,
|
|
'SCALE_MAX': 4.0
|
|
}).items()
|
|
}
|
|
|
|
zero_terminal_snr = cfg.get('ZERO_TERMINAL_SNR', False)
|
|
if zero_terminal_snr:
|
|
assert parameterization == 'v', 'Now zero_terminal_snr only support v-prediction mode.'
|
|
sigmas = noise_schedule(schedule=schedule_args.pop('name'),
|
|
n=num_timesteps,
|
|
zero_terminal_snr=zero_terminal_snr,
|
|
**schedule_args)
|
|
diffusion = GaussianDiffusion(sigmas=sigmas,
|
|
prediction_type=parameterization)
|
|
return diffusion
|
|
|
|
def get_batch(self, value_dict, num_samples=1):
|
|
batch = {}
|
|
batch_uc = {}
|
|
N = num_samples
|
|
device = we.device_id
|
|
for key in value_dict:
|
|
if key == 'prompt':
|
|
if not self.tokenizer:
|
|
batch['prompt'] = value_dict['prompt']
|
|
batch_uc['prompt'] = value_dict['negative_prompt']
|
|
else:
|
|
batch['tokens'] = self.tokenizer(value_dict['prompt']).to(
|
|
we.device_id)
|
|
batch_uc['tokens'] = self.tokenizer(
|
|
value_dict['negative_prompt']).to(we.device_id)
|
|
elif key == 'original_size_as_tuple':
|
|
batch['original_size_as_tuple'] = (torch.tensor(
|
|
value_dict['original_size_as_tuple']).to(device).repeat(
|
|
N, 1))
|
|
elif key == 'crop_coords_top_left':
|
|
batch['crop_coords_top_left'] = (torch.tensor(
|
|
value_dict['crop_coords_top_left']).to(device).repeat(
|
|
N, 1))
|
|
elif key == 'aesthetic_score':
|
|
batch['aesthetic_score'] = (torch.tensor(
|
|
[value_dict['aesthetic_score']]).to(device).repeat(N, 1))
|
|
batch_uc['aesthetic_score'] = (torch.tensor([
|
|
value_dict['negative_aesthetic_score']
|
|
]).to(device).repeat(N, 1))
|
|
|
|
elif key == 'target_size_as_tuple':
|
|
batch['target_size_as_tuple'] = (torch.tensor(
|
|
value_dict['target_size_as_tuple']).to(device).repeat(
|
|
N, 1))
|
|
elif key == 'image':
|
|
batch[key] = self.load_image(value_dict[key], num_samples=N)
|
|
else:
|
|
batch[key] = value_dict[key]
|
|
|
|
for key in batch.keys():
|
|
if key not in batch_uc and isinstance(batch[key], torch.Tensor):
|
|
batch_uc[key] = torch.clone(batch[key])
|
|
return batch, batch_uc
|
|
|
|
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']
|
|
|
|
def encode_first_stage(self, x, **kwargs):
|
|
_, dtype = self.get_function_info(self.first_stage_model, 'encode')
|
|
with torch.autocast('cuda',
|
|
enabled=dtype == 'float16',
|
|
dtype=getattr(torch, dtype)):
|
|
z = get_model(self.first_stage_model).encode(x)
|
|
return self.first_stage_model['paras']['scale_factor'] * z
|
|
|
|
def decode_first_stage(self, z):
|
|
_, dtype = self.get_function_info(self.first_stage_model, 'encode')
|
|
with torch.autocast('cuda',
|
|
enabled=dtype == 'float16',
|
|
dtype=getattr(torch, dtype)):
|
|
z = 1. / self.first_stage_model['paras']['scale_factor'] * z
|
|
return get_model(self.first_stage_model).decode(z)
|
|
|
|
@torch.no_grad()
|
|
def __call__(self,
|
|
input,
|
|
num_samples=1,
|
|
intermediate_callback=None,
|
|
refine_strength=0,
|
|
img_to_img_strength=0,
|
|
cat_uc=True,
|
|
tuner_model=None,
|
|
control_model=None,
|
|
**kwargs):
|
|
|
|
value_input = copy.deepcopy(self.input)
|
|
value_input.update(input)
|
|
print(value_input)
|
|
height, width = value_input['target_size_as_tuple']
|
|
value_output = copy.deepcopy(self.output)
|
|
batch, batch_uc = self.get_batch(value_input, num_samples=1)
|
|
|
|
# first stage encode
|
|
task = kwargs.get('largen_task', 'Text_Guided_Inpainting')
|
|
image_scale = kwargs.get('largen_image_scale', 1.0)
|
|
tar_image = kwargs.get('largen_tar_image', None)
|
|
tar_mask = kwargs.get('largen_tar_mask', None)
|
|
masked_image = kwargs.get('largen_masked_image', None)
|
|
ref_image = kwargs.get('largen_ref_image', None)
|
|
ref_mask = kwargs.get('largen_ref_mask', None)
|
|
ref_clip = kwargs.get('largen_ref_clip', None)
|
|
|
|
base_image = kwargs.get('largen_base_image', None)
|
|
extra_sizes = kwargs.get('largen_extra_sizes', None)
|
|
bbox_yyxx = kwargs.get('largen_bbox_yyxx', None)
|
|
|
|
device = we.device_id
|
|
tar_image = tar_image.to(device)
|
|
tar_mask = tar_mask.to(device)
|
|
masked_image = masked_image.to(device)
|
|
if 'Subject' in task:
|
|
ref_image = ref_image.to(device)
|
|
ref_mask = ref_mask.to(device)
|
|
ref_clip = ref_clip.to(device)
|
|
|
|
self.dynamic_load(self.first_stage_model, 'first_stage_model')
|
|
|
|
tar_x0 = self.encode_first_stage(tar_image)
|
|
masked_x0 = self.encode_first_stage(masked_image)
|
|
b, _, h, w = tar_x0.shape
|
|
tar_mask_latent = TF.resize(tar_mask, (h, w), antialias=True)
|
|
tar_mask_latent = (tar_mask_latent > 0.5).float()
|
|
|
|
batch.update({
|
|
'tar_x0': tar_x0,
|
|
'tar_mask_latent': tar_mask_latent,
|
|
'masked_x0': masked_x0,
|
|
'task': task
|
|
})
|
|
batch_uc.update({
|
|
'tar_x0': tar_x0,
|
|
'tar_mask_latent': tar_mask_latent,
|
|
'masked_x0': masked_x0,
|
|
'task': task
|
|
})
|
|
|
|
if 'Subject' in task and ref_image is not None:
|
|
ref_x0 = self.encode_first_stage(ref_image)
|
|
batch.update({
|
|
'ref_ip': ref_clip,
|
|
'ref_detail': ref_clip,
|
|
'ref_x0': ref_x0,
|
|
'ref_mask': ref_mask,
|
|
'image_scale': image_scale,
|
|
})
|
|
batch_uc.update({
|
|
'ref_ip': torch.zeros_like(ref_clip),
|
|
'ref_detail': ref_clip,
|
|
'ref_x0': ref_x0,
|
|
'ref_mask': ref_mask,
|
|
'image_scale': image_scale,
|
|
})
|
|
|
|
self.dynamic_unload(self.first_stage_model,
|
|
'first_stage_model',
|
|
skip_loaded=False)
|
|
|
|
# cond stage
|
|
self.dynamic_load(self.cond_stage_model, 'cond_stage_model')
|
|
function_name, dtype = self.get_function_info(self.cond_stage_model)
|
|
with torch.autocast('cuda',
|
|
enabled=dtype == 'float16',
|
|
dtype=getattr(torch, dtype)):
|
|
context = getattr(get_model(self.cond_stage_model),
|
|
function_name)(batch)
|
|
null_context = getattr(get_model(self.cond_stage_model),
|
|
function_name)(batch_uc)
|
|
|
|
self.dynamic_unload(self.cond_stage_model,
|
|
'cond_stage_model',
|
|
skip_loaded=False)
|
|
|
|
# get noise
|
|
seed = kwargs.pop('seed', -1)
|
|
g = torch.Generator(device=we.device_id)
|
|
seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
|
|
g.manual_seed(seed)
|
|
if 'seed' in value_output:
|
|
value_output['seed'] = seed
|
|
for sample_id in range(num_samples):
|
|
if self.diffusion_model is not None:
|
|
noise = torch.empty(
|
|
1,
|
|
4,
|
|
height // self.first_stage_model['paras']['size_factor'],
|
|
width // self.first_stage_model['paras']['size_factor'],
|
|
device=we.device_id).normal_(generator=g)
|
|
|
|
self.dynamic_load(self.diffusion_model, 'diffusion_model')
|
|
# UNet use input n_prompt
|
|
function_name, dtype = self.get_function_info(
|
|
self.diffusion_model)
|
|
with torch.autocast('cuda',
|
|
enabled=dtype == 'float16',
|
|
dtype=getattr(torch, dtype)):
|
|
latent = self.diffusion.sample(
|
|
noise=noise,
|
|
x=None,
|
|
denoising_strength=1.0,
|
|
refine_strength=refine_strength,
|
|
solver=value_input.get('sample', 'ddim'),
|
|
model=get_model(self.diffusion_model),
|
|
model_kwargs=[{
|
|
'cond': context
|
|
}, {
|
|
'cond': null_context
|
|
}],
|
|
steps=value_input.get('sample_steps', 50),
|
|
guide_scale=value_input.get('guide_scale', 7.5),
|
|
guide_rescale=value_input.get('guide_rescale', 0.5),
|
|
discretization=value_input.get('discretization',
|
|
'trailing'),
|
|
show_progress=True,
|
|
seed=seed,
|
|
condition_fn=None,
|
|
clamp=None,
|
|
percentile=None,
|
|
t_max=None,
|
|
t_min=None,
|
|
discard_penultimate_step=None,
|
|
intermediate_callback=intermediate_callback,
|
|
cat_uc=cat_uc,
|
|
**kwargs)
|
|
|
|
self.dynamic_unload(self.diffusion_model,
|
|
'diffusion_model',
|
|
skip_loaded=True)
|
|
|
|
if 'latent' in value_output:
|
|
if value_output['latent'] is None or (
|
|
isinstance(value_output['latent'], list)
|
|
and len(value_output['latent']) < 1):
|
|
value_output['latent'] = []
|
|
value_output['latent'].append(latent)
|
|
|
|
self.dynamic_load(self.first_stage_model, 'first_stage_model')
|
|
x_samples = self.decode_first_stage(latent).float()
|
|
self.dynamic_unload(self.first_stage_model,
|
|
'first_stage_model',
|
|
skip_loaded=False)
|
|
images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
|
|
if base_image is not None:
|
|
stitch_images = []
|
|
for img in images:
|
|
stitch_img = crop_back(img, copy.deepcopy(base_image),
|
|
extra_sizes, bbox_yyxx)
|
|
stitch_images.append(stitch_img)
|
|
images = torch.stack(stitch_images, dim=0)
|
|
if 'images' in value_output:
|
|
if value_output['images'] is None or (
|
|
isinstance(value_output['images'], list)
|
|
and len(value_output['images']) < 1):
|
|
value_output['images'] = []
|
|
value_output['images'].append(images)
|
|
|
|
for k, v in value_output.items():
|
|
if isinstance(v, list):
|
|
value_output[k] = torch.cat(v, dim=0)
|
|
if isinstance(v, torch.Tensor):
|
|
value_output[k] = v.cpu()
|
|
|
|
return value_output
|