238 lines
9.1 KiB
Python
238 lines
9.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import numbers
|
|
import random
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from scepter.modules.model.network.diffusion.diffusion import \
|
|
GaussianDiffusionRF
|
|
from scepter.modules.model.network.diffusion.schedules import noise_schedule
|
|
from scepter.modules.model.network.ldm import LatentDiffusion
|
|
from scepter.modules.model.registry import MODELS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.distribute import we
|
|
|
|
|
|
@MODELS.register_class()
|
|
class LatentDiffusionSD3(LatentDiffusion):
|
|
para_dict = LatentDiffusion.para_dict
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
|
|
self.shift_factor = cfg.get('SHIFT_FACTOR', 0)
|
|
self.t_weight_type = cfg.get('T_WEIGHT', 'logit_normal')
|
|
self.logit_mean = cfg.get('LOGIT_MEAN', 0.0)
|
|
self.logit_std = cfg.get('LOGIT_STD', 1.0)
|
|
|
|
def init_params(self):
|
|
self.parameterization = self.cfg.get('PARAMETERIZATION', 'rf')
|
|
assert self.parameterization in [
|
|
'eps', 'x0', 'v', 'rf'
|
|
], 'currently only supporting "eps" and "x0" and "v" and "rf"'
|
|
self.num_timesteps = self.cfg.get('TIMESTEPS', 1000)
|
|
|
|
self.schedule_args = {
|
|
k.lower(): v
|
|
for k, v in self.cfg.get('SCHEDULE_ARGS', {
|
|
'NAME': 'logsnr_cosine_interp',
|
|
'SCALE_MIN': 2.0,
|
|
'SCALE_MAX': 4.0
|
|
}).items()
|
|
}
|
|
|
|
self.min_snr_gamma = self.cfg.get('MIN_SNR_GAMMA', None)
|
|
|
|
self.zero_terminal_snr = self.cfg.get('ZERO_TERMINAL_SNR', False)
|
|
if self.zero_terminal_snr:
|
|
assert self.parameterization == 'v', 'Now zero_terminal_snr only support v-prediction mode.'
|
|
|
|
self.sigmas = noise_schedule(schedule=self.schedule_args.pop('name'),
|
|
n=self.num_timesteps,
|
|
zero_terminal_snr=self.zero_terminal_snr,
|
|
**self.schedule_args)
|
|
|
|
self.diffusion = GaussianDiffusionRF(
|
|
sigmas=self.sigmas, prediction_type=self.parameterization)
|
|
|
|
self.pretrained_model = self.cfg.get('PRETRAINED_MODEL', None)
|
|
self.ignore_keys = self.cfg.get('IGNORE_KEYS', [])
|
|
|
|
self.model_config = self.cfg.DIFFUSION_MODEL
|
|
self.first_stage_config = self.cfg.FIRST_STAGE_MODEL
|
|
self.cond_stage_config = self.cfg.COND_STAGE_MODEL
|
|
self.tokenizer_config = self.cfg.get('TOKENIZER', None)
|
|
self.loss_config = self.cfg.get('LOSS', None)
|
|
|
|
self.scale_factor = self.cfg.get('SCALE_FACTOR', 0.18215)
|
|
self.size_factor = self.cfg.get('SIZE_FACTOR', 8)
|
|
self.default_n_prompt = self.cfg.get('DEFAULT_N_PROMPT', '')
|
|
self.default_n_prompt = '' if self.default_n_prompt is None else self.default_n_prompt
|
|
self.p_zero = self.cfg.get('P_ZERO', 0.0)
|
|
self.train_n_prompt = self.cfg.get('TRAIN_N_PROMPT', '')
|
|
if self.default_n_prompt is None:
|
|
self.default_n_prompt = ''
|
|
if self.train_n_prompt is None:
|
|
self.train_n_prompt = ''
|
|
self.use_ema = self.cfg.get('USE_EMA', False)
|
|
self.model_ema_config = self.cfg.get('DIFFUSION_MODEL_EMA', None)
|
|
|
|
def noise_sample(self, batch_size, h, w, g, c=4):
|
|
noise = torch.empty(batch_size, c, h, w,
|
|
device=we.device_id).normal_(generator=g)
|
|
return noise
|
|
|
|
def forward_train(self, image=None, noise=None, prompt=None, **kwargs):
|
|
n, c, h, w = image.shape
|
|
x_start = self.encode_first_stage(image, **kwargs)
|
|
if self.t_weight_type == 'uniform':
|
|
t = torch.randint(0,
|
|
self.num_timesteps, (n, ),
|
|
device=x_start.device).long()
|
|
elif self.t_weight_type == 'logit_normal':
|
|
density = F.sigmoid(
|
|
torch.normal(mean=self.logit_mean,
|
|
std=self.logit_std,
|
|
size=(n, ),
|
|
device=x_start.device))
|
|
t = (density * (self.num_timesteps - 1)).round().long()
|
|
sigma = (t + 1) / self.num_timesteps
|
|
shift = self.schedule_args['shift']
|
|
if shift > 1.:
|
|
sigma = shift * sigma / (1 + (shift - 1) * sigma)
|
|
t = sigma * self.num_timesteps
|
|
|
|
context = {}
|
|
if prompt and self.cond_stage_model:
|
|
ctx, pooled = getattr(self.cond_stage_model, 'encode')(prompt)
|
|
context['crossattn'] = ctx.float()
|
|
context['y'] = pooled
|
|
else:
|
|
assert False
|
|
|
|
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}}
|
|
return ret
|
|
|
|
@torch.no_grad()
|
|
def forward_test(self,
|
|
image=None,
|
|
prompt=None,
|
|
sampler='ddim',
|
|
sample_steps=20,
|
|
seed=2023,
|
|
guide_scale=4.5,
|
|
guide_rescale=0.0,
|
|
**kwargs):
|
|
g = torch.Generator(device=we.device_id)
|
|
seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
|
|
g.manual_seed(seed)
|
|
num_samples = len(prompt)
|
|
context = {}
|
|
null_context = {}
|
|
|
|
if prompt and self.cond_stage_model:
|
|
ctx, pooled = getattr(self.cond_stage_model, 'encode')(prompt)
|
|
null_ctx, null_pooled = getattr(self.cond_stage_model,
|
|
'encode')([''] * len(prompt))
|
|
context['crossattn'] = ctx.float()
|
|
context['y'] = pooled.float()
|
|
null_context['crossattn'] = null_ctx.float()
|
|
null_context['y'] = null_pooled.float()
|
|
else:
|
|
assert False
|
|
|
|
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 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,
|
|
c=16)
|
|
# UNet use input n_prompt
|
|
samples = self.diffusion.sample(
|
|
solver=sampler,
|
|
noise=noise,
|
|
model=self.model,
|
|
model_kwargs=[{
|
|
'cond': context
|
|
}, {
|
|
'cond': null_context
|
|
}] if guide_scale is not None and guide_scale > 0 else {
|
|
'cond': context,
|
|
},
|
|
cat_uc=False,
|
|
steps=sample_steps,
|
|
guide_scale=guide_scale,
|
|
guide_rescale=guide_rescale,
|
|
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)
|
|
outputs = list()
|
|
|
|
for i, (p, img) in enumerate(zip(prompt, x_samples)):
|
|
one_tup = {'prompt': str(p), 'n_prompt': '', 'image': img}
|
|
outputs.append(one_tup)
|
|
|
|
return outputs
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('MODEL',
|
|
__class__.__name__,
|
|
LatentDiffusionSD3.para_dict,
|
|
set_name=True)
|
|
|
|
@torch.no_grad()
|
|
def encode_first_stage(self, x, **kwargs):
|
|
z = self.first_stage_model.encode(x)
|
|
return self.scale_factor * (z - self.shift_factor)
|
|
|
|
@torch.no_grad()
|
|
def decode_first_stage(self, z):
|
|
z = 1. / self.scale_factor * z + self.shift_factor
|
|
return self.first_stage_model.decode(z)
|