304 lines
12 KiB
Python
304 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import copy
|
|
import numbers
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
from scepter.modules.model.network.ldm.ldm import LatentDiffusion
|
|
from scepter.modules.model.registry import MODELS
|
|
from scepter.modules.model.utils.basic_utils import default
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.distribute import we
|
|
|
|
|
|
@MODELS.register_class()
|
|
class LatentDiffusionEdit(LatentDiffusion):
|
|
para_dict = {
|
|
'CONCAT_NO_SCALE_FACTOR': {
|
|
'value': False,
|
|
'description': 'Whether concat input scaled after VAE.'
|
|
},
|
|
'I_ZERO': {
|
|
'value': 0.0,
|
|
'description': 'P-zero of concat image.'
|
|
},
|
|
}
|
|
para_dict.update(LatentDiffusion.para_dict)
|
|
|
|
def __init__(self, cfg, logger):
|
|
super().__init__(cfg, logger=logger)
|
|
self.concat_no_scale_factor = self.cfg.get('CONCAT_NO_SCALE_FACTOR',
|
|
False)
|
|
self.i_zero = self.cfg.get('I_ZERO', 0.0)
|
|
# # overwrite the original diffusion
|
|
# self.diffusion = GaussianDiffusion(
|
|
# sigmas=self.sigmas, prediction_type=self.parameterization)
|
|
|
|
def forward_train(self, image=None, noise=None, prompt=None, **kwargs):
|
|
condition_attn = None
|
|
if 'condition_attn' in kwargs and kwargs['condition_attn'] is not None:
|
|
condition_attn = kwargs.pop('condition_attn')
|
|
if np.random.uniform() < self.i_zero:
|
|
condition_attn = torch.zeros_like(condition_attn)
|
|
|
|
condition_cat = None
|
|
if 'condition_cat' in kwargs and kwargs['condition_cat'] is not None:
|
|
condition_cat = kwargs.pop('condition_cat')
|
|
if np.random.uniform() < self.i_zero:
|
|
condition_cat = torch.zeros_like(condition_cat)
|
|
|
|
###############################
|
|
x_start = self.encode_first_stage(image, **kwargs)
|
|
t = torch.randint(0,
|
|
self.num_timesteps, (x_start.shape[0], ),
|
|
device=x_start.device).long()
|
|
|
|
context = {}
|
|
if prompt and self.cond_stage_model:
|
|
zeros = (torch.rand(len(prompt)) < self.p_zero).numpy().tolist()
|
|
prompt = [
|
|
self.train_n_prompt if zeros[idx] else p
|
|
for idx, p in enumerate(prompt)
|
|
]
|
|
self.register_probe({'after_prompt': prompt})
|
|
with torch.autocast(device_type='cuda', enabled=False):
|
|
if condition_attn is None:
|
|
context['crossattn'] = self.encode_condition(prompt)
|
|
else:
|
|
image_feature = self.encode_condition(condition_attn,
|
|
type='image')
|
|
context['crossattn'] = self.encode_condition(prompt,
|
|
image_feature,
|
|
type='hybrid')
|
|
if isinstance(context['crossattn'], dict):
|
|
attn = context.pop('crossattn')
|
|
context.update(attn)
|
|
|
|
if condition_cat is not None:
|
|
cat = self.encode_first_stage(condition_cat, **kwargs)
|
|
if self.concat_no_scale_factor:
|
|
cat /= self.scale_factor
|
|
context['concat'] = cat
|
|
|
|
if 'hint' in kwargs and kwargs['hint'] is not None:
|
|
hint = kwargs.pop('hint')
|
|
context['hint'] = hint
|
|
|
|
if self.min_snr_gamma is not None:
|
|
alphas = self.diffusion.alphas.to(we.device_id)[t]
|
|
sigmas = self.diffusion.sigmas.pow(2).to(we.device_id)[t]
|
|
snrs = (alphas / sigmas).clamp(min=1e-20)
|
|
min_snrs = snrs.clamp(max=self.min_snr_gamma)
|
|
weights = min_snrs / snrs
|
|
else:
|
|
weights = 1
|
|
self.register_probe({'snrs_weights': weights})
|
|
|
|
loss = self.diffusion.loss(x0=x_start,
|
|
t=t,
|
|
model=self.model,
|
|
model_kwargs={'cond': context},
|
|
noise=noise,
|
|
**kwargs)
|
|
loss = loss * weights
|
|
loss = loss.mean()
|
|
ret = {'loss': loss, 'probe_data': {'prompt': prompt}}
|
|
# ret = {'loss': loss, 'probe_data': {'prompt': prompt, 'concat': condition_cat, 'attn': condition_attn}}
|
|
return ret
|
|
|
|
@torch.no_grad()
|
|
@torch.autocast('cuda', dtype=torch.float16)
|
|
def forward_test(self,
|
|
image=None,
|
|
prompt=None,
|
|
n_prompt=None,
|
|
sampler='ddim',
|
|
sample_steps=50,
|
|
seed=2023,
|
|
guide_scale=7.5,
|
|
guide_rescale=0.5,
|
|
discretization='trailing',
|
|
run_train_n=True,
|
|
**kwargs):
|
|
condition_attn = None
|
|
if 'condition_attn' in kwargs and kwargs['condition_attn'] is not None:
|
|
condition_attn = kwargs.pop('condition_attn')
|
|
|
|
condition_cat = None
|
|
if 'condition_cat' in kwargs and kwargs['condition_cat'] is not None:
|
|
condition_cat = kwargs.pop('condition_cat')
|
|
|
|
image = None
|
|
if 'image' in kwargs and kwargs['image'] is not None:
|
|
image = kwargs.pop('image')
|
|
|
|
###############################
|
|
g = torch.Generator(device=we.device_id)
|
|
seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
|
|
g.manual_seed(seed)
|
|
# torch.manual_seed(seed)
|
|
num_samples = len(prompt)
|
|
|
|
# if 'dynamic_encode_text' in kwargs and kwargs.pop(
|
|
# 'dynamic_encode_text'):
|
|
# method = 'dynamic_encode_text'
|
|
# else:
|
|
# method = 'encode_text'
|
|
|
|
n_prompt = default(n_prompt, [self.default_n_prompt] * len(prompt))
|
|
assert isinstance(prompt, list) and \
|
|
isinstance(n_prompt, list) and \
|
|
len(prompt) == len(n_prompt)
|
|
|
|
context = {}
|
|
with torch.autocast(device_type='cuda', enabled=False):
|
|
if condition_attn is None:
|
|
context['crossattn'] = self.encode_condition(prompt)
|
|
else:
|
|
image_feature = self.encode_condition(condition_attn,
|
|
type='image')
|
|
context['crossattn'] = self.encode_condition(prompt,
|
|
image_feature,
|
|
type='hybrid')
|
|
if isinstance(context['crossattn'], dict):
|
|
attn = context.pop('crossattn')
|
|
context.update(attn)
|
|
|
|
null_context = {}
|
|
null_context['crossattn'] = self.encode_condition(n_prompt)
|
|
if isinstance(null_context['crossattn'], dict):
|
|
attn = null_context.pop('crossattn')
|
|
null_context.update(attn)
|
|
|
|
if 'hint' in kwargs and kwargs['hint'] is not None:
|
|
hint = kwargs.pop('hint')
|
|
context['hint'] = hint
|
|
null_context['hint'] = hint
|
|
else:
|
|
hint = None
|
|
|
|
model_kwargs = [{'cond': context}]
|
|
if condition_cat is not None:
|
|
cat = self.encode_first_stage(condition_cat, **kwargs)
|
|
if self.concat_no_scale_factor:
|
|
cat /= self.scale_factor
|
|
context['concat'] = cat
|
|
null_context['concat'] = torch.zeros_like(cat)
|
|
|
|
mid_context = {}
|
|
mid_context.update(null_context)
|
|
mid_context.update({'concat': cat})
|
|
model_kwargs.append({'cond': mid_context})
|
|
model_kwargs.append({'cond': null_context})
|
|
|
|
if 'index' in kwargs:
|
|
kwargs.pop('index')
|
|
image_size = None
|
|
|
|
if 'meta' in kwargs:
|
|
meta = kwargs.pop('meta')
|
|
if 'image_size' in meta:
|
|
h = int(meta['image_size'][0][0])
|
|
w = int(meta['image_size'][1][0])
|
|
image_size = [h, w]
|
|
if 'image_size' in kwargs:
|
|
image_size = kwargs.pop('image_size')
|
|
if condition_cat is not None:
|
|
image_size = condition_cat.shape[2:4]
|
|
if isinstance(image_size, numbers.Number):
|
|
image_size = [image_size, image_size]
|
|
if image_size is None:
|
|
image_size = [1024, 1024]
|
|
height, width = image_size
|
|
|
|
noise = self.noise_sample(num_samples, height // self.size_factor,
|
|
width // self.size_factor, g)
|
|
# UNet use input n_prompt
|
|
samples = self.diffusion.sample(solver=sampler,
|
|
noise=noise,
|
|
model=self.model,
|
|
model_kwargs=model_kwargs,
|
|
steps=sample_steps,
|
|
guide_scale=guide_scale,
|
|
guide_rescale=guide_rescale,
|
|
discretization=discretization,
|
|
show_progress=True,
|
|
seed=seed,
|
|
condition_fn=None,
|
|
clamp=None,
|
|
percentile=None,
|
|
t_max=None,
|
|
t_min=None,
|
|
discard_penultimate_step=None,
|
|
return_intermediate=None,
|
|
**kwargs)
|
|
x_samples = self.decode_first_stage(samples).float()
|
|
x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
|
|
if image is not None:
|
|
x_samples = torch.cat([image, x_samples], dim=-1)
|
|
if condition_attn is not None:
|
|
condition_attn = torch.clamp((condition_attn + 1.8) / 4,
|
|
min=0.0,
|
|
max=1.0)
|
|
x_samples = torch.cat([condition_attn, x_samples], dim=-1)
|
|
if condition_cat is not None:
|
|
condition_cat = torch.clamp((condition_cat + 1.0) / 2.0,
|
|
min=0.0,
|
|
max=1.0)
|
|
x_samples = torch.cat([condition_cat, x_samples], dim=-1)
|
|
# UNet use train n_prompt
|
|
# deleted!
|
|
train_n_prompt = ['' for _ in prompt]
|
|
t_x_samples = [None for _ in prompt]
|
|
|
|
outputs = list()
|
|
for i, (p, n_p, tnp, img, t_img) in enumerate(
|
|
zip(prompt, n_prompt, train_n_prompt, x_samples, t_x_samples)):
|
|
one_tup = {'prompt': p, 'n_prompt': n_p, 'image': img}
|
|
|
|
if hint is not None:
|
|
one_tup.update({'hint': hint[i]})
|
|
if t_img is not None:
|
|
one_tup['train_n_prompt'] = tnp
|
|
one_tup['train_n_image'] = t_img
|
|
outputs.append(one_tup)
|
|
|
|
return outputs
|
|
|
|
def encode_condition(self, data, data2=None, type='text'):
|
|
assert hasattr(self, 'tokenizer')
|
|
if type == 'image' and (
|
|
hasattr(self.cond_stage_model, 'build_new_tokens')
|
|
and not hasattr(self.cond_stage_model, 'new_tokens_to_ids')):
|
|
self.cond_stage_model.build_new_tokens(self.tokenizer)
|
|
|
|
if type == 'text':
|
|
text = self.tokenizer(data).to(we.device_id)
|
|
return self.cond_stage_model.encode_text(text)
|
|
elif type == 'image':
|
|
return self.cond_stage_model.encode_image(data)
|
|
elif type == 'hybrid':
|
|
text = self.tokenizer(data).to(we.device_id)
|
|
return self.cond_stage_model.encode_text(text, data2)
|
|
|
|
def save_pretrained(self,
|
|
*args,
|
|
destination=None,
|
|
prefix='',
|
|
keep_vars=False):
|
|
return self.model.state_dict(*args,
|
|
destination=destination,
|
|
keep_vars=keep_vars)
|
|
|
|
def save_pretrained_config(self):
|
|
return copy.deepcopy(self.cfg.COND_STAGE_MODEL.cfg_dict)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('MODELS',
|
|
__class__.__name__,
|
|
LatentDiffusionEdit.para_dict,
|
|
set_name=True)
|