# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import torch import torch.nn.functional as F import copy import math import random from contextlib import nullcontext from einops import rearrange from scepter.modules.model.network.ldm import LatentDiffusion from scepter.modules.model.registry import MODELS, DIFFUSIONS, BACKBONES, LOSSES, TOKENIZERS, EMBEDDERS from scepter.modules.model.utils.basic_utils import check_list_of_list, to_device, pack_imagelist_into_tensor, \ limit_batch_data, unpack_tensor_into_imagelist, count_params, disabled_train from scepter.modules.utils.config import dict_to_yaml from scepter.modules.utils.distribute import we @MODELS.register_class() class LatentDiffusionACEPlus(LatentDiffusion): para_dict = LatentDiffusion.para_dict def __init__(self, cfg, logger=None): super().__init__(cfg, logger=logger) self.guide_scale = cfg.get('GUIDE_SCALE', 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"' diffusion_cfg = self.cfg.get("DIFFUSION", None) assert diffusion_cfg is not None if self.cfg.have("WORK_DIR"): diffusion_cfg.WORK_DIR = self.cfg.WORK_DIR self.diffusion = DIFFUSIONS.build(diffusion_cfg, logger=self.logger) 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', 16) 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 construct_network(self): # embedding_context = torch.device("meta") if self.model_config.get("PRETRAINED_MODEL", None) else nullcontext() # with embedding_context: self.model = BACKBONES.build(self.model_config, logger=self.logger).to(torch.bfloat16) self.logger.info('all parameters:{}'.format(count_params(self.model))) if self.use_ema: if self.model_ema_config: self.model_ema = BACKBONES.build(self.model_ema_config, logger=self.logger) else: self.model_ema = copy.deepcopy(self.model) self.model_ema = self.model_ema.eval() for param in self.model_ema.parameters(): param.requires_grad = False if self.loss_config: self.loss = LOSSES.build(self.loss_config, logger=self.logger) if self.tokenizer_config is not None: self.tokenizer = TOKENIZERS.build(self.tokenizer_config, logger=self.logger) if self.first_stage_config: self.first_stage_model = MODELS.build(self.first_stage_config, logger=self.logger) self.first_stage_model = self.first_stage_model.eval() self.first_stage_model.train = disabled_train for param in self.first_stage_model.parameters(): param.requires_grad = False else: self.first_stage_model = None if self.tokenizer_config is not None: self.cond_stage_config.KWARGS = { 'vocab_size': self.tokenizer.vocab_size } if self.cond_stage_config == '__is_unconditional__': print( f'Training {self.__class__.__name__} as an unconditional model.' ) self.cond_stage_model = None else: model = EMBEDDERS.build(self.cond_stage_config, logger=self.logger) self.cond_stage_model = model.eval().requires_grad_(False) self.cond_stage_model.train = disabled_train @torch.no_grad() def encode_first_stage(self, x, **kwargs): def run_one_image(u): zu = self.first_stage_model.encode(u) if isinstance(zu, (tuple, list)): zu = zu[0] return zu z = [run_one_image(u.unsqueeze(0) if u.dim() == 3 else u) for u in x] return z @torch.no_grad() def decode_first_stage(self, z): return [self.first_stage_model.decode(zu) for zu in z] def noise_sample(self, num_samples, h, w, seed, dtype=torch.bfloat16): noise = torch.randn( num_samples, 16, # allow for packing 2 * math.ceil(h / 16), 2 * math.ceil(w / 16), device=we.device_id, dtype=dtype, generator=torch.Generator(device=we.device_id).manual_seed(seed), ) return noise def resize_func(self, x, size): if x is None: return x return F.interpolate(x.unsqueeze(0), size = size, mode='nearest-exact') def parse_ref_and_edit(self, src_image, modify_image, src_image_mask, text_embedding, #text_mask, edit_id): edit_image = [] modi_image = [] edit_mask = [] ref_image = [] ref_mask = [] ref_context = [] ref_y = [] ref_id = [] txt = [] txt_y = [] for sample_id, (one_src, one_modify, one_src_mask, one_text_embedding, one_text_y, # one_text_mask, one_edit_id) in enumerate(zip(src_image, modify_image, src_image_mask, text_embedding["context"], text_embedding["y"], #text_mask, edit_id) ): ref_id.append([i for i in range(len(one_src))]) if hasattr(self, "ref_cond_stage_model") and self.ref_cond_stage_model: ref_image.append(self.ref_cond_stage_model.encode_list([((i + 1.0) / 2.0 * 255).type(torch.uint8) for i in one_src])) else: ref_image.append(one_src) ref_mask.append(one_src_mask) # process edit image & edit image mask current_edit_image = to_device([one_src[i] for i in one_edit_id], strict=False) current_edit_image = [v.squeeze(0) for v in self.encode_first_stage(current_edit_image)] # process modi image current_modify_image = to_device([one_modify[i] for i in one_edit_id], strict=False) current_modify_image = [ v.squeeze(0) for v in self.encode_first_stage(current_modify_image) ] current_edit_image_mask = to_device( [one_src_mask[i] for i in one_edit_id], strict=False) current_edit_image_mask = [ self.reshape_func(m).squeeze(0) for m in current_edit_image_mask ] edit_image.append(current_edit_image) modi_image.append(current_modify_image) edit_mask.append(current_edit_image_mask) ref_context.append(one_text_embedding[:len(ref_id[-1])]) ref_y.append(one_text_y[:len(ref_id[-1])]) if not sum(len(src_) for src_ in src_image) > 0: ref_image = None ref_context = None ref_y = None for sample_id, (one_text_embedding, one_text_y) in enumerate(zip(text_embedding["context"], text_embedding["y"])): txt.append(one_text_embedding[-1].squeeze(0)) txt_y.append(one_text_y[-1]) return { "edit": edit_image, 'modify': modi_image, "edit_mask": edit_mask, "edit_id": edit_id, "ref_context": ref_context, "ref_y": ref_y, "context": txt, "y": txt_y, "ref_x": ref_image, "ref_mask": ref_mask, "ref_id": ref_id } def reshape_func(self, mask): mask = mask.to(torch.bfloat16) mask = mask.view((-1, mask.shape[-2], mask.shape[-1])) mask = rearrange( mask, "c (h ph) (w pw) -> c (ph pw) h w", ph=8, pw=8, ) return mask def forward_train(self, src_image_list=[], modify_image_list=[], src_mask_list=[], edit_id=[], image=None, image_mask=None, noise=None, prompt=[], **kwargs): ''' Args: src_image: list of list of src_image src_image_mask: list of list of src_image_mask image: target image image_mask: target image mask noise: default is None, generate automaticly ref_prompt: list of list of text prompt: list of text **kwargs: Returns: ''' assert check_list_of_list(src_image_list) and check_list_of_list( src_mask_list) assert self.cond_stage_model is not None gc_seg = kwargs.pop("gc_seg", []) gc_seg = int(gc_seg[0]) if len(gc_seg) > 0 else 0 align = kwargs.pop("align", []) prompt_ = [[pp] if isinstance(pp, str) else pp for pp in prompt] if len(align) < 1: align = [0] * len(prompt_) context = getattr(self.cond_stage_model, 'encode_list_of_list')(prompt_) guide_scale = self.guide_scale if guide_scale is not None: guide_scale = torch.full((len(prompt_),), guide_scale, device=we.device_id) else: guide_scale = None # image and image_mask # print("is list of list", check_list_of_list(image)) if check_list_of_list(image): image = [to_device(ix) for ix in image] x_start = [self.encode_first_stage(ix, **kwargs) for ix in image] noise = [[torch.randn_like(ii) for ii in ix] for ix in x_start] x_start = [torch.cat(ix, dim=-1) for ix in x_start] noise = [torch.cat(ix, dim=-1) for ix in noise] noise, _ = pack_imagelist_into_tensor(noise) image_mask = [to_device(im, strict=False) for im in image_mask] x_mask = [[self.reshape_func(i).squeeze(0) for i in im] if im is not None else [None] * len(ix) for ix, im in zip(image, image_mask)] x_mask = [torch.cat(im, dim=-1) for im in x_mask] else: image = to_device(image) x_start = self.encode_first_stage(image, **kwargs) image_mask = to_device(image_mask, strict=False) x_mask = [self.reshape_func(i).squeeze(0) for i in image_mask] if image_mask is not None else [None] * len( image) loss_mask, _ = pack_imagelist_into_tensor( tuple(torch.ones_like(ix, dtype=torch.bool, device=ix.device) for ix in x_start)) x_start, x_shapes = pack_imagelist_into_tensor(x_start) context['x_shapes'] = x_shapes context['align'] = align # process image mask context['x_mask'] = x_mask ref_edit_context = self.parse_ref_and_edit(src_image_list, modify_image_list, src_mask_list, context, edit_id) context.update(ref_edit_context) teacher_context = copy.deepcopy(context) teacher_context["context"] = torch.cat(teacher_context["context"], dim=0) teacher_context["y"] = torch.cat(teacher_context["y"], dim=0) loss = self.diffusion.loss(x_0=x_start, model=self.model, model_kwargs={"cond": context, "gc_seg": gc_seg, "guidance": guide_scale}, noise=noise, reduction='none', **kwargs) loss = loss[loss_mask].mean() ret = {'loss': loss, 'probe_data': {'prompt': prompt}} return ret @torch.no_grad() def forward_test(self, src_image_list=[], modify_image_list=[], src_mask_list=[], edit_id=[], image=None, image_mask=None, prompt=[], sampler='flow_euler', sample_steps=20, seed=2023, guide_scale=3.5, guide_rescale=0.0, show_process=False, log_num=-1, **kwargs): outputs = self.forward_editing( src_image_list=src_image_list, src_mask_list=src_mask_list, modify_image_list=modify_image_list, edit_id=edit_id, image=image, image_mask=image_mask, prompt=prompt, sampler=sampler, sample_steps=sample_steps, seed=seed, guide_scale=guide_scale, guide_rescale=guide_rescale, show_process=show_process, log_num=log_num, **kwargs ) return outputs @torch.no_grad() def forward_editing(self, src_image_list=[], modify_image_list=None, src_mask_list=[], edit_id=[], image=None, image_mask=None, prompt=[], sampler='flow_euler', sample_steps=20, seed=2023, guide_scale=3.5, log_num=-1, **kwargs ): # gc_seg is unused prompt, image, image_mask, src_image, modify_image, src_image_mask, edit_id = limit_batch_data( [prompt, image, image_mask, src_image_list, modify_image_list, src_mask_list, edit_id], log_num) assert check_list_of_list(src_image) and check_list_of_list(src_image_mask) assert self.cond_stage_model is not None align = kwargs.pop("align", []) prompt_ = [[pp] if isinstance(pp, str) else pp for pp in prompt] if len(align) < 1: align = [0] * len(prompt_) context = getattr(self.cond_stage_model, 'encode_list_of_list')(prompt_) guide_scale = guide_scale or self.guide_scale if guide_scale is not None: guide_scale = torch.full((len(prompt),), guide_scale, device=we.device_id) else: guide_scale = None # image and image_mask seed = seed if seed >= 0 else random.randint(0, 2 ** 32 - 1) if image is not None: if check_list_of_list(image): image = [torch.cat(ix, dim=-1) for ix in image] image_mask = [torch.cat(im, dim=-1) for im in image_mask] noise = [self.noise_sample(1, ix.shape[1], ix.shape[2], seed) for ix in image] else: height, width = kwargs.pop("height"), kwargs.pop("width") noise = [self.noise_sample(1, height, width, seed) for _ in prompt] noise, x_shapes = pack_imagelist_into_tensor(noise) context['x_shapes'] = x_shapes context['align'] = align # process image mask image_mask = to_device(image_mask, strict=False) x_mask = [self.reshape_func(i).squeeze(0) for i in image_mask] context['x_mask'] = x_mask ref_edit_context = self.parse_ref_and_edit(src_image, modify_image, src_image_mask, context, edit_id) context.update(ref_edit_context) # UNet use input n_prompt # model = self.model_ema if self.use_ema and self.eval_ema else self.model # import pdb;pdb.set_trace() model = self.model embedding_context = model.no_sync if isinstance(model, torch.distributed.fsdp.FullyShardedDataParallel) \ else nullcontext with embedding_context(): samples = self.diffusion.sample( noise=noise, sampler=sampler, model=self.model, model_kwargs={"cond": context, "guidance": guide_scale, "gc_seg": -1 }, steps=sample_steps, show_progress=True, guide_scale=guide_scale, return_intermediate=None, **kwargs).float() samples = unpack_tensor_into_imagelist(samples, x_shapes) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): x_samples = self.decode_first_stage(samples) outputs = list() for i in range(len(prompt)): rec_img = torch.clamp((x_samples[i].float() + 1.0) / 2.0, min=0.0, max=1.0) rec_img = rec_img.squeeze(0) edit_imgs, modify_imgs, edit_img_masks = [], [], [] if src_image is not None and src_image[i] is not None: if src_image_mask[i] is None: src_image_mask[i] = [None] * len(src_image[i]) for edit_img, modify_img, edit_mask in zip(src_image[i], modify_image_list[i], src_image_mask[i]): edit_img = torch.clamp((edit_img.float() + 1.0) / 2.0, min=0.0, max=1.0) edit_imgs.append(edit_img.squeeze(0)) modify_img = torch.clamp((modify_img.float() + 1.0) / 2.0, min=0.0, max=1.0) modify_imgs.append(modify_img.squeeze(0)) if edit_mask is None: edit_mask = torch.ones_like(edit_img[[0], :, :]) edit_img_masks.append(edit_mask) one_tup = { 'reconstruct_image': rec_img, 'instruction': prompt[i], 'edit_image': edit_imgs if len(edit_imgs) > 0 else None, 'modify_image': modify_imgs if len(modify_imgs) > 0 else None, 'edit_mask': edit_img_masks if len(edit_imgs) > 0 else None } if image is not None: if image_mask is None: image_mask = [None] * len(image) ori_img = torch.clamp((image[i] + 1.0) / 2.0, min=0.0, max=1.0) one_tup['target_image'] = ori_img.squeeze(0) one_tup['target_mask'] = image_mask[i] if image_mask[i] is not None else torch.ones_like( ori_img[[0], :, :]) outputs.append(one_tup) return outputs @staticmethod def get_config_template(): return dict_to_yaml('MODEL', __class__.__name__, LatentDiffusionACEPlus.para_dict, set_name=True)