351 lines
14 KiB
Python
351 lines
14 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
import os.path as osp
|
|
from typing import Any, Dict
|
|
|
|
import json
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models import Model
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.models.multi_modal.diffusion.diffusion import (
|
|
GaussianDiffusion, beta_schedule)
|
|
from modelscope.models.multi_modal.diffusion.structbert import (BertConfig,
|
|
BertModel)
|
|
from modelscope.models.multi_modal.diffusion.tokenizer import FullTokenizer
|
|
from modelscope.models.multi_modal.diffusion.unet_generator import \
|
|
DiffusionGenerator
|
|
from modelscope.models.multi_modal.diffusion.unet_upsampler_256 import \
|
|
SuperResUNet256
|
|
from modelscope.models.multi_modal.diffusion.unet_upsampler_1024 import \
|
|
SuperResUNet1024
|
|
from modelscope.utils.constant import ModelFile, Tasks
|
|
from modelscope.utils.device import create_device
|
|
from modelscope.utils.logger import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
__all__ = ['DiffusionForTextToImageSynthesis']
|
|
|
|
|
|
def make_diffusion(schedule,
|
|
num_timesteps=1000,
|
|
init_beta=None,
|
|
last_beta=None,
|
|
var_type='fixed_small'):
|
|
betas = beta_schedule(schedule, num_timesteps, init_beta, last_beta)
|
|
diffusion = GaussianDiffusion(betas, var_type=var_type)
|
|
return diffusion
|
|
|
|
|
|
class Tokenizer(object):
|
|
|
|
def __init__(self, vocab_file, seq_len=64):
|
|
self.vocab_file = vocab_file
|
|
self.seq_len = seq_len
|
|
self.tokenizer = FullTokenizer(
|
|
vocab_file=vocab_file, do_lower_case=True)
|
|
|
|
def __call__(self, text):
|
|
# tokenization
|
|
tokens = self.tokenizer.tokenize(text)
|
|
tokens = ['[CLS]'] + tokens[:self.seq_len - 2] + ['[SEP]']
|
|
input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
|
|
input_mask = [1] * len(input_ids)
|
|
segment_ids = [0] * len(input_ids)
|
|
|
|
# padding
|
|
input_ids += [0] * (self.seq_len - len(input_ids))
|
|
input_mask += [0] * (self.seq_len - len(input_mask))
|
|
segment_ids += [0] * (self.seq_len - len(segment_ids))
|
|
assert len(input_ids) == len(input_mask) == len(
|
|
segment_ids) == self.seq_len
|
|
|
|
# convert to tensors
|
|
input_ids = torch.LongTensor(input_ids)
|
|
input_mask = torch.LongTensor(input_mask)
|
|
segment_ids = torch.LongTensor(segment_ids)
|
|
return input_ids, segment_ids, input_mask
|
|
|
|
|
|
class DiffusionModel(nn.Module):
|
|
|
|
def __init__(self, model_dir):
|
|
super(DiffusionModel, self).__init__()
|
|
# including text and generator config
|
|
model_config = json.load(
|
|
open('{}/model_config.json'.format(model_dir), encoding='utf-8'))
|
|
|
|
# text encoder
|
|
text_config = model_config['text_config']
|
|
self.text_encoder = BertModel(BertConfig.from_dict(text_config))
|
|
|
|
# generator (64x64)
|
|
generator_config = model_config['generator_config']
|
|
self.unet_generator = DiffusionGenerator(**generator_config)
|
|
|
|
# upsampler (256x256)
|
|
upsampler_256_config = model_config['upsampler_256_config']
|
|
self.unet_upsampler_256 = SuperResUNet256(**upsampler_256_config)
|
|
|
|
# upsampler (1024x1024)
|
|
upsampler_1024_config = model_config['upsampler_1024_config']
|
|
self.unet_upsampler_1024 = SuperResUNet1024(**upsampler_1024_config)
|
|
|
|
def forward(self, noise, timesteps, input_ids, token_type_ids,
|
|
attention_mask):
|
|
context, y = self.text_encoder(
|
|
input_ids=input_ids,
|
|
token_type_ids=token_type_ids,
|
|
attention_mask=attention_mask)
|
|
context = context[-1]
|
|
x = self.unet_generator(noise, timesteps, y, context, attention_mask)
|
|
x = self.unet_upsampler_256(noise, timesteps, x,
|
|
torch.zeros_like(timesteps), y, context,
|
|
attention_mask)
|
|
x = self.unet_upsampler_1024(x, t, x)
|
|
return x
|
|
|
|
|
|
@MODELS.register_module(
|
|
Tasks.text_to_image_synthesis, module_name=Models.diffusion)
|
|
class DiffusionForTextToImageSynthesis(Model):
|
|
|
|
def __init__(self, model_dir, device='gpu', **kwargs):
|
|
device = 'gpu' if torch.cuda.is_available() else 'cpu'
|
|
super().__init__(model_dir=model_dir, device=device, **kwargs)
|
|
diffusion_model = DiffusionModel(model_dir=model_dir)
|
|
pretrained_params = torch.load(
|
|
osp.join(model_dir, ModelFile.TORCH_MODEL_BIN_FILE), 'cpu')
|
|
diffusion_model.load_state_dict(pretrained_params)
|
|
diffusion_model.eval().to()
|
|
|
|
self.device = create_device(device)
|
|
diffusion_model.to(self.device)
|
|
|
|
# modules
|
|
self.text_encoder = diffusion_model.text_encoder
|
|
self.unet_generator = diffusion_model.unet_generator
|
|
self.unet_upsampler_256 = diffusion_model.unet_upsampler_256
|
|
self.unet_upsampler_1024 = diffusion_model.unet_upsampler_1024
|
|
|
|
# text tokenizer
|
|
vocab_path = f'{model_dir}/{ModelFile.VOCAB_FILE}'
|
|
self.tokenizer = Tokenizer(vocab_file=vocab_path, seq_len=64)
|
|
|
|
# diffusion process
|
|
diffusion_params = json.load(
|
|
open(
|
|
'{}/diffusion_config.json'.format(model_dir),
|
|
encoding='utf-8'))
|
|
self.diffusion_generator = make_diffusion(
|
|
**diffusion_params['generator_config'])
|
|
self.diffusion_upsampler_256 = make_diffusion(
|
|
**diffusion_params['upsampler_256_config'])
|
|
self.diffusion_upsampler_1024 = make_diffusion(
|
|
**diffusion_params['upsampler_1024_config'])
|
|
|
|
def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:
|
|
if not all([key in input for key in ('text', 'noise', 'timesteps')]):
|
|
raise ValueError(
|
|
f'input should contains "text", "noise", and "timesteps", but got {input.keys()}'
|
|
)
|
|
input_ids, token_type_ids, attention_mask = self.tokenizer(
|
|
input['text'])
|
|
input_ids = input_ids.to(self.device).unsqueeze(0)
|
|
token_type_ids = token_type_ids.to(self.device).unsqueeze(0)
|
|
attention_mask = attention_mask.to(self.device).unsqueeze(0)
|
|
context, y = self.text_encoder(
|
|
input_ids=input_ids,
|
|
token_type_ids=token_type_ids,
|
|
attention_mask=attention_mask)
|
|
context = context[-1]
|
|
x = self.unet_generator(noise, timesteps, y, context, attention_mask)
|
|
x = self.unet_upsampler_256(noise, timesteps, x,
|
|
torch.zeros_like(timesteps), y, context,
|
|
attention_mask)
|
|
x = self.unet_upsampler_1024(x, t, x)
|
|
img = x.clamp(-1, 1).add(1).mul(127.5)
|
|
img = img.squeeze(0).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
|
|
return img
|
|
|
|
def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
return inputs
|
|
|
|
@torch.no_grad()
|
|
def generate(self, input: Dict[str, Any]) -> Dict[str, Any]:
|
|
if 'text' not in input:
|
|
raise ValueError(
|
|
f'input should contain "text", but got {input.keys()}')
|
|
|
|
# encode text
|
|
input_ids, token_type_ids, attention_mask = self.tokenizer(
|
|
input['text'])
|
|
input_ids = input_ids.to(self.device).unsqueeze(0)
|
|
token_type_ids = token_type_ids.to(self.device).unsqueeze(0)
|
|
attention_mask = attention_mask.to(self.device).unsqueeze(0)
|
|
context, y = self.text_encoder(
|
|
input_ids=input_ids,
|
|
token_type_ids=token_type_ids,
|
|
attention_mask=attention_mask)
|
|
context = context[-1]
|
|
|
|
# choose a proper solver
|
|
solver = input.get('solver', 'dpm-solver')
|
|
if solver == 'dpm-solver':
|
|
# generation
|
|
img = self.diffusion_generator.dpm_solver_sample_loop(
|
|
noise=torch.randn(1, 3, 64, 64).to(self.device),
|
|
model=self.unet_generator,
|
|
model_kwargs=[{
|
|
'y': y,
|
|
'context': context,
|
|
'mask': attention_mask
|
|
}, {
|
|
'y': torch.zeros_like(y),
|
|
'context': torch.zeros_like(context),
|
|
'mask': attention_mask
|
|
}],
|
|
percentile=input.get('generator_percentile', 0.995),
|
|
guide_scale=input.get('generator_guide_scale', 5.0),
|
|
dpm_solver_timesteps=input.get('dpm_solver_timesteps', 20),
|
|
order=3,
|
|
skip_type='logSNR',
|
|
method='singlestep',
|
|
t_start=0.9946)
|
|
|
|
# upsampling (64->256)
|
|
if not input.get('debug', False):
|
|
img = F.interpolate(
|
|
img,
|
|
scale_factor=4.0,
|
|
mode='bilinear',
|
|
align_corners=False)
|
|
img = self.diffusion_upsampler_256.dpm_solver_sample_loop(
|
|
noise=torch.randn_like(img),
|
|
model=self.unet_upsampler_256,
|
|
model_kwargs=[{
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': y,
|
|
'context': context,
|
|
'mask': attention_mask
|
|
}, {
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': torch.zeros_like(y),
|
|
'context': torch.zeros_like(context),
|
|
'mask': torch.zeros_like(attention_mask)
|
|
}],
|
|
percentile=input.get('upsampler_256_percentile', 0.995),
|
|
guide_scale=input.get('upsampler_256_guide_scale', 5.0),
|
|
dpm_solver_timesteps=input.get('dpm_solver_timesteps', 20),
|
|
order=3,
|
|
skip_type='logSNR',
|
|
method='singlestep',
|
|
t_start=0.9946)
|
|
|
|
# upsampling (256->1024)
|
|
if not input.get('debug', False):
|
|
img = F.interpolate(
|
|
img,
|
|
scale_factor=4.0,
|
|
mode='bilinear',
|
|
align_corners=False)
|
|
img = self.diffusion_upsampler_1024.dpm_solver_sample_loop(
|
|
noise=torch.randn_like(img),
|
|
model=self.unet_upsampler_256,
|
|
model_kwargs=[{
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': y,
|
|
'context': context,
|
|
'mask': attention_mask
|
|
}, {
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': torch.zeros_like(y),
|
|
'context': torch.zeros_like(context),
|
|
'mask': torch.zeros_like(attention_mask)
|
|
}],
|
|
percentile=input.get('upsampler_256_percentile', 0.995),
|
|
guide_scale=input.get('upsampler_256_guide_scale', 5.0),
|
|
dpm_solver_timesteps=input.get('dpm_solver_timesteps', 10),
|
|
order=3,
|
|
skip_type='logSNR',
|
|
method='singlestep',
|
|
t_start=None)
|
|
elif solver == 'ddim':
|
|
# generation
|
|
img = self.diffusion_generator.ddim_sample_loop(
|
|
noise=torch.randn(1, 3, 64, 64).to(self.device),
|
|
model=self.unet_generator,
|
|
model_kwargs=[{
|
|
'y': y,
|
|
'context': context,
|
|
'mask': attention_mask
|
|
}, {
|
|
'y': torch.zeros_like(y),
|
|
'context': torch.zeros_like(context),
|
|
'mask': attention_mask
|
|
}],
|
|
percentile=input.get('generator_percentile', 0.995),
|
|
guide_scale=input.get('generator_guide_scale', 5.0),
|
|
ddim_timesteps=input.get('generator_ddim_timesteps', 250),
|
|
eta=input.get('generator_ddim_eta', 0.0))
|
|
|
|
# upsampling (64->256)
|
|
if not input.get('debug', False):
|
|
img = F.interpolate(
|
|
img,
|
|
scale_factor=4.0,
|
|
mode='bilinear',
|
|
align_corners=False)
|
|
img = self.diffusion_upsampler_256.ddim_sample_loop(
|
|
noise=torch.randn_like(img),
|
|
model=self.unet_upsampler_256,
|
|
model_kwargs=[{
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': y,
|
|
'context': context,
|
|
'mask': attention_mask
|
|
}, {
|
|
'lx': img,
|
|
'lt': torch.zeros(1).to(self.device),
|
|
'y': torch.zeros_like(y),
|
|
'context': torch.zeros_like(context),
|
|
'mask': torch.zeros_like(attention_mask)
|
|
}],
|
|
percentile=input.get('upsampler_256_percentile', 0.995),
|
|
guide_scale=input.get('upsampler_256_guide_scale', 5.0),
|
|
ddim_timesteps=input.get('upsampler_256_ddim_timesteps', 50),
|
|
eta=input.get('upsampler_256_ddim_eta', 0.0))
|
|
|
|
# upsampling (256->1024)
|
|
if not input.get('debug', False):
|
|
img = F.interpolate(
|
|
img,
|
|
scale_factor=4.0,
|
|
mode='bilinear',
|
|
align_corners=False)
|
|
img = self.diffusion_upsampler_1024.ddim_sample_loop(
|
|
noise=torch.randn_like(img),
|
|
model=self.unet_upsampler_1024,
|
|
model_kwargs={'concat': img},
|
|
percentile=input.get('upsampler_1024_percentile', 0.995),
|
|
ddim_timesteps=input.get('upsampler_1024_ddim_timesteps', 20),
|
|
eta=input.get('upsampler_1024_ddim_eta', 0.0))
|
|
else:
|
|
raise ValueError(
|
|
'currently only supports "ddim" and "dpm-solve" solvers')
|
|
|
|
# output
|
|
img = img.clamp(-1, 1).add(1).mul(127.5).squeeze(0).permute(
|
|
1, 2, 0).cpu().numpy().astype(np.uint8)
|
|
return img
|