update v0.0.4
This commit is contained in:
LouieStark
@@ -244,12 +244,18 @@ class Text2ImageDataset(BaseDataset):
|
||||
image_size = [image_size, image_size]
|
||||
assert isinstance(image_size, Iterable) and len(image_size) == 2
|
||||
|
||||
prompt_file = cfg.PROMPT_FILE
|
||||
with FS.get_object(prompt_file) as local_data:
|
||||
if cfg.PROMPT_FILE is not None and cfg.PROMPT_FILE != '':
|
||||
prompt_file = cfg.PROMPT_FILE
|
||||
with FS.get_object(prompt_file) as local_data:
|
||||
rows = [
|
||||
i.split(delimiter,
|
||||
len(fields) - 1)
|
||||
for i in local_data.decode('utf-8').strip().split('\n')
|
||||
]
|
||||
else:
|
||||
rows = [
|
||||
i.split(delimiter,
|
||||
len(fields) - 1)
|
||||
for i in local_data.decode('utf-8').strip().split('\n')
|
||||
len(fields) - 1) for i in cfg.PROMPT_DATA
|
||||
]
|
||||
|
||||
self.items = list()
|
||||
@@ -263,10 +269,9 @@ class Text2ImageDataset(BaseDataset):
|
||||
item['meta']['img_path'] = os.path.join(path_prefix, value)
|
||||
elif key in ['width', 'height']:
|
||||
item['meta'][key] = int(value)
|
||||
elif key != 'meta':
|
||||
item[key] = value
|
||||
else:
|
||||
continue
|
||||
item['meta'][key] = value
|
||||
|
||||
self.items.append(item)
|
||||
if use_num > 0:
|
||||
self.items = self.items[:use_num]
|
||||
|
||||
@@ -2,18 +2,23 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import copy
|
||||
import os
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as TT
|
||||
from PIL.Image import Image
|
||||
from swift import SwiftModel
|
||||
|
||||
from scepter.modules.model.registry import TUNERS
|
||||
from scepter.modules.utils.config import Config
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
try:
|
||||
from swift import SwiftModel
|
||||
except Exception:
|
||||
warnings.warn('Import swift failed, please check it.')
|
||||
|
||||
|
||||
class ControlInference():
|
||||
def __init__(self, logger=None):
|
||||
|
||||
@@ -474,10 +474,14 @@ class DiffusionInference():
|
||||
enabled=dtype == 'float16',
|
||||
dtype=getattr(torch, dtype)):
|
||||
if self.tokenizer:
|
||||
if not hasattr(get_model(self.cond_stage_model), 'tokenizer'):
|
||||
setattr(get_model(self.cond_stage_model), 'tokenizer',
|
||||
self.tokenizer)
|
||||
context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch['tokens'])
|
||||
null_context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch_uc['tokens'])
|
||||
|
||||
else:
|
||||
context = getattr(get_model(self.cond_stage_model),
|
||||
function_name)(batch)
|
||||
@@ -558,12 +562,13 @@ class DiffusionInference():
|
||||
seed=seed,
|
||||
condition_fn=None,
|
||||
clamp=None,
|
||||
sharpness=value_input.get('sharpness', 0.0),
|
||||
percentile=None,
|
||||
t_max=None,
|
||||
t_min=None,
|
||||
discard_penultimate_step=None,
|
||||
intermediate_callback=intermediate_callback,
|
||||
cat_uc=cat_uc,
|
||||
cat_uc=value_input.get('cat_uc', cat_uc),
|
||||
**kwargs)
|
||||
|
||||
self.dynamic_unload(self.diffusion_model,
|
||||
|
||||
@@ -0,0 +1,595 @@
|
||||
# -*- 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
|
||||
|
||||
|
||||
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
|
||||
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':
|
||||
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=True)
|
||||
|
||||
# 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=True)
|
||||
|
||||
# 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=True)
|
||||
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
|
||||
@@ -1,4 +1,5 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.backbone.autoencoder.ae_module import (Decoder,
|
||||
Encoder)
|
||||
Encoder,
|
||||
RDecoder)
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import repeat
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
|
||||
from scepter.modules.model.backbone.autoencoder.ae_utils import (
|
||||
@@ -244,6 +245,7 @@ class Decoder(BaseModel):
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = self.ch * self.ch_mult[self.num_resolutions - 1]
|
||||
self.block_in = block_in
|
||||
curr_res = 1
|
||||
# z to block_in
|
||||
self.conv_in = torch.nn.Conv2d(self.z_channels,
|
||||
@@ -340,3 +342,48 @@ class Decoder(BaseModel):
|
||||
__class__.__name__,
|
||||
Decoder.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@BACKBONES.register_class()
|
||||
class RDecoder(Decoder):
|
||||
def construct_model(self):
|
||||
super().construct_model()
|
||||
self.resize_level = nn.Sequential(
|
||||
nn.Linear(self.block_in, self.block_in),
|
||||
nn.SiLU(),
|
||||
nn.Linear(self.block_in, self.block_in),
|
||||
)
|
||||
|
||||
def forward(self, z, rembed=None):
|
||||
# timestep embedding
|
||||
temb = None
|
||||
h = self.conv_in(z)
|
||||
bs, channel, hdim, wdim = h.size()
|
||||
if rembed is not None:
|
||||
rembed = self.resize_level(rembed)
|
||||
rembed = repeat(rembed, 'b e-> b e hd wd', hd=hdim, wd=wdim)
|
||||
h = h + rembed
|
||||
|
||||
# middle
|
||||
if not self.use_checkpoint:
|
||||
h = self.mid_upsclae_transform(h, temb)
|
||||
else:
|
||||
h = checkpoint(self.mid_upsclae_transform, h, temb)
|
||||
|
||||
# end
|
||||
if self.give_pre_end:
|
||||
return h
|
||||
|
||||
h = self.norm_out(h)
|
||||
h = nonlinearity(h)
|
||||
h = self.conv_out(h)
|
||||
if self.tanh_out:
|
||||
h = torch.tanh(h)
|
||||
return h
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('BACKBONE',
|
||||
__class__.__name__,
|
||||
Decoder.para_dict,
|
||||
set_name=True)
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.backbone.unet.unet_module import DiffusionUNet
|
||||
from scepter.modules.model.backbone.unet.unet_module import (DiffusionUNet,
|
||||
DiffusionUNetXL,
|
||||
LargenUNetXL)
|
||||
|
||||
@@ -8,8 +8,9 @@ import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from scepter.modules.model.backbone.unet.unet_utils import (
|
||||
Downsample, ResBlock, SpatialTransformer, Timestep,
|
||||
TimestepEmbedSequential, Upsample, conv_nd, linear, normalization,
|
||||
BasicTransformerBlock, Downsample, ResBlock, SpatialTransformer,
|
||||
SpatialTransformerV2, Timestep, TimestepEmbedSequential,
|
||||
TransformerBlockV2, Upsample, conv_nd, linear, normalization,
|
||||
timestep_embedding, zero_module)
|
||||
from scepter.modules.model.base_model import BaseModel
|
||||
from scepter.modules.model.registry import BACKBONES
|
||||
@@ -951,3 +952,443 @@ class DiffusionUNetXL(DiffusionUNet):
|
||||
__class__.__name__,
|
||||
DiffusionUNetXL.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@BACKBONES.register_class()
|
||||
class LargenUNetXL(DiffusionUNetXL):
|
||||
para_dict = {
|
||||
'TRANSFORMER_BLOCK_TYPE': {
|
||||
'value': 'att_v1'
|
||||
},
|
||||
'IMAGE_SCALE': {
|
||||
'value': 0.0,
|
||||
},
|
||||
}
|
||||
para_dict.update(DiffusionUNetXL.para_dict)
|
||||
|
||||
def __init__(self, cfg, logger):
|
||||
super().__init__(cfg, logger=logger)
|
||||
self.init_params(cfg)
|
||||
self.construct_network()
|
||||
|
||||
def init_params(self, cfg):
|
||||
super().init_params(cfg)
|
||||
self.transformer_block_type = cfg.get('TRANSFORMER_BLOCK_TYPE',
|
||||
'att_v1')
|
||||
TRANSFORMER_BLOCKS = {
|
||||
'att_v1': BasicTransformerBlock,
|
||||
'att_v2': TransformerBlockV2,
|
||||
}
|
||||
assert self.transformer_block_type in list(TRANSFORMER_BLOCKS.keys())
|
||||
self.transformer_block = TRANSFORMER_BLOCKS[
|
||||
self.transformer_block_type]
|
||||
self.image_scale = cfg.get('IMAGE_SCALE', 0.0)
|
||||
self.use_refine = cfg.get('USE_REFINE', False)
|
||||
|
||||
def construct_network(self):
|
||||
in_channels = self.in_channels
|
||||
model_channels = self.model_channels
|
||||
out_channels = self.out_channels
|
||||
attention_resolutions = self.attention_resolutions
|
||||
channel_mult = self.channel_mult
|
||||
num_classes = self.num_classes
|
||||
num_heads = self.num_heads
|
||||
num_head_channels = self.num_head_channels
|
||||
dims = self.dims
|
||||
dropout = self.dropout
|
||||
use_checkpoint = self.use_checkpoint
|
||||
use_scale_shift_norm = self.use_scale_shift_norm
|
||||
disable_self_attentions = self.disable_self_attentions
|
||||
disable_middle_self_attn = self.disable_middle_self_attn
|
||||
transformer_depth = self.transformer_depth
|
||||
transformer_depth_middle = self.transformer_depth_middle
|
||||
context_dim = self.context_dim
|
||||
use_linear_in_transformer = self.use_linear_in_transformer
|
||||
resblock_updown = self.resblock_updown
|
||||
conv_resample = self.conv_resample
|
||||
adm_in_channels = self.adm_in_channels
|
||||
transformer_block = self.transformer_block
|
||||
|
||||
time_embed_dim = model_channels * 4
|
||||
self.time_embed = nn.Sequential(
|
||||
linear(model_channels, time_embed_dim),
|
||||
nn.SiLU(),
|
||||
linear(time_embed_dim, time_embed_dim),
|
||||
)
|
||||
|
||||
if self.num_classes is not None:
|
||||
if isinstance(self.num_classes, int):
|
||||
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
|
||||
elif self.num_classes == 'continuous':
|
||||
print('setting up linear c_adm embedding layer')
|
||||
self.label_emb = nn.Linear(1, time_embed_dim)
|
||||
elif self.num_classes == 'timestep':
|
||||
self.label_emb = nn.Sequential(
|
||||
Timestep(model_channels),
|
||||
nn.Sequential(
|
||||
linear(model_channels, time_embed_dim),
|
||||
nn.SiLU(),
|
||||
linear(time_embed_dim, time_embed_dim),
|
||||
),
|
||||
)
|
||||
elif self.num_classes == 'sequential':
|
||||
assert adm_in_channels is not None
|
||||
self.label_emb = nn.Sequential(
|
||||
nn.Sequential(
|
||||
linear(adm_in_channels, time_embed_dim),
|
||||
nn.SiLU(),
|
||||
linear(time_embed_dim, time_embed_dim),
|
||||
))
|
||||
else:
|
||||
raise ValueError()
|
||||
|
||||
self.input_blocks = nn.ModuleList([
|
||||
TimestepEmbedSequential(
|
||||
conv_nd(dims, in_channels, model_channels, 3, padding=1))
|
||||
])
|
||||
self._feature_size = model_channels
|
||||
input_block_chans = [model_channels]
|
||||
input_down_flag = [False]
|
||||
ch = model_channels
|
||||
ds = 1
|
||||
for level, mult in enumerate(channel_mult):
|
||||
for nr in range(self.num_res_blocks[level]):
|
||||
layers = [
|
||||
ResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=mult * model_channels,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
)
|
||||
]
|
||||
ch = mult * model_channels
|
||||
if ds in attention_resolutions:
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
else:
|
||||
num_heads = ch // num_head_channels
|
||||
dim_head = num_head_channels
|
||||
disabled_sa = disable_self_attentions[level] if exists(
|
||||
disable_self_attentions) else False
|
||||
|
||||
layers.append(
|
||||
SpatialTransformerV2(
|
||||
ch,
|
||||
num_heads,
|
||||
dim_head,
|
||||
transformer_block=transformer_block,
|
||||
depth=transformer_depth[level],
|
||||
context_dim=context_dim,
|
||||
disable_self_attn=disabled_sa,
|
||||
use_linear=use_linear_in_transformer,
|
||||
use_checkpoint=use_checkpoint))
|
||||
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
||||
self._feature_size += ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(False)
|
||||
if level != len(channel_mult) - 1:
|
||||
out_ch = ch
|
||||
self.input_blocks.append(
|
||||
TimestepEmbedSequential(
|
||||
ResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=out_ch,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
down=True,
|
||||
) if resblock_updown else Downsample(
|
||||
ch, conv_resample, dims=dims, out_channels=out_ch))
|
||||
)
|
||||
ch = out_ch
|
||||
input_block_chans.append(ch)
|
||||
input_down_flag.append(True)
|
||||
ds *= 2
|
||||
self._feature_size += ch
|
||||
self._input_block_chans = copy.deepcopy(input_block_chans)
|
||||
self._input_down_flag = input_down_flag
|
||||
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
else:
|
||||
num_heads = ch // num_head_channels
|
||||
dim_head = num_head_channels
|
||||
self.middle_block = TimestepEmbedSequential(
|
||||
ResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
),
|
||||
SpatialTransformerV2(ch,
|
||||
num_heads,
|
||||
dim_head,
|
||||
transformer_block=transformer_block,
|
||||
depth=transformer_depth_middle,
|
||||
context_dim=context_dim,
|
||||
disable_self_attn=disable_middle_self_attn,
|
||||
use_linear=use_linear_in_transformer,
|
||||
use_checkpoint=use_checkpoint),
|
||||
ResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
),
|
||||
)
|
||||
self._feature_size += ch
|
||||
self._middle_block_chans = [ch]
|
||||
|
||||
self._output_block_chans = []
|
||||
self.output_blocks = nn.ModuleList([])
|
||||
for level, mult in list(enumerate(channel_mult))[::-1]:
|
||||
for i in range(self.num_res_blocks[level] + 1):
|
||||
ich = input_block_chans.pop()
|
||||
layers = [
|
||||
ResBlock(
|
||||
ch + ich,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=model_channels * mult,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
)
|
||||
]
|
||||
ch = model_channels * mult
|
||||
if ds in attention_resolutions:
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
else:
|
||||
num_heads = ch // num_head_channels
|
||||
dim_head = num_head_channels
|
||||
disabled_sa = disable_self_attentions[level] if exists(
|
||||
disable_self_attentions) else False
|
||||
layers.append(
|
||||
SpatialTransformerV2(
|
||||
ch,
|
||||
num_heads,
|
||||
dim_head,
|
||||
transformer_block=transformer_block,
|
||||
depth=transformer_depth[level],
|
||||
context_dim=context_dim,
|
||||
disable_self_attn=disabled_sa,
|
||||
use_linear=use_linear_in_transformer,
|
||||
use_checkpoint=use_checkpoint))
|
||||
if level and i == self.num_res_blocks[level]:
|
||||
out_ch = ch
|
||||
layers.append(
|
||||
ResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=out_ch,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
up=True,
|
||||
) if resblock_updown else Upsample(
|
||||
ch, conv_resample, dims=dims, out_channels=out_ch))
|
||||
ds //= 2
|
||||
|
||||
self.output_blocks.append(TimestepEmbedSequential(*layers))
|
||||
|
||||
self._feature_size += ch
|
||||
self._output_block_chans.append(ch)
|
||||
|
||||
self.out = nn.Sequential(
|
||||
normalization(ch),
|
||||
nn.SiLU(),
|
||||
zero_module(
|
||||
conv_nd(dims, model_channels, out_channels, 3, padding=1)),
|
||||
)
|
||||
|
||||
if self.use_refine:
|
||||
self.ref_time_embed = copy.deepcopy(self.time_embed)
|
||||
self.ref_label_emb = copy.deepcopy(self.label_emb)
|
||||
self.ref_input_blocks = copy.deepcopy(self.input_blocks)
|
||||
self.ref_input_blocks[0] = TimestepEmbedSequential(
|
||||
conv_nd(dims, 4, model_channels, 3, padding=1))
|
||||
self.ref_middle_block = copy.deepcopy(self.middle_block)
|
||||
self.ref_output_blocks = copy.deepcopy(self.output_blocks)
|
||||
|
||||
def load_pretrained_model(self, pretrained_model):
|
||||
if pretrained_model is not None:
|
||||
with FS.get_from(pretrained_model,
|
||||
wait_finish=True) as local_model:
|
||||
self.init_from_ckpt(local_model, ignore_keys=self.ignore_keys)
|
||||
|
||||
def init_from_ckpt(self, path, 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:
|
||||
if k == 'input_blocks.0.0.weight':
|
||||
if we.rank == 0:
|
||||
self.logger.info(
|
||||
'Partial initial key {} from state_dict.'.format(
|
||||
k))
|
||||
new_v = torch.empty(320, self.in_channels, 3, 3)
|
||||
nn.init.zeros_(new_v)
|
||||
new_v[:, :v.shape[1]] = v
|
||||
new_sd[k] = new_v
|
||||
if self.use_refine:
|
||||
new_sd['ref_' + k] = v
|
||||
else:
|
||||
new_sd[k] = v
|
||||
if self.use_refine:
|
||||
new_sd['ref_' + k] = v
|
||||
|
||||
missing, unexpected = self.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 forward(self, x, t=None, cond=dict(), **kwargs):
|
||||
t_emb = timestep_embedding(t,
|
||||
self.model_channels,
|
||||
repeat_only=False,
|
||||
legacy=True)
|
||||
emb = self.time_embed(t_emb)
|
||||
|
||||
if isinstance(cond, dict):
|
||||
if 'y' in cond:
|
||||
assert self.num_classes is not None
|
||||
emb = emb + self.label_emb(cond['y'])
|
||||
if self.use_refine:
|
||||
ref_emb = self.ref_time_embed(t_emb)
|
||||
assert 'null_y' in cond
|
||||
cond_y = cond['y'].clone()
|
||||
cond_y[:, :cond['null_y'].shape[1]] = cond['null_y']
|
||||
ref_emb = ref_emb + self.ref_label_emb(cond_y)
|
||||
|
||||
if 'concat' in cond:
|
||||
c = cond['concat']
|
||||
x = torch.cat([x, c], dim=1)
|
||||
|
||||
context = cond.get('crossattn', None)
|
||||
img_context = cond.get('img_crossattn', None)
|
||||
|
||||
task = cond['task']
|
||||
image_scale = cond.get('image_scale', self.image_scale)
|
||||
if 'Subject' in task and img_context is not None:
|
||||
ip_enc_scale = image_scale
|
||||
ip_dec_scale = image_scale
|
||||
num_img_tokens = img_context.shape[1]
|
||||
context = torch.cat([context, img_context], dim=1)
|
||||
else:
|
||||
ip_enc_scale = None
|
||||
ip_dec_scale = None
|
||||
num_img_tokens = None
|
||||
|
||||
ref = cond.get('ref_xt', None)
|
||||
ref_context = cond.get('ref_crossattn', None)
|
||||
else:
|
||||
raise TypeError
|
||||
|
||||
hs = []
|
||||
refs = []
|
||||
h = x
|
||||
|
||||
if self.use_refine:
|
||||
assert ref is not None and ref_context is not None
|
||||
for i, (ref_module, module) in enumerate(
|
||||
zip(self.ref_input_blocks, self.input_blocks)):
|
||||
ref = ref_module(ref, ref_emb, ref_context, caching=None)
|
||||
h = module(h,
|
||||
emb,
|
||||
context,
|
||||
caching=None,
|
||||
scale=ip_enc_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
refs.append(ref)
|
||||
hs.append(h)
|
||||
|
||||
ref = self.ref_middle_block(ref,
|
||||
ref_emb,
|
||||
ref_context,
|
||||
caching=None)
|
||||
h = self.middle_block(h,
|
||||
emb,
|
||||
context,
|
||||
caching=None,
|
||||
scale=ip_enc_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
|
||||
for i, (ref_module, module) in enumerate(
|
||||
zip(self.ref_output_blocks, self.output_blocks)):
|
||||
cache = []
|
||||
ref = torch.cat([ref, refs.pop()], dim=1)
|
||||
ref = ref_module(ref,
|
||||
ref_emb,
|
||||
ref_context,
|
||||
caching='write',
|
||||
cache=cache)
|
||||
h = torch.cat([h, hs.pop()], dim=1)
|
||||
h = module(h,
|
||||
emb,
|
||||
context,
|
||||
caching='read',
|
||||
cache=cache,
|
||||
scale=ip_dec_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
else:
|
||||
for module in self.input_blocks:
|
||||
h = module(h,
|
||||
emb,
|
||||
context,
|
||||
caching=None,
|
||||
scale=ip_enc_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
hs.append(h)
|
||||
h = self.middle_block(h,
|
||||
emb,
|
||||
context,
|
||||
caching=None,
|
||||
scale=ip_enc_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
for module in self.output_blocks:
|
||||
h = torch.cat([h, hs.pop()], dim=1)
|
||||
h = module(h,
|
||||
emb,
|
||||
context,
|
||||
caching=None,
|
||||
scale=ip_dec_scale,
|
||||
num_img_token=num_img_tokens)
|
||||
|
||||
out = self.out(h)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('BACKBONE',
|
||||
__class__.__name__,
|
||||
LargenUNetXL.para_dict,
|
||||
set_name=True)
|
||||
|
||||
@@ -10,6 +10,7 @@ import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms.functional as TF
|
||||
from einops import rearrange, repeat
|
||||
from packaging import version
|
||||
|
||||
@@ -171,12 +172,14 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
|
||||
A sequential module that passes timestep embeddings to the children that
|
||||
support it as an extra input.
|
||||
"""
|
||||
def forward(self, x, emb, context=None, target_size=None):
|
||||
def forward(self, x, emb, context=None, target_size=None, **kwargs):
|
||||
for layer in self:
|
||||
if isinstance(layer, TimestepBlock):
|
||||
x = layer(x, emb)
|
||||
elif isinstance(layer, SpatialTransformer):
|
||||
x = layer(x, context)
|
||||
elif isinstance(layer, SpatialTransformerV2):
|
||||
x = layer(x, context, **kwargs)
|
||||
elif isinstance(layer, Upsample):
|
||||
x = layer(x, target_size)
|
||||
else:
|
||||
@@ -864,6 +867,92 @@ class MemoryEfficientCrossAttention(nn.Module):
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class XFormersMHA_IP(nn.Module):
|
||||
def __init__(self,
|
||||
query_dim,
|
||||
context_dim=None,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.0):
|
||||
super().__init__()
|
||||
inner_dim = dim_head * heads
|
||||
context_dim = default(context_dim, query_dim)
|
||||
|
||||
self.heads = heads
|
||||
self.dim_head = dim_head
|
||||
|
||||
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
|
||||
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
|
||||
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
|
||||
self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
|
||||
self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
|
||||
|
||||
self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim),
|
||||
nn.Dropout(dropout))
|
||||
self.attention_op = None
|
||||
|
||||
def forward(self,
|
||||
x,
|
||||
context=None,
|
||||
mask=None,
|
||||
scale=None,
|
||||
num_img_token=None):
|
||||
q = self.to_q(x)
|
||||
context = default(context, x)
|
||||
|
||||
if scale is not None and num_img_token is not None:
|
||||
eos = context.shape[1] - num_img_token
|
||||
txt_context = context[:, :eos, :]
|
||||
img_context = context[:, eos:, :]
|
||||
|
||||
k = self.to_k(txt_context)
|
||||
v = self.to_v(txt_context)
|
||||
k_i = self.to_k_ip(img_context)
|
||||
v_i = self.to_v_ip(img_context)
|
||||
|
||||
b, _, _ = q.shape
|
||||
q, k, v, k_i, v_i = map(
|
||||
lambda t: t.unsqueeze(3).reshape(b, t.shape[
|
||||
1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
|
||||
b * self.heads, t.shape[1], self.dim_head).contiguous(
|
||||
),
|
||||
(q, k, v, k_i, v_i),
|
||||
)
|
||||
|
||||
# actually compute the attention, what we cannot get enough of
|
||||
txt_out = xformers.ops.memory_efficient_attention(
|
||||
q, k, v, attn_bias=None, op=self.attention_op)
|
||||
img_out = xformers.ops.memory_efficient_attention(
|
||||
q, k_i, v_i, attn_bias=None, op=self.attention_op)
|
||||
out = txt_out + scale * img_out
|
||||
else:
|
||||
k = self.to_k(context)
|
||||
v = self.to_v(context)
|
||||
b, _, _ = q.shape
|
||||
q, k, v = map(
|
||||
lambda t: t.unsqueeze(3).reshape(b, t.shape[
|
||||
1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
|
||||
b * self.heads, t.shape[1], self.dim_head).contiguous(
|
||||
),
|
||||
(q, k, v),
|
||||
)
|
||||
out = xformers.ops.memory_efficient_attention(q,
|
||||
k,
|
||||
v,
|
||||
attn_bias=None,
|
||||
op=self.attention_op)
|
||||
|
||||
# TODO: Use this directly in the attention operation, as a bias
|
||||
if exists(mask):
|
||||
raise NotImplementedError
|
||||
out = (out.unsqueeze(0).reshape(
|
||||
b, self.heads, out.shape[1],
|
||||
self.dim_head).permute(0, 2, 1,
|
||||
3).reshape(b, out.shape[1],
|
||||
self.heads * self.dim_head))
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class BasicTransformerBlock(nn.Module):
|
||||
def __init__(self,
|
||||
dim,
|
||||
@@ -908,6 +997,65 @@ class BasicTransformerBlock(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
class TransformerBlockV2(nn.Module):
|
||||
def __init__(self,
|
||||
query_dim,
|
||||
n_heads,
|
||||
d_head,
|
||||
dropout=0.,
|
||||
context_dim=None,
|
||||
gated_ff=True,
|
||||
use_checkpoint=False,
|
||||
disable_self_attn=False):
|
||||
super().__init__()
|
||||
self.disable_self_attn = disable_self_attn
|
||||
self.attn1 = MemoryEfficientCrossAttention(query_dim=query_dim,
|
||||
heads=n_heads,
|
||||
dim_head=d_head,
|
||||
dropout=dropout,
|
||||
context_dim=None)
|
||||
self.ff = FeedForward(query_dim, dropout=dropout, glu=gated_ff)
|
||||
self.attn2 = XFormersMHA_IP(query_dim=query_dim,
|
||||
heads=n_heads,
|
||||
dim_head=d_head,
|
||||
context_dim=context_dim)
|
||||
self.norm1 = nn.LayerNorm(query_dim)
|
||||
self.norm2 = nn.LayerNorm(query_dim)
|
||||
self.norm3 = nn.LayerNorm(query_dim)
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
def forward(self,
|
||||
x,
|
||||
context,
|
||||
caching=None,
|
||||
cache=None,
|
||||
scale=None,
|
||||
num_img_token=None,
|
||||
**kwargs):
|
||||
y = self.norm1(x)
|
||||
if caching == 'write':
|
||||
assert isinstance(cache, list)
|
||||
cache.append(y)
|
||||
x = self.attn1(y, context=None) + x
|
||||
elif caching == 'read':
|
||||
assert isinstance(cache, list) and len(cache) > 0
|
||||
c = cache.pop(0)
|
||||
self_ctx = torch.cat([y, c], dim=1)
|
||||
x = self.attn1(y, context=self_ctx) + x
|
||||
elif caching is None:
|
||||
x = self.attn1(y, context=None) + x
|
||||
else:
|
||||
assert False
|
||||
|
||||
x = self.attn2(self.norm2(x),
|
||||
context=context,
|
||||
scale=scale,
|
||||
num_img_token=num_img_token) + x
|
||||
x = self.ff(self.norm3(x)) + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SpatialTransformer(nn.Module):
|
||||
"""
|
||||
Transformer block for image-like data.
|
||||
@@ -1003,3 +1151,108 @@ class SpatialTransformer(nn.Module):
|
||||
if not self.use_linear:
|
||||
x = self.proj_out(x)
|
||||
return x + x_in
|
||||
|
||||
|
||||
class SpatialTransformerV2(nn.Module):
|
||||
"""
|
||||
Transformer block for image-like data.
|
||||
First, project the input (aka embedding)
|
||||
and reshape to b, t, d.
|
||||
Then apply standard transformer action.
|
||||
Finally, reshape to image
|
||||
NEW: use_linear for more efficiency instead of the 1x1 convs
|
||||
"""
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
n_heads,
|
||||
d_head,
|
||||
transformer_block,
|
||||
depth=1,
|
||||
dropout=0.,
|
||||
context_dim=None,
|
||||
disable_self_attn=False,
|
||||
use_linear=False,
|
||||
use_checkpoint=True):
|
||||
super().__init__()
|
||||
if exists(context_dim) and not isinstance(context_dim, list):
|
||||
context_dim = [context_dim]
|
||||
|
||||
if exists(context_dim) and not isinstance(context_dim, (list)):
|
||||
context_dim = [context_dim]
|
||||
if exists(context_dim) and isinstance(context_dim, list):
|
||||
if depth != len(context_dim):
|
||||
print(
|
||||
f'WARNING: {self.__class__.__name__}: Found context dims {context_dim} of'
|
||||
f" depth {len(context_dim)}, which does not match the specified 'depth' of"
|
||||
f' {depth}. Setting context_dim to {depth * [context_dim[0]]} now.'
|
||||
)
|
||||
# depth does not match context dims.
|
||||
assert all(
|
||||
map(lambda x: x == context_dim[0], context_dim)
|
||||
), 'need homogenous context_dim to match depth automatically'
|
||||
context_dim = depth * [context_dim[0]]
|
||||
elif context_dim is None:
|
||||
context_dim = [None] * depth
|
||||
|
||||
self.in_channels = in_channels
|
||||
inner_dim = n_heads * d_head
|
||||
self.norm = normalization(in_channels)
|
||||
if not use_linear:
|
||||
self.proj_in = nn.Conv2d(in_channels,
|
||||
inner_dim,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0)
|
||||
else:
|
||||
self.proj_in = nn.Linear(in_channels, inner_dim)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList([
|
||||
transformer_block(inner_dim,
|
||||
n_heads,
|
||||
d_head,
|
||||
dropout=dropout,
|
||||
context_dim=context_dim[d],
|
||||
disable_self_attn=disable_self_attn,
|
||||
use_checkpoint=use_checkpoint)
|
||||
for d in range(depth)
|
||||
])
|
||||
if not use_linear:
|
||||
self.proj_out = zero_module(
|
||||
nn.Conv2d(inner_dim,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0))
|
||||
else:
|
||||
self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
|
||||
self.use_linear = use_linear
|
||||
|
||||
def forward(self, x, context=None, **kwargs):
|
||||
# note: if no context is given, cross-attention defaults to self-attention
|
||||
if not isinstance(context, list):
|
||||
context = [context]
|
||||
b, c, h, w = x.shape
|
||||
|
||||
ref_mask = kwargs.pop('ref_mask', None)
|
||||
if ref_mask is not None:
|
||||
ref_mask = TF.resize(ref_mask, (h, w), antialias=True)
|
||||
ref_mask = (ref_mask > 0.5).float()
|
||||
ref_mask = rearrange(ref_mask, 'b c h w -> b (h w) c').contiguous()
|
||||
|
||||
x_in = x
|
||||
x = self.norm(x)
|
||||
if not self.use_linear:
|
||||
x = self.proj_in(x)
|
||||
x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
|
||||
if self.use_linear:
|
||||
x = self.proj_in(x)
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
if i > 0 and len(context) == 1:
|
||||
i = 0 # use same context for each block
|
||||
x = block(x, context=context[i], ref_mask=ref_mask, **kwargs)
|
||||
if self.use_linear:
|
||||
x = self.proj_out(x)
|
||||
x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
|
||||
if not self.use_linear:
|
||||
x = self.proj_out(x)
|
||||
return x + x_in
|
||||
|
||||
@@ -1,17 +1,5 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional
|
||||
from einops import rearrange
|
||||
|
||||
from scepter.modules.model.backbone.video.init_helper import (
|
||||
_init_transformer_weights, trunc_normal_)
|
||||
from scepter.modules.model.registry import BACKBONES, BRICKS, STEMS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
'''
|
||||
The implementations of vivit as https://arxiv.org/abs/2103.15691.
|
||||
The following setting alined the proposed model in the paper above.
|
||||
@@ -39,6 +27,18 @@ TimesFormer:
|
||||
complexity: (n_h * n_w) ** 2 + O(attn_temp)
|
||||
'''
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional
|
||||
from einops import rearrange
|
||||
|
||||
from scepter.modules.model.backbone.video.init_helper import (
|
||||
_init_transformer_weights, trunc_normal_)
|
||||
from scepter.modules.model.registry import BACKBONES, BRICKS, STEMS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
|
||||
|
||||
@BACKBONES.register_class()
|
||||
class VideoTransformer(nn.Module):
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from scepter.modules.model.embedder.embedder import (ConcatTimestepEmbedderND,
|
||||
FrozenCLIPEmbedder,
|
||||
FrozenOpenCLIPEmbedder,
|
||||
FrozenOpenCLIPEmbedder2,
|
||||
GeneralConditioner)
|
||||
GeneralConditioner,
|
||||
IPAdapterPlusEmbedder,
|
||||
RefCrossEmbedder)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -22,9 +22,10 @@ from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
from .base_embedder import BaseEmbedder
|
||||
from .resampler import Resampler
|
||||
|
||||
try:
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
|
||||
except Exception as e:
|
||||
warnings.warn(
|
||||
f'Import transformers error, please deal with this problem: {e}')
|
||||
@@ -513,6 +514,93 @@ class ConcatTimestepEmbedderND(BaseEmbedder):
|
||||
set_name=True)
|
||||
|
||||
|
||||
@EMBEDDERS.register_class()
|
||||
class IPAdapterPlusEmbedder(BaseEmbedder):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
with FS.get_dir_to_local_dir(cfg.CLIP_DIR,
|
||||
wait_finish=True) as local_path:
|
||||
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
local_path)
|
||||
|
||||
self.image_proj_model = Resampler(
|
||||
dim=self.cfg.get('IN_DIM', 768),
|
||||
depth=self.cfg.get('DEPTH', 4),
|
||||
dim_head=64,
|
||||
heads=self.cfg.get('HEADS', 12),
|
||||
num_queries=self.cfg.get('NUM_TOKENS', 16),
|
||||
embedding_dim=self.image_encoder.config.hidden_size,
|
||||
output_dim=self.cfg.get('CROSSATTN_DIM', 768),
|
||||
ff_mult=4,
|
||||
)
|
||||
|
||||
with FS.get_from(cfg.PRETRAINED_MODEL, wait_finish=True) as local_path:
|
||||
ckpt = torch.load(local_path, map_location='cpu')
|
||||
self.image_proj_model.load_state_dict(ckpt['image_proj'],
|
||||
strict=True)
|
||||
|
||||
self.patch_projector = nn.Linear(self.image_encoder.config.hidden_size,
|
||||
self.cfg.get('CROSSATTN_DIM', 768))
|
||||
|
||||
def encode(self, ref_ip, ref_detail):
|
||||
encoder_output = self.image_encoder(ref_ip, output_hidden_states=True)
|
||||
image_prompt_embeds = self.image_proj_model(
|
||||
encoder_output.hidden_states[-2])
|
||||
encoder_output_2 = self.image_encoder(ref_detail,
|
||||
output_hidden_states=True)
|
||||
image_patch_embeds = self.patch_projector(
|
||||
encoder_output_2.last_hidden_state)
|
||||
out = {
|
||||
'img_crossattn': image_prompt_embeds,
|
||||
'ref_crossattn': image_patch_embeds,
|
||||
}
|
||||
return out
|
||||
|
||||
def forward(self, ref_ip, ref_detail):
|
||||
return self.encode(ref_ip, ref_detail)
|
||||
|
||||
|
||||
class RefCrossEmbedder(BaseEmbedder):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg, logger=logger)
|
||||
|
||||
with FS.get_dir_to_local_dir(cfg.CLIP_DIR,
|
||||
wait_finish=True) as local_path:
|
||||
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
local_path)
|
||||
|
||||
self.patch_projector = nn.Linear(self.image_encoder.config.hidden_size,
|
||||
self.cfg.get('CROSSATTN_DIM', 768))
|
||||
|
||||
def encode(self, img):
|
||||
encoder_output = self.image_encoder(img, output_hidden_states=True)
|
||||
image_patch_embeds = self.patch_projector(
|
||||
encoder_output.last_hidden_state)
|
||||
out = {
|
||||
'ref_crossattn': image_patch_embeds,
|
||||
}
|
||||
return out
|
||||
|
||||
def forward(self, img):
|
||||
return self.encode(img)
|
||||
|
||||
|
||||
@EMBEDDERS.register_class()
|
||||
class TransparentEmbedder(BaseEmbedder):
|
||||
def forward(self, *args):
|
||||
out = dict()
|
||||
for key, val in zip(self.input_keys, args):
|
||||
out[key] = val
|
||||
return out
|
||||
|
||||
|
||||
@EMBEDDERS.register_class()
|
||||
class NoiseConcatEmbedder(BaseEmbedder):
|
||||
def forward(self, *args):
|
||||
return {'concat': torch.cat(args, dim=1)}
|
||||
|
||||
|
||||
@EMBEDDERS.register_class()
|
||||
class GeneralConditioner(BaseEmbedder):
|
||||
OUTPUT_DIM2KEYS = {2: 'y', 3: 'crossattn', 4: 'concat', 5: 'concat'}
|
||||
@@ -598,42 +686,60 @@ class GeneralConditioner(BaseEmbedder):
|
||||
with embedding_context():
|
||||
if hasattr(embedder, 'input_key') and (embedder.input_key
|
||||
is not None):
|
||||
if embedder.input_key not in batch:
|
||||
continue
|
||||
if embedder.legacy_ucg_val is not None:
|
||||
batch = self.possibly_get_ucg_val(embedder, batch)
|
||||
emb_out = embedder(batch[embedder.input_key])
|
||||
elif hasattr(embedder, 'input_keys'):
|
||||
if any([k not in batch for k in embedder.input_keys]):
|
||||
continue
|
||||
emb_out = embedder(
|
||||
*[batch[k] for k in embedder.input_keys])
|
||||
assert isinstance(
|
||||
emb_out, (torch.Tensor, list, tuple)
|
||||
), f'encoder outputs must be tensors or a sequence, but got {type(emb_out)}'
|
||||
if not isinstance(emb_out, (list, tuple)):
|
||||
emb_out = [emb_out]
|
||||
for emb in emb_out:
|
||||
# print("emb.shape", emb.shape)
|
||||
# print("emb.input_keys", embedder.input_keys)
|
||||
out_key = self.OUTPUT_DIM2KEYS[emb.dim()]
|
||||
if embedder.ucg_rate > 0.0 and embedder.legacy_ucg_val is None:
|
||||
emb = (expand_dims_like(
|
||||
torch.bernoulli(
|
||||
(1.0 - embedder.ucg_rate) *
|
||||
torch.ones(emb.shape[0], device=emb.device)),
|
||||
emb,
|
||||
) * emb)
|
||||
if (hasattr(embedder, 'input_keys')):
|
||||
if np.sum(
|
||||
np.array([
|
||||
key in force_zero_embeddings
|
||||
for key in embedder.input_keys
|
||||
])) > 0:
|
||||
emb = torch.zeros_like(emb)
|
||||
if out_key in output:
|
||||
output[out_key] = torch.cat((output[out_key], emb),
|
||||
self.KEY2CATDIM[out_key])
|
||||
else:
|
||||
output[out_key] = emb
|
||||
# if "y" in output:
|
||||
# print("out.shape", output["y"].shape)
|
||||
|
||||
if isinstance(emb_out, dict):
|
||||
for key, val in emb_out.items():
|
||||
if key in output:
|
||||
# 重复出现的key必须在（y, crossattn, concat）中，否则raise error
|
||||
assert key in self.KEY2CATDIM
|
||||
output[key] = torch.cat([output[key], val],
|
||||
dim=self.KEY2CATDIM[key])
|
||||
else:
|
||||
output[key] = val
|
||||
else:
|
||||
assert isinstance(
|
||||
emb_out, (torch.Tensor, list, tuple)
|
||||
), f'encoder outputs must be tensors or a sequence, but got {type(emb_out)}'
|
||||
|
||||
if not isinstance(emb_out, (list, tuple)):
|
||||
emb_out = [emb_out]
|
||||
|
||||
for emb in emb_out:
|
||||
# 根据emb的维度，判断该cond归属于 (y, concat, crossattn）中的哪一种
|
||||
out_key = self.OUTPUT_DIM2KEYS[emb.dim()]
|
||||
|
||||
if embedder.ucg_rate > 0.0 and embedder.legacy_ucg_val is None:
|
||||
emb = (expand_dims_like(
|
||||
torch.bernoulli(
|
||||
(1.0 - embedder.ucg_rate) *
|
||||
torch.ones(emb.shape[0], device=emb.device)),
|
||||
emb,
|
||||
) * emb)
|
||||
|
||||
if (hasattr(embedder, 'input_keys')):
|
||||
if np.sum(
|
||||
np.array([
|
||||
key in force_zero_embeddings
|
||||
for key in embedder.input_keys
|
||||
])) > 0:
|
||||
emb = torch.zeros_like(emb)
|
||||
|
||||
if out_key in output:
|
||||
output[out_key] = torch.cat((output[out_key], emb),
|
||||
self.KEY2CATDIM[out_key])
|
||||
else:
|
||||
output[out_key] = emb
|
||||
|
||||
return output
|
||||
|
||||
def get_unconditional_conditioning(self,
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from einops.layers.torch import Rearrange
|
||||
|
||||
|
||||
# FFN
|
||||
def FeedForward(dim, mult=4):
|
||||
inner_dim = int(dim * mult)
|
||||
return nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
nn.Linear(dim, inner_dim, bias=False),
|
||||
nn.GELU(),
|
||||
nn.Linear(inner_dim, dim, bias=False),
|
||||
)
|
||||
|
||||
|
||||
def reshape_tensor(x, heads):
|
||||
bs, length, width = x.shape
|
||||
# (bs, length, width) --> (bs, length, n_heads, dim_per_head)
|
||||
x = x.view(bs, length, heads, -1)
|
||||
# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
|
||||
x = x.transpose(1, 2)
|
||||
# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
|
||||
x = x.reshape(bs, heads, length, -1)
|
||||
return x
|
||||
|
||||
|
||||
class PerceiverAttention(nn.Module):
|
||||
def __init__(self, *, dim, dim_head=64, heads=8):
|
||||
super().__init__()
|
||||
self.scale = dim_head**-0.5
|
||||
self.dim_head = dim_head
|
||||
self.heads = heads
|
||||
inner_dim = dim_head * heads
|
||||
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
|
||||
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
||||
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
|
||||
self.to_out = nn.Linear(inner_dim, dim, bias=False)
|
||||
|
||||
def forward(self, x, latents):
|
||||
"""
|
||||
Args:
|
||||
x (torch.Tensor): image features
|
||||
shape (b, n1, D)
|
||||
latent (torch.Tensor): latent features
|
||||
shape (b, n2, D)
|
||||
"""
|
||||
x = self.norm1(x)
|
||||
latents = self.norm2(latents)
|
||||
|
||||
b, l, _ = latents.shape
|
||||
|
||||
q = self.to_q(latents)
|
||||
kv_input = torch.cat((x, latents), dim=-2)
|
||||
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
|
||||
|
||||
q = reshape_tensor(q, self.heads)
|
||||
k = reshape_tensor(k, self.heads)
|
||||
v = reshape_tensor(v, self.heads)
|
||||
|
||||
# attention
|
||||
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
|
||||
weight = (q * scale) @ (k * scale).transpose(
|
||||
-2, -1) # More stable with f16 than dividing afterwards
|
||||
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
|
||||
out = weight @ v
|
||||
|
||||
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
|
||||
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class Resampler(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=1024,
|
||||
depth=8,
|
||||
dim_head=64,
|
||||
heads=16,
|
||||
num_queries=8,
|
||||
embedding_dim=768,
|
||||
output_dim=1024,
|
||||
ff_mult=4,
|
||||
max_seq_len: int = 257, # CLIP tokens + CLS token
|
||||
apply_pos_emb: bool = False,
|
||||
num_latents_mean_pooled:
|
||||
int = 0, # number of latents derived from mean pooled representation of the sequence
|
||||
):
|
||||
super().__init__()
|
||||
self.pos_emb = nn.Embedding(max_seq_len,
|
||||
embedding_dim) if apply_pos_emb else None
|
||||
|
||||
self.latents = nn.Parameter(
|
||||
torch.randn(1, num_queries, dim) / dim**0.5)
|
||||
|
||||
self.proj_in = nn.Linear(embedding_dim, dim)
|
||||
|
||||
self.proj_out = nn.Linear(dim, output_dim)
|
||||
self.norm_out = nn.LayerNorm(output_dim)
|
||||
|
||||
self.to_latents_from_mean_pooled_seq = (nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
nn.Linear(dim, dim * num_latents_mean_pooled),
|
||||
Rearrange('b (n d) -> b n d', n=num_latents_mean_pooled),
|
||||
) if num_latents_mean_pooled > 0 else None)
|
||||
|
||||
self.layers = nn.ModuleList([])
|
||||
for _ in range(depth):
|
||||
self.layers.append(
|
||||
nn.ModuleList([
|
||||
PerceiverAttention(dim=dim, dim_head=dim_head,
|
||||
heads=heads),
|
||||
FeedForward(dim=dim, mult=ff_mult),
|
||||
]))
|
||||
|
||||
def forward(self, x):
|
||||
if self.pos_emb is not None:
|
||||
n, device = x.shape[1], x.device
|
||||
pos_emb = self.pos_emb(torch.arange(n, device=device))
|
||||
x = x + pos_emb
|
||||
|
||||
latents = self.latents.repeat(x.size(0), 1, 1)
|
||||
|
||||
x = self.proj_in(x)
|
||||
|
||||
if self.to_latents_from_mean_pooled_seq:
|
||||
meanpooled_seq = masked_mean(x,
|
||||
dim=1,
|
||||
mask=torch.ones(x.shape[:2],
|
||||
device=x.device,
|
||||
dtype=torch.bool))
|
||||
meanpooled_latents = self.to_latents_from_mean_pooled_seq(
|
||||
meanpooled_seq)
|
||||
latents = torch.cat((meanpooled_latents, latents), dim=-2)
|
||||
|
||||
for attn, ff in self.layers:
|
||||
latents = attn(x, latents) + latents
|
||||
latents = ff(latents) + latents
|
||||
|
||||
latents = self.proj_out(latents)
|
||||
return self.norm_out(latents)
|
||||
|
||||
|
||||
def masked_mean(t, *, dim, mask=None):
|
||||
if mask is None:
|
||||
return t.mean(dim=dim)
|
||||
|
||||
denom = mask.sum(dim=dim, keepdim=True)
|
||||
mask = rearrange(mask, 'b n -> b n 1')
|
||||
masked_t = t.masked_fill(~mask, 0.0)
|
||||
|
||||
return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)
|
||||
@@ -1,5 +1,8 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.head.classifier_head import (
|
||||
ClassifierHead, CosineLinearHead, TransformerHead, TransformerHeadx2,
|
||||
VideoClassifierHead, VideoClassifierHeadx2)
|
||||
from scepter.modules.model.head.classifier_head import (ClassifierHead,
|
||||
CosineLinearHead,
|
||||
TransformerHead,
|
||||
TransformerHeadx2,
|
||||
VideoClassifierHead,
|
||||
VideoClassifierHeadx2)
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
from scepter.modules.model.metric.classification import (AccuracyMetric,
|
||||
EnsembleAccuracyMetric
|
||||
)
|
||||
from scepter.modules.model.metric.classification import (
|
||||
AccuracyMetric, EnsembleAccuracyMetric)
|
||||
|
||||
@@ -13,8 +13,7 @@ from .schedules import karras_schedule
|
||||
from .solvers import (sample_ddim, sample_dpm_2, sample_dpm_2_ancestral,
|
||||
sample_dpmpp_2m, sample_dpmpp_2m_sde,
|
||||
sample_dpmpp_2s_ancestral, sample_dpmpp_sde,
|
||||
sample_euler, sample_euler_ancestral, sample_heun,
|
||||
sample_img2img_euler, sample_img2img_euler_ancestral)
|
||||
sample_euler, sample_euler_ancestral, sample_heun)
|
||||
|
||||
__all__ = ['GaussianDiffusion']
|
||||
|
||||
@@ -27,6 +26,148 @@ def _i(tensor, t, x):
|
||||
return tensor[t.to(tensor.device)].view(shape).to(x.device)
|
||||
|
||||
|
||||
def _unpack_2d_ks(kernel_size):
|
||||
if isinstance(kernel_size, int):
|
||||
ky = kx = kernel_size
|
||||
else:
|
||||
assert len(
|
||||
kernel_size) == 2, '2D Kernel size should have a length of 2.'
|
||||
ky, kx = kernel_size
|
||||
|
||||
ky = int(ky)
|
||||
kx = int(kx)
|
||||
return ky, kx
|
||||
|
||||
|
||||
def _compute_zero_padding(kernel_size):
|
||||
ky, kx = _unpack_2d_ks(kernel_size)
|
||||
return (ky - 1) // 2, (kx - 1) // 2
|
||||
|
||||
|
||||
def _bilateral_blur(
|
||||
input,
|
||||
guidance,
|
||||
kernel_size,
|
||||
sigma_color,
|
||||
sigma_space,
|
||||
border_type='reflect',
|
||||
color_distance_type='l1',
|
||||
):
|
||||
|
||||
if isinstance(sigma_color, torch.Tensor):
|
||||
sigma_color = sigma_color.to(device=input.device,
|
||||
dtype=input.dtype).view(-1, 1, 1, 1, 1)
|
||||
|
||||
ky, kx = _unpack_2d_ks(kernel_size)
|
||||
pad_y, pad_x = _compute_zero_padding(kernel_size)
|
||||
|
||||
padded_input = torch.nn.functional.pad(input, (pad_x, pad_x, pad_y, pad_y),
|
||||
mode=border_type)
|
||||
unfolded_input = padded_input.unfold(2, ky, 1).unfold(3, kx, 1).flatten(
|
||||
-2) # (B, C, H, W, Ky x Kx)
|
||||
|
||||
if guidance is None:
|
||||
guidance = input
|
||||
unfolded_guidance = unfolded_input
|
||||
else:
|
||||
padded_guidance = torch.nn.functional.pad(guidance,
|
||||
(pad_x, pad_x, pad_y, pad_y),
|
||||
mode=border_type)
|
||||
unfolded_guidance = padded_guidance.unfold(2, ky, 1).unfold(
|
||||
3, kx, 1).flatten(-2) # (B, C, H, W, Ky x Kx)
|
||||
|
||||
diff = unfolded_guidance - guidance.unsqueeze(-1)
|
||||
if color_distance_type == 'l1':
|
||||
color_distance_sq = diff.abs().sum(1, keepdim=True).square()
|
||||
elif color_distance_type == 'l2':
|
||||
color_distance_sq = diff.square().sum(1, keepdim=True)
|
||||
else:
|
||||
raise ValueError('color_distance_type only acceps l1 or l2')
|
||||
color_kernel = (-0.5 / sigma_color**2 *
|
||||
color_distance_sq).exp() # (B, 1, H, W, Ky x Kx)
|
||||
|
||||
space_kernel = get_gaussian_kernel2d(kernel_size,
|
||||
sigma_space,
|
||||
device=input.device,
|
||||
dtype=input.dtype)
|
||||
space_kernel = space_kernel.view(-1, 1, 1, 1, kx * ky)
|
||||
|
||||
kernel = space_kernel * color_kernel
|
||||
out = (unfolded_input * kernel).sum(-1) / kernel.sum(-1)
|
||||
return out
|
||||
|
||||
|
||||
def get_gaussian_kernel1d(
|
||||
kernel_size,
|
||||
sigma,
|
||||
force_even,
|
||||
*,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
|
||||
return gaussian(kernel_size, sigma, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def gaussian(window_size, sigma, *, device=None, dtype=None):
|
||||
|
||||
batch_size = sigma.shape[0]
|
||||
|
||||
x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) -
|
||||
window_size // 2).expand(batch_size, -1)
|
||||
|
||||
if window_size % 2 == 0:
|
||||
x = x + 0.5
|
||||
|
||||
gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
|
||||
|
||||
return gauss / gauss.sum(-1, keepdim=True)
|
||||
|
||||
|
||||
def get_gaussian_kernel2d(
|
||||
kernel_size,
|
||||
sigma,
|
||||
force_even=False,
|
||||
*,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
|
||||
sigma = torch.Tensor([[sigma, sigma]]).to(device=device, dtype=dtype)
|
||||
|
||||
ksize_y, ksize_x = _unpack_2d_ks(kernel_size)
|
||||
sigma_y, sigma_x = sigma[:, 0, None], sigma[:, 1, None]
|
||||
|
||||
kernel_y = get_gaussian_kernel1d(ksize_y,
|
||||
sigma_y,
|
||||
force_even,
|
||||
device=device,
|
||||
dtype=dtype)[..., None]
|
||||
kernel_x = get_gaussian_kernel1d(ksize_x,
|
||||
sigma_x,
|
||||
force_even,
|
||||
device=device,
|
||||
dtype=dtype)[..., None]
|
||||
|
||||
return kernel_y * kernel_x.view(-1, 1, ksize_x)
|
||||
|
||||
|
||||
def adaptive_anisotropic_filter(x, g=None):
|
||||
if g is None:
|
||||
g = x
|
||||
s, m = torch.std_mean(g, dim=(1, 2, 3), keepdim=True)
|
||||
s = s + 1e-5
|
||||
guidance = (g - m) / s
|
||||
y = _bilateral_blur(x,
|
||||
guidance,
|
||||
kernel_size=(13, 13),
|
||||
sigma_color=3.0,
|
||||
sigma_space=3.0,
|
||||
border_type='reflect',
|
||||
color_distance_type='l1')
|
||||
return y
|
||||
|
||||
|
||||
class GaussianDiffusion(object):
|
||||
def __init__(self, sigmas, prediction_type='eps'):
|
||||
assert prediction_type in {'x0', 'eps', 'v'}
|
||||
@@ -53,6 +194,7 @@ class GaussianDiffusion(object):
|
||||
guide_scale=None,
|
||||
guide_rescale=None,
|
||||
clamp=None,
|
||||
sharpness=0.0,
|
||||
percentile=None,
|
||||
cat_uc=False,
|
||||
**kwargs):
|
||||
@@ -79,54 +221,99 @@ class GaussianDiffusion(object):
|
||||
|
||||
# prediction
|
||||
if guide_scale is None:
|
||||
assert isinstance(model_kwargs, dict)
|
||||
out = model(xt, t=t, **model_kwargs, **kwargs)
|
||||
if isinstance(model_kwargs, dict):
|
||||
out = model(xt, t=t, **model_kwargs, **kwargs)
|
||||
elif isinstance(model_kwargs, list) and len(model_kwargs) > 0:
|
||||
out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
else:
|
||||
raise Exception('Error')
|
||||
else:
|
||||
# classifier-free guidance (arXiv:2207.12598)
|
||||
# model_kwargs[0]: conditional kwargs
|
||||
# model_kwargs[1]: non-conditional kwargs
|
||||
assert isinstance(model_kwargs, list) and len(model_kwargs) == 2
|
||||
|
||||
if guide_scale == 1.:
|
||||
out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
else:
|
||||
if cat_uc:
|
||||
|
||||
def parse_model_kwargs(prev_value, value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
prev_value = torch.cat([prev_value, value], dim=0)
|
||||
elif isinstance(value, dict):
|
||||
for k, v in value.items():
|
||||
prev_value[k] = parse_model_kwargs(
|
||||
prev_value[k], v)
|
||||
elif isinstance(value, list):
|
||||
for idx, v in enumerate(value):
|
||||
prev_value[idx] = parse_model_kwargs(
|
||||
prev_value[idx], v)
|
||||
return prev_value
|
||||
|
||||
all_model_kwargs = copy.deepcopy(model_kwargs[0])
|
||||
for model_kwarg in model_kwargs[1:]:
|
||||
for key, value in model_kwarg.items():
|
||||
all_model_kwargs[key] = parse_model_kwargs(
|
||||
all_model_kwargs[key], value)
|
||||
all_out = model(xt.repeat(2, 1, 1, 1),
|
||||
t=t.repeat(2),
|
||||
**all_model_kwargs,
|
||||
**kwargs)
|
||||
y_out, u_out = all_out.chunk(2)
|
||||
assert isinstance(model_kwargs, list) and len(model_kwargs) >= 2
|
||||
if isinstance(guide_scale, float) or isinstance(guide_scale, int):
|
||||
assert len(model_kwargs) == 2
|
||||
if guide_scale == 1.:
|
||||
out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
else:
|
||||
y_out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
u_out = model(xt, t=t, **model_kwargs[1], **kwargs)
|
||||
out = u_out + guide_scale * (y_out - u_out)
|
||||
if cat_uc:
|
||||
|
||||
# rescale the output according to arXiv:2305.08891
|
||||
if guide_rescale is not None:
|
||||
assert guide_rescale >= 0 and guide_rescale <= 1
|
||||
ratio = (y_out.flatten(1).std(dim=1) /
|
||||
(out.flatten(1).std(dim=1) +
|
||||
1e-12)).view((-1, ) + (1, ) * (y_out.ndim - 1))
|
||||
out *= guide_rescale * ratio + (1 - guide_rescale) * 1.0
|
||||
def parse_model_kwargs(prev_value, value):
|
||||
if isinstance(value, torch.Tensor):
|
||||
prev_value = torch.cat([prev_value, value],
|
||||
dim=0)
|
||||
elif isinstance(value, dict):
|
||||
for k, v in value.items():
|
||||
prev_value[k] = parse_model_kwargs(
|
||||
prev_value[k], v)
|
||||
elif isinstance(value, list):
|
||||
for idx, v in enumerate(value):
|
||||
prev_value[idx] = parse_model_kwargs(
|
||||
prev_value[idx], v)
|
||||
return prev_value
|
||||
|
||||
all_model_kwargs = copy.deepcopy(model_kwargs[0])
|
||||
for model_kwarg in model_kwargs[1:]:
|
||||
for key, value in model_kwarg.items():
|
||||
all_model_kwargs[key] = parse_model_kwargs(
|
||||
all_model_kwargs[key], value)
|
||||
all_out = model(xt.repeat(2, 1, 1, 1),
|
||||
t=t.repeat(2),
|
||||
**all_model_kwargs,
|
||||
**kwargs)
|
||||
y_out, u_out = all_out.chunk(2)
|
||||
else:
|
||||
y_out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
u_out = model(xt, t=t, **model_kwargs[1], **kwargs)
|
||||
# todo sharpness
|
||||
# sharpness sampling
|
||||
if sharpness is not None and sharpness > 0:
|
||||
positive_x0 = alphas * xt - sigmas * y_out
|
||||
negative_x0 = alphas * xt - sigmas * u_out
|
||||
|
||||
positive_eps = xt - positive_x0
|
||||
negative_eps = xt - negative_x0
|
||||
|
||||
global_diffusion_progress = (
|
||||
1 - t / 999.0).detach().cpu().numpy().tolist()[0]
|
||||
alpha = 0.001 * sharpness * global_diffusion_progress
|
||||
positive_eps_degraded = adaptive_anisotropic_filter(
|
||||
x=positive_eps, g=positive_x0)
|
||||
positive_eps_degraded_weighted = positive_eps_degraded * alpha + positive_eps * (
|
||||
1.0 - alpha)
|
||||
|
||||
final_eps = negative_eps + guide_scale * (
|
||||
positive_eps_degraded_weighted - negative_eps)
|
||||
final_x0 = xt - final_eps
|
||||
out = (alphas * xt - final_x0) / sigmas
|
||||
else:
|
||||
out = u_out + guide_scale * (y_out - u_out)
|
||||
elif isinstance(guide_scale, dict):
|
||||
assert len(model_kwargs) == 3
|
||||
y_out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
m_out = model(xt, t=t, **model_kwargs[1], **kwargs)
|
||||
u_out = model(xt, t=t, **model_kwargs[2], **kwargs)
|
||||
out = u_out + guide_scale['image'] * (
|
||||
m_out - u_out) + guide_scale['text'] * (y_out - m_out)
|
||||
elif isinstance(guide_scale, list):
|
||||
assert len(guide_scale) == len(model_kwargs) - 1
|
||||
y_out = model(xt, t=t, **model_kwargs[0], **kwargs)
|
||||
outs = [y_out]
|
||||
for i in range(1, len(model_kwargs)):
|
||||
outs.append(model(xt, t=t, **model_kwargs[i], **kwargs))
|
||||
out = outs[-1]
|
||||
for i in range(len(guide_scale)):
|
||||
out += guide_scale[i] * (outs[-i - 2] - outs[-i - 1])
|
||||
|
||||
# rescale the output according to arXiv:2305.08891
|
||||
if guide_rescale is not None and guide_rescale > 0.0:
|
||||
assert guide_rescale >= 0 and guide_rescale <= 1
|
||||
ratio = (
|
||||
y_out.flatten(1).std(dim=1) /
|
||||
(out.flatten(1).std(dim=1) + 1e-12)).view((-1, ) + (1, ) *
|
||||
(y_out.ndim - 1))
|
||||
out *= guide_rescale * ratio + (1 - guide_rescale) * 1.0
|
||||
# compute x0
|
||||
if self.prediction_type == 'x0':
|
||||
x0 = out
|
||||
@@ -197,6 +384,7 @@ class GaussianDiffusion(object):
|
||||
guide_scale=None,
|
||||
guide_rescale=None,
|
||||
clamp=None,
|
||||
sharpness=0.0,
|
||||
percentile=None,
|
||||
solver='euler_a',
|
||||
steps=20,
|
||||
@@ -209,12 +397,16 @@ class GaussianDiffusion(object):
|
||||
seed=-1,
|
||||
intermediate_callback=None,
|
||||
cat_uc=False,
|
||||
add_noise=False,
|
||||
free_steps=None,
|
||||
step_offset=None,
|
||||
**kwargs):
|
||||
# sanity check
|
||||
assert isinstance(steps, (int, torch.LongTensor))
|
||||
assert t_max is None or (t_max > 0 and t_max <= self.num_timesteps - 1)
|
||||
assert t_min is None or (t_min >= 0 and t_min < self.num_timesteps - 1)
|
||||
assert discretization in (None, 'leading', 'linspace', 'trailing')
|
||||
assert discretization in (None, 'leading', 'linspace', 'trailing',
|
||||
'free')
|
||||
assert discard_penultimate_step in (None, True, False)
|
||||
assert return_intermediate in (None, 'x0', 'xt')
|
||||
|
||||
@@ -255,17 +447,50 @@ class GaussianDiffusion(object):
|
||||
def model_fn(xt, sigma):
|
||||
# denoising
|
||||
t = self._sigma_to_t(sigma).repeat(len(xt)).round().long()
|
||||
x0 = self.denoise(xt,
|
||||
t,
|
||||
None,
|
||||
model,
|
||||
model_kwargs,
|
||||
guide_scale,
|
||||
guide_rescale,
|
||||
clamp,
|
||||
percentile,
|
||||
cat_uc=cat_uc,
|
||||
**kwargs)[-2]
|
||||
|
||||
if isinstance(
|
||||
model_kwargs[0]['cond'], dict) and \
|
||||
'tar_x0' in model_kwargs[0]['cond'] and \
|
||||
'tar_mask_latent' in model_kwargs[0]['cond']:
|
||||
tar_x0 = model_kwargs[0]['cond']['tar_x0']
|
||||
tar_mask = model_kwargs[0]['cond']['tar_mask_latent']
|
||||
|
||||
tar_xt = self.diffuse(x0=tar_x0, t=t)
|
||||
xt = tar_xt * (1.0 - tar_mask) + xt * tar_mask
|
||||
|
||||
if isinstance(model_kwargs[0]['cond'],
|
||||
dict) and 'ref_x0' in model_kwargs[0]['cond']:
|
||||
model_kwargs[0]['cond']['ref_xt'] = self.diffuse(
|
||||
x0=model_kwargs[0]['cond']['ref_x0'], t=t)
|
||||
model_kwargs[1]['cond']['ref_xt'] = self.diffuse(
|
||||
x0=model_kwargs[1]['cond']['ref_x0'], t=t)
|
||||
|
||||
if solver in ('onestep', 'multistep', 'multistep2', 'multistep3'):
|
||||
x0 = self.denoise(xt,
|
||||
t,
|
||||
None,
|
||||
model,
|
||||
model_kwargs,
|
||||
guide_scale,
|
||||
guide_rescale,
|
||||
clamp,
|
||||
sharpness,
|
||||
percentile,
|
||||
cat_uc=cat_uc,
|
||||
**kwargs)[-3]
|
||||
else:
|
||||
x0 = self.denoise(xt,
|
||||
t,
|
||||
None,
|
||||
model,
|
||||
model_kwargs,
|
||||
guide_scale,
|
||||
guide_rescale,
|
||||
clamp,
|
||||
sharpness,
|
||||
percentile,
|
||||
cat_uc=cat_uc,
|
||||
**kwargs)[-2]
|
||||
|
||||
# collect intermediate outputs
|
||||
if return_intermediate == 'xt':
|
||||
@@ -291,10 +516,14 @@ class GaussianDiffusion(object):
|
||||
elif discretization == 'trailing':
|
||||
steps = torch.arange(t_max, t_min - 1,
|
||||
-((t_max - t_min + 1) / steps))
|
||||
elif discretization == 'free':
|
||||
steps = torch.tensor(free_steps)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'{discretization} discretization not implemented')
|
||||
steps = steps.clamp_(t_min, t_max)
|
||||
elif isinstance(steps, list):
|
||||
steps = torch.tensor(steps)
|
||||
steps = torch.as_tensor(steps,
|
||||
dtype=torch.float32,
|
||||
device=noise.device)
|
||||
@@ -335,6 +564,23 @@ class GaussianDiffusion(object):
|
||||
if discard_penultimate_step:
|
||||
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
|
||||
kwargs['seed'] = seed
|
||||
# add noise to x0
|
||||
if add_noise:
|
||||
if 'dm_steps' in kwargs:
|
||||
if step_offset:
|
||||
add_noise_step = -kwargs['dm_steps'] + step_offset
|
||||
if add_noise_step < 0:
|
||||
noise = self.diffuse(
|
||||
noise,
|
||||
torch.full((noise.shape[0], 1),
|
||||
steps[add_noise_step],
|
||||
dtype=torch.int))
|
||||
else:
|
||||
noise = self.diffuse(
|
||||
noise,
|
||||
torch.full((noise.shape[0], 1),
|
||||
steps[-kwargs['dm_steps'] - 1],
|
||||
dtype=torch.int))
|
||||
# sampling
|
||||
x0 = solver_fn(noise,
|
||||
model_fn,
|
||||
@@ -373,168 +619,6 @@ class GaussianDiffusion(object):
|
||||
| torch.isinf(log_sigma)] = float('inf')
|
||||
return log_sigma.exp()
|
||||
|
||||
@torch.no_grad()
|
||||
def stochastic_encode(self, x0, t, steps):
|
||||
# fast, but does not allow for exact reconstruction
|
||||
# t serves as an index to gather the correct alphas
|
||||
|
||||
t_max = None
|
||||
t_min = None
|
||||
|
||||
# discretization method
|
||||
discretization = 'trailing' if self.prediction_type == 'v' else 'leading'
|
||||
|
||||
# timesteps
|
||||
if isinstance(steps, int):
|
||||
t_max = self.num_timesteps - 1 if t_max is None else t_max
|
||||
t_min = 0 if t_min is None else t_min
|
||||
steps = discretize_timesteps(t_max, t_min, steps, discretization)
|
||||
steps = torch.as_tensor(steps).round().long().flip(0).to(x0.device)
|
||||
# steps = torch.as_tensor(steps).round().long().to(x0.device)
|
||||
|
||||
# self.alphas_bar = torch.cumprod(1 - self.sigmas ** 2, dim=0)
|
||||
# print('sigma: ', self.sigmas, len(self.sigmas))
|
||||
# print('alpha_bar: ', self.alphas_bar, len(self.alphas_bar))
|
||||
# print('steps: ', steps, len(steps))
|
||||
# sqrt_alphas_cumprod = torch.sqrt(self.alphas_bar).to(x0.device)[steps]
|
||||
# sqrt_one_minus_alphas_cumprod = torch.sqrt(1 - self.alphas_bar).to(x0.device)[steps]
|
||||
|
||||
sqrt_alphas_cumprod = self.alphas.to(x0.device)[steps]
|
||||
sqrt_one_minus_alphas_cumprod = self.sigmas.to(x0.device)[steps]
|
||||
# print('sigma: ', self.sigmas, len(self.sigmas))
|
||||
# print('alpha: ', self.alphas, len(self.alphas))
|
||||
# print('steps: ', steps, len(steps))
|
||||
|
||||
noise = torch.randn_like(x0)
|
||||
return (
|
||||
extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
|
||||
extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) *
|
||||
noise)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_img2img(self,
|
||||
x,
|
||||
noise,
|
||||
model,
|
||||
denoising_strength=1,
|
||||
model_kwargs={},
|
||||
condition_fn=None,
|
||||
guide_scale=None,
|
||||
guide_rescale=None,
|
||||
clamp=None,
|
||||
percentile=None,
|
||||
solver='euler_a',
|
||||
steps=20,
|
||||
t_max=None,
|
||||
t_min=None,
|
||||
discretization=None,
|
||||
discard_penultimate_step=None,
|
||||
return_intermediate=None,
|
||||
show_progress=False,
|
||||
seed=-1,
|
||||
**kwargs):
|
||||
# sanity check
|
||||
assert isinstance(steps, (int, torch.LongTensor))
|
||||
assert t_max is None or (t_max > 0 and t_max <= self.num_timesteps - 1)
|
||||
assert t_min is None or (t_min >= 0 and t_min < self.num_timesteps - 1)
|
||||
assert discretization in (None, 'leading', 'linspace', 'trailing')
|
||||
assert discard_penultimate_step in (None, True, False)
|
||||
assert return_intermediate in (None, 'x0', 'xt')
|
||||
# function of diffusion solver
|
||||
solver_fn = {
|
||||
'euler_ancestral': sample_img2img_euler_ancestral,
|
||||
'euler': sample_img2img_euler,
|
||||
}[solver]
|
||||
# options
|
||||
schedule = 'karras' if 'karras' in solver else None
|
||||
discretization = discretization or 'linspace'
|
||||
seed = seed if seed >= 0 else random.randint(0, 2**31)
|
||||
if isinstance(steps, torch.LongTensor):
|
||||
discard_penultimate_step = False
|
||||
if discard_penultimate_step is None:
|
||||
discard_penultimate_step = True if solver in (
|
||||
'dpm2', 'dpm2_ancestral', 'dpmpp_2m_sde', 'dpm2_karras',
|
||||
'dpm2_ancestral_karras', 'dpmpp_2m_sde_karras') else False
|
||||
|
||||
# function for denoising xt to get x0
|
||||
intermediates = []
|
||||
|
||||
def get_scalings(sigma):
|
||||
c_out = -sigma
|
||||
c_in = 1 / (sigma**2 + 1.**2)**0.5
|
||||
return c_out, c_in
|
||||
|
||||
def model_fn(xt, sigma):
|
||||
# denoising
|
||||
c_out, c_in = get_scalings(sigma)
|
||||
t = self._sigma_to_t(sigma).repeat(len(xt)).round().long()
|
||||
|
||||
x0 = self.denoise(xt * c_in, t, None, model, model_kwargs,
|
||||
guide_scale, guide_rescale, clamp, percentile,
|
||||
**kwargs)[-2]
|
||||
# collect intermediate outputs
|
||||
if return_intermediate == 'xt':
|
||||
intermediates.append(xt)
|
||||
elif return_intermediate == 'x0':
|
||||
intermediates.append(x0)
|
||||
return xt + x0 * c_out
|
||||
|
||||
# get timesteps
|
||||
if isinstance(steps, int):
|
||||
steps += 1 if discard_penultimate_step else 0
|
||||
t_max = self.num_timesteps - 1 if t_max is None else t_max
|
||||
t_min = 0 if t_min is None else t_min
|
||||
# discretize timesteps
|
||||
if discretization == 'leading':
|
||||
steps = torch.arange(t_min, t_max + 1,
|
||||
(t_max - t_min + 1) / steps).flip(0)
|
||||
elif discretization == 'linspace':
|
||||
steps = torch.linspace(t_max, t_min, steps)
|
||||
elif discretization == 'trailing':
|
||||
steps = torch.arange(t_max, t_min - 1,
|
||||
-((t_max - t_min + 1) / steps))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'{discretization} discretization not implemented')
|
||||
steps = steps.clamp_(t_min, t_max)
|
||||
steps = torch.as_tensor(steps, dtype=torch.float32, device=x.device)
|
||||
# get sigmas
|
||||
sigmas = self._t_to_sigma(steps)
|
||||
sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
|
||||
t_enc = int(min(denoising_strength, 0.999) * len(steps))
|
||||
sigmas = sigmas[len(steps) - t_enc - 1:]
|
||||
noise = x + noise * sigmas[0]
|
||||
|
||||
if schedule == 'karras':
|
||||
if sigmas[0] == float('inf'):
|
||||
sigmas = karras_schedule(
|
||||
n=len(steps) - 1,
|
||||
sigma_min=sigmas[sigmas > 0].min().item(),
|
||||
sigma_max=sigmas[sigmas < float('inf')].max().item(),
|
||||
rho=7.).to(sigmas)
|
||||
sigmas = torch.cat([
|
||||
sigmas.new_tensor([float('inf')]), sigmas,
|
||||
sigmas.new_zeros([1])
|
||||
])
|
||||
else:
|
||||
sigmas = karras_schedule(
|
||||
n=len(steps),
|
||||
sigma_min=sigmas[sigmas > 0].min().item(),
|
||||
sigma_max=sigmas.max().item(),
|
||||
rho=7.).to(sigmas)
|
||||
sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
|
||||
if discard_penultimate_step:
|
||||
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
|
||||
|
||||
# sampling
|
||||
x0 = solver_fn(noise,
|
||||
model_fn,
|
||||
sigmas,
|
||||
seed=seed,
|
||||
show_progress=show_progress,
|
||||
**kwargs)
|
||||
return (x0, intermediates) if return_intermediate is not None else x0
|
||||
|
||||
|
||||
def extract_into_tensor(a, t, x_shape):
|
||||
b, *_ = t.shape
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torchvision.transforms.functional as TF
|
||||
|
||||
|
||||
def get_bbox_from_mask(mask):
|
||||
h, w = mask.shape[0], mask.shape[1]
|
||||
if mask.sum() < 10:
|
||||
return 0, h, 0, w
|
||||
rows = np.any(mask, axis=1)
|
||||
cols = np.any(mask, axis=0)
|
||||
y1, y2 = np.where(rows)[0][[0, -1]]
|
||||
x1, x2 = np.where(cols)[0][[0, -1]]
|
||||
return (y1, y2, x1, x2)
|
||||
|
||||
|
||||
def pad_to_square(image, pad_value=255, random=False):
|
||||
H, W = image.shape[0], image.shape[1]
|
||||
if H == W:
|
||||
return image, 0, 0
|
||||
|
||||
padd = abs(H - W)
|
||||
if random:
|
||||
padd_1 = int(np.random.randint(0, padd))
|
||||
else:
|
||||
padd_1 = int(padd / 2)
|
||||
padd_2 = padd - padd_1
|
||||
|
||||
if H > W:
|
||||
pad_param = ((0, 0), (padd_1, padd_2), (0, 0))
|
||||
else:
|
||||
pad_param = ((padd_1, padd_2), (0, 0), (0, 0))
|
||||
|
||||
# print(pad_param, pad_value)
|
||||
image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
|
||||
return image, padd_1, padd_2
|
||||
|
||||
|
||||
def box_in_box(small_box, big_box):
|
||||
y1, y2, x1, x2 = small_box
|
||||
y1_b, _, x1_b, _ = big_box
|
||||
y1, y2, x1, x2 = y1 - y1_b, y2 - y1_b, x1 - x1_b, x2 - x1_b
|
||||
return (y1, y2, x1, x2)
|
||||
|
||||
|
||||
def box2squre(image, box):
|
||||
H, W = image.shape[0], image.shape[1]
|
||||
y1, y2, x1, x2 = box
|
||||
cx = (x1 + x2) // 2
|
||||
cy = (y1 + y2) // 2
|
||||
h, w = y2 - y1, x2 - x1
|
||||
|
||||
if h >= w:
|
||||
x1 = cx - h // 2
|
||||
x2 = x1 + h
|
||||
else:
|
||||
y1 = cy - w // 2
|
||||
y2 = y1 + w
|
||||
x1 = max(0, x1)
|
||||
x2 = min(W, x2)
|
||||
y1 = max(0, y1)
|
||||
y2 = min(H, y2)
|
||||
return (y1, y2, x1, x2)
|
||||
|
||||
|
||||
def expand_bbox(mask,
|
||||
yyxx,
|
||||
ratio=1.0,
|
||||
min_crop=0,
|
||||
expand_type='center',
|
||||
to_square=False):
|
||||
y1, y2, x1, x2 = yyxx
|
||||
h = y2 - y1 + 1
|
||||
w = x2 - x1 + 1
|
||||
|
||||
H, W = mask.shape[0], mask.shape[1]
|
||||
xc, yc = 0.5 * (x1 + x2), 0.5 * (y1 + y2)
|
||||
|
||||
def expand(k):
|
||||
if isinstance(ratio, tuple) or isinstance(ratio, list):
|
||||
r = np.random.uniform(*ratio)
|
||||
k = k * r
|
||||
else:
|
||||
k = ratio * k
|
||||
return k
|
||||
|
||||
new_h = expand(h)
|
||||
new_w = expand(w)
|
||||
new_h = max(new_h, min_crop)
|
||||
new_w = max(new_w, min_crop)
|
||||
|
||||
if to_square:
|
||||
if new_w / new_h < 0.334:
|
||||
new_w = new_w + 1.0 / 3.0 * new_h
|
||||
elif new_h / new_w < 0.334:
|
||||
new_h = new_h + 1.0 / 3.0 * new_w
|
||||
|
||||
if expand_type == 'center':
|
||||
x1 = max(0, int(xc - new_w * 0.5))
|
||||
x2 = min(W, int(xc + new_w * 0.5))
|
||||
y1 = max(0, int(yc - new_h * 0.5))
|
||||
y2 = min(H, int(yc + new_h * 0.5))
|
||||
else:
|
||||
x1 = max(0, min(x1,
|
||||
int(x2 - new_w * np.random.uniform(w / new_w, 1.0))))
|
||||
x2 = min(W, max(x2, x1 + new_w))
|
||||
y1 = max(0, min(y1,
|
||||
int(y2 - new_h * np.random.uniform(h / new_h, 1.0))))
|
||||
y2 = min(H, max(y2, y1 + new_h))
|
||||
|
||||
return (int(y1), int(y2), int(x1), int(x2))
|
||||
|
||||
|
||||
def crop_back(pred, tar_image, extra_sizes, tar_box_yyxx_crop):
|
||||
H1, W1, H2, W2, pad1, pad2 = extra_sizes
|
||||
y1, y2, x1, x2 = tar_box_yyxx_crop
|
||||
pred = TF.resize(pred, (H2, W2), antialias=True)
|
||||
# if W1 == H1:
|
||||
# tar_image[:, y1:y2, x1:x2] = pred
|
||||
# return tar_image
|
||||
# if W1 < W2:
|
||||
# pad1 = int((W2 - W1) / 2)
|
||||
# pad2 = W2 - W1 - pad1
|
||||
# pred = pred[:, :,pad1:-pad2]
|
||||
# else:
|
||||
# pad1 = int((H2 - H1) / 2)
|
||||
# pad2 = H2 - H1 - pad1
|
||||
# pred = pred[:, pad1:-pad2, :]
|
||||
if W1 < W2:
|
||||
# pad width
|
||||
assert H1 == H2 and (pad1 + W1) == (W2 - pad2)
|
||||
pred = pred[:, :, pad1 + 2:(W2 - pad2 - 2)]
|
||||
tar_image[:, y1:y2, x1 + 2:x2 - 2] = pred
|
||||
elif H1 < H2:
|
||||
# pad height
|
||||
assert W1 == W2 and (pad1 + H1) == (H2 - pad2)
|
||||
pred = pred[:, pad1 + 2:(H2 - pad2 - 2), :]
|
||||
tar_image[:, y1 + 2:y2 - 2, x1:x2] = pred
|
||||
else:
|
||||
tar_image[:, y1:y2, x1:x2] = pred
|
||||
return tar_image
|
||||
|
||||
|
||||
def save_image(image, save_path):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
cv2.imwrite(save_path, image)
|
||||
@@ -1,7 +1,10 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from scepter.modules.opt.optimizers.official_optimizers import (
|
||||
ASGD, LBFGS, SGD, Adadelta, Adagrad, Adam, Adamax, AdamW, RMSprop, Rprop,
|
||||
SparseAdam)
|
||||
from scepter.modules.opt.optimizers.official_optimizers import (ASGD, LBFGS,
|
||||
SGD, Adadelta,
|
||||
Adagrad, Adam,
|
||||
Adamax, AdamW,
|
||||
RMSprop, Rprop,
|
||||
SparseAdam)
|
||||
from scepter.modules.opt.optimizers.registry import OPTIMIZERS
|
||||
|
||||
@@ -333,12 +333,12 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
data_iter = iter(self.datas[self._mode].dataloader)
|
||||
self.print_memory_status()
|
||||
for step in range(self.max_steps):
|
||||
if 'eval' in self._mode_set and (step % self.eval_interval == 0
|
||||
or step == self.max_steps - 1):
|
||||
if 'eval' in self._mode_set and (self.eval_interval > 0 and
|
||||
step % self.eval_interval == 0):
|
||||
self.run_eval()
|
||||
self.train_mode()
|
||||
self.before_iter(self.hooks_dict[self._mode])
|
||||
batch_data = next(data_iter)
|
||||
self.before_iter(self.hooks_dict[self._mode])
|
||||
if self.sample_args:
|
||||
batch_data.update(self.sample_args.get_lowercase_dict())
|
||||
if 'meta' in batch_data:
|
||||
@@ -364,6 +364,9 @@ class LatentDiffusionSolver(BaseSolver):
|
||||
self.after_iter(self.hooks_dict[self._mode])
|
||||
if we.debug:
|
||||
self.print_trainable_params_status(prefix='model.')
|
||||
if 'eval' in self._mode_set and (self.eval_interval > 0
|
||||
and step == self.max_steps - 1):
|
||||
self.run_eval()
|
||||
self.after_all_iter(self.hooks_dict[self._mode])
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
@@ -4,12 +4,14 @@
|
||||
from scepter.modules.solver.hooks.backward import BackwardHook
|
||||
from scepter.modules.solver.hooks.checkpoint import CheckpointHook
|
||||
from scepter.modules.solver.hooks.data_probe import ProbeDataHook
|
||||
from scepter.modules.solver.hooks.ema import ModelEmaHook
|
||||
from scepter.modules.solver.hooks.hook import Hook
|
||||
from scepter.modules.solver.hooks.log import LogHook, TensorboardLogHook
|
||||
from scepter.modules.solver.hooks.lr import LrHook
|
||||
from scepter.modules.solver.hooks.registry import HOOKS
|
||||
from scepter.modules.solver.hooks.safetensors import SafetensorsHook
|
||||
from scepter.modules.solver.hooks.sampler import DistSamplerHook
|
||||
|
||||
"""
|
||||
Normally, hooks have priorities, below we recommend priority that runs fine (low score MEANS high priority)
|
||||
BackwardHook: 0
|
||||
@@ -47,5 +49,6 @@ after solve:
|
||||
|
||||
__all__ = [
|
||||
'HOOKS', 'BackwardHook', 'CheckpointHook', 'Hook', 'LrHook', 'LogHook',
|
||||
'TensorboardLogHook', 'DistSamplerHook', 'ProbeDataHook', 'SafetensorsHook'
|
||||
'TensorboardLogHook', 'DistSamplerHook', 'ProbeDataHook',
|
||||
'SafetensorsHook', 'ModelEmaHook'
|
||||
]
|
||||
|
||||
@@ -147,7 +147,9 @@ class CheckpointHook(Hook):
|
||||
|
||||
if self.save_last and solver.total_iter == solver.max_steps - 1:
|
||||
with FS.get_fs_client(save_path) as client:
|
||||
last_path = osp.join(solver.work_dir, 'checkpoint.pth')
|
||||
last_path = osp.join(
|
||||
solver.work_dir,
|
||||
f'checkpoints/{self.save_name_prefix}-last')
|
||||
client.make_link(last_path, save_path)
|
||||
self.last_ckpt = save_path
|
||||
|
||||
|
||||
@@ -30,7 +30,10 @@ class ProbeDataHook(Hook):
|
||||
def __init__(self, cfg, logger=None):
|
||||
super(ProbeDataHook, self).__init__(cfg, logger=logger)
|
||||
self.priority = cfg.get('PRIORITY', _DEFAULT_PROBE_PRIORITY)
|
||||
self.log_interval = cfg.get('PROB_INTERVAL', 1000)
|
||||
self.prob_interval = cfg.get('PROB_INTERVAL', 1000)
|
||||
self.save_name_prefix = cfg.get('SAVE_NAME_PREFIX', 'step')
|
||||
self.save_probe_prefix = cfg.get('SAVE_PROBE_PREFIX', None)
|
||||
self.save_last = cfg.get('SAVE_LAST', False)
|
||||
|
||||
def before_all_iter(self, solver):
|
||||
pass
|
||||
@@ -39,19 +42,23 @@ class ProbeDataHook(Hook):
|
||||
pass
|
||||
|
||||
def after_iter(self, solver):
|
||||
if solver.mode == 'train' and solver.total_iter % self.log_interval == 0:
|
||||
if solver.mode == 'train' and solver.total_iter % self.prob_interval == 0:
|
||||
probe_dict = solver.probe_data
|
||||
if we.rank == 0:
|
||||
save_folder = os.path.join(
|
||||
solver.work_dir,
|
||||
f'{solver.mode}_probe/step_{solver.total_iter}')
|
||||
f'{solver.mode}_probe/{self.save_probe_prefix}-{solver.total_iter}'
|
||||
)
|
||||
ret_data = {}
|
||||
for k, v in probe_dict.items():
|
||||
ret_one = v.to_log(
|
||||
os.path.join(
|
||||
if self.save_probe_prefix is not None:
|
||||
ret_prefix = os.path.join(save_folder,
|
||||
self.save_probe_prefix)
|
||||
else:
|
||||
ret_prefix = os.path.join(
|
||||
save_folder,
|
||||
k.replace('/', '_') +
|
||||
f'_step_{solver.total_iter}'))
|
||||
k.replace('/', '_') + f'_step_{solver.total_iter}')
|
||||
ret_one = v.to_log(ret_prefix)
|
||||
if (isinstance(ret_one, list)
|
||||
or isinstance(ret_one, dict)) and len(ret_one) < 1:
|
||||
continue
|
||||
@@ -71,13 +78,19 @@ class ProbeDataHook(Hook):
|
||||
if we.rank == 0:
|
||||
step = solver._total_iter[
|
||||
'train'] if 'train' in solver._total_iter else 0
|
||||
save_folder = os.path.join(solver.work_dir,
|
||||
f'{solver.mode}_probe/step_{step}')
|
||||
save_folder = os.path.join(
|
||||
solver.work_dir,
|
||||
f'{solver.mode}_probe/{self.save_name_prefix}-{step}')
|
||||
ret_data = {}
|
||||
for k, v in probe_dict.items():
|
||||
ret_one = v.to_log(
|
||||
os.path.join(save_folder,
|
||||
k.replace('/', '_') + f'_step_{step}'))
|
||||
if self.save_probe_prefix is not None:
|
||||
ret_prefix = os.path.join(save_folder,
|
||||
self.save_probe_prefix)
|
||||
else:
|
||||
ret_prefix = os.path.join(
|
||||
save_folder,
|
||||
k.replace('/', '_') + f'_step_{step}')
|
||||
ret_one = v.to_log(ret_prefix)
|
||||
if (isinstance(ret_one, list)
|
||||
or isinstance(ret_one, dict)) and len(ret_one) < 1:
|
||||
continue
|
||||
@@ -87,6 +100,16 @@ class ProbeDataHook(Hook):
|
||||
json.dump(ret_data,
|
||||
open(local_path, 'w'),
|
||||
ensure_ascii=False)
|
||||
|
||||
if self.save_last and step == solver.max_steps:
|
||||
with FS.get_fs_client(save_folder) as client:
|
||||
last_save_folder = os.path.join(
|
||||
solver.work_dir,
|
||||
f'{solver.mode}_probe/{self.save_name_prefix}-last'
|
||||
)
|
||||
print(last_save_folder, save_folder)
|
||||
client.make_link(last_save_folder, save_folder)
|
||||
|
||||
solver.clear_probe()
|
||||
torch.cuda.synchronize()
|
||||
barrier()
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
from torch.distributed.fsdp import (FullStateDictConfig,
|
||||
FullyShardedDataParallel, StateDictType)
|
||||
|
||||
from scepter.modules.solver.hooks.hook import Hook
|
||||
from scepter.modules.solver.hooks.registry import HOOKS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.distribute import we
|
||||
|
||||
|
||||
@HOOKS.register_class()
|
||||
class ModelEmaHook(Hook):
|
||||
para_dict = [{
|
||||
'PRIORITY': {
|
||||
'value': 100,
|
||||
'description': 'the priority for processing!'
|
||||
},
|
||||
'BETA': {
|
||||
'value': 0.9999,
|
||||
'description': ''
|
||||
},
|
||||
}]
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super(ModelEmaHook, self).__init__(cfg, logger=logger)
|
||||
self.priority = cfg.get('PRIORITY', 100)
|
||||
self.beta = cfg.get('BETA', 0.9999)
|
||||
|
||||
def after_iter(self, solver):
|
||||
if solver.model.use_ema:
|
||||
model_ema = solver.model.model_ema
|
||||
model = solver.model.model
|
||||
self.ema(model_ema, model, use_fsdp=solver.use_fsdp)
|
||||
|
||||
@torch.no_grad()
|
||||
def ema(self, net_ema, net, use_fsdp=True):
|
||||
if we.is_distributed:
|
||||
if use_fsdp:
|
||||
save_policy = FullStateDictConfig(offload_to_cpu=False,
|
||||
rank0_only=False)
|
||||
with FullyShardedDataParallel.state_dict_type(
|
||||
net, StateDictType.FULL_STATE_DICT, save_policy):
|
||||
nonema_state = net.state_dict()
|
||||
elif hasattr(net, 'module'):
|
||||
nonema_state = net.module.state_dict()
|
||||
else:
|
||||
nonema_state = net.state_dict()
|
||||
else:
|
||||
nonema_state = net.state_dict()
|
||||
|
||||
for k, v in net_ema.named_parameters():
|
||||
v.copy_(nonema_state[k].lerp(v, self.beta))
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('hook',
|
||||
__class__.__name__,
|
||||
ModelEmaHook.para_dict,
|
||||
set_name=True)
|
||||
@@ -10,11 +10,16 @@ import torchvision.transforms as transforms
|
||||
import torchvision.transforms.functional as TF
|
||||
|
||||
from scepter.modules.transform.registry import TRANSFORMS
|
||||
from scepter.modules.transform.utils import (
|
||||
BACKEND_CV2, BACKEND_PILLOW, BACKEND_TORCHVISION, INPUT_CV2_TYPE_WARNING,
|
||||
INPUT_PIL_TYPE_WARNING, INPUT_TENSOR_TYPE_WARNING, INTERPOLATION_STYLE,
|
||||
INTERPOLATION_STYLE_CV2, TORCHVISION_CAPABILITY, is_cv2_image,
|
||||
is_pil_image, is_tensor)
|
||||
from scepter.modules.transform.utils import (BACKEND_CV2, BACKEND_PILLOW,
|
||||
BACKEND_TORCHVISION,
|
||||
INPUT_CV2_TYPE_WARNING,
|
||||
INPUT_PIL_TYPE_WARNING,
|
||||
INPUT_TENSOR_TYPE_WARNING,
|
||||
INTERPOLATION_STYLE,
|
||||
INTERPOLATION_STYLE_CV2,
|
||||
TORCHVISION_CAPABILITY,
|
||||
is_cv2_image, is_pil_image,
|
||||
is_tensor)
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
|
||||
if TORCHVISION_CAPABILITY:
|
||||
|
||||
@@ -6,13 +6,15 @@ import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from PIL import Image, ImageFile
|
||||
|
||||
from scepter.modules.transform.registry import TRANSFORMS
|
||||
from scepter.modules.utils.config import dict_to_yaml
|
||||
from scepter.modules.utils.distribute import we
|
||||
from scepter.modules.utils.file_system import DATA_FS as FS
|
||||
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
|
||||
|
||||
def pillow_convert(image, rgb_order):
|
||||
if image.mode != rgb_order:
|
||||
|
||||
@@ -603,3 +603,6 @@ class Config(object):
|
||||
return cfg_new
|
||||
else:
|
||||
return cfg
|
||||
|
||||
def pop(self, name):
|
||||
self.cfg_dict.pop(name)
|
||||
|
||||
@@ -4,10 +4,10 @@ import io
|
||||
from io import BytesIO
|
||||
|
||||
import onnx
|
||||
import onnxruntime
|
||||
import torch
|
||||
from torch.onnx import OperatorExportTypes
|
||||
|
||||
import onnxruntime
|
||||
from scepter.modules.utils.distribute import we
|
||||
|
||||
type_map = {
|
||||
|
||||
@@ -293,6 +293,9 @@ class LocalFs(BaseFs):
|
||||
return True
|
||||
|
||||
def get_logging_handler(self, target_logging_path):
|
||||
dirname = os.path.dirname(target_logging_path)
|
||||
if not os.path.exists(dirname):
|
||||
os.makedirs(dirname, exist_ok=True)
|
||||
return logging.FileHandler(target_logging_path)
|
||||
|
||||
def put_dir_from_local_dir(self,
|
||||
@@ -303,11 +306,11 @@ class LocalFs(BaseFs):
|
||||
target_dir = self.reconstruct_path(target_dir)
|
||||
if local_dir == target_dir:
|
||||
return True
|
||||
# cp -f local_dir/* target_dir/*
|
||||
if not osp.exists(target_dir):
|
||||
status = os.system(f'mkdir -p {target_dir}')
|
||||
if status != 0:
|
||||
return False
|
||||
# # cp -f local_dir/* target_dir/*
|
||||
# if not osp.exists(target_dir):
|
||||
# status = os.system(f'mkdir -p {target_dir}')
|
||||
# if status != 0:
|
||||
# return False
|
||||
try:
|
||||
shutil.copytree(local_dir, target_dir, symlinks=True)
|
||||
except Exception:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
import io
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
@@ -327,6 +328,10 @@ class FileSystem(object):
|
||||
target_path_list):
|
||||
if local_path is None or target_path is None:
|
||||
flg = False
|
||||
elif isinstance(local_path, io.BytesIO):
|
||||
flg = FS.put_object(local_path.getvalue(), target_path)
|
||||
elif isinstance(local_path, bytes):
|
||||
flg = FS.put_object(local_path, target_path)
|
||||
elif self.exists(local_path):
|
||||
local_cache = self.get_from(local_path,
|
||||
local_path + f'{time.time()}',
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os.path as osp
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from scepter.modules.utils.file_system import FS
|
||||
|
||||
|
||||
def init_1level_llfs(list_file,
|
||||
max_lines=1024,
|
||||
index_name='index',
|
||||
delimiter='\n'):
|
||||
r"""Construct large-list-file index.
|
||||
"""
|
||||
index_dir = osp.splitext(list_file)[0]
|
||||
print(list_file)
|
||||
with FS.get_from(list_file, wait_finish=True) as local_path:
|
||||
print(local_path)
|
||||
_num_split, index_stack, save_files = 0, [], []
|
||||
with open(local_path, 'r', buffering=1000000) as f:
|
||||
for line in tqdm(f):
|
||||
index_stack.append(line.strip())
|
||||
if len(index_stack) >= max_lines:
|
||||
save_file = f'{index_dir}/{index_name}/{_num_split + 1:09d}.txt'
|
||||
with FS.put_to(save_file) as cache_path:
|
||||
with open(cache_path, 'w') as f_w:
|
||||
f_w.write('\n'.join(index_stack))
|
||||
index_stack = []
|
||||
_num_split += 1
|
||||
save_files.append(save_file)
|
||||
|
||||
if len(index_stack) > 0:
|
||||
save_file = f'{index_dir}/{index_name}/{_num_split + 1:06d}.txt'
|
||||
with FS.put_to(save_file) as cache_path:
|
||||
with open(cache_path, 'w') as f_w:
|
||||
f_w.write('\n'.join(index_stack))
|
||||
save_files.append(save_file)
|
||||
|
||||
# output meta-file
|
||||
index_file = osp.join(index_dir, f'{index_name}.txt')
|
||||
with FS.put_to(index_file) as cache_path:
|
||||
with open(cache_path, 'w') as f_w:
|
||||
f_w.write('\n'.join(save_files))
|
||||
return index_file
|
||||
@@ -263,7 +263,8 @@ class ProbeData():
|
||||
url = FS.get_url(one_path,
|
||||
lifecycle=3600 * 365 * 24).replace(
|
||||
'.oss-internal.aliyun-inc.',
|
||||
'.oss.aliyuncs.')
|
||||
'.oss.aliyuncs.').replace(
|
||||
'-internal', '')
|
||||
one_rank += (
|
||||
f'<td align="center"><input type="image" src="{url}" >'
|
||||
f'<br><font size="4"><strong>{save_id}-{idx}|{one_label}<strong></font><br/></td>'
|
||||
|
||||
Reference in New Issue
Block a user