import os import pathlib from PIL import Image from functools import partial import torch import torch.nn.functional as F from torchvision import transforms from torchvision.datasets.folder import is_image_file from tqdm.auto import tqdm import numpy as np from skimage.color import rgb2gray import diffusers from libs.engine import ModelState from libs.metric.lpips_origin import LPIPS from libs.metric.piq.perceptual import DISTS as DISTS_PIQ from libs.metric.clip_score import CLIPScoreWrapper from methods.painter.diffsketchedit import ( Painter, SketchPainterOptimizer, Token2AttnMixinASDSPipeline, Token2AttnMixinASDSSDXLPipeline) from methods.painter.diffsketchedit.sketch_utils import ( log_tensor_img, plt_batch, plt_attn, save_tensor_img, fix_image_scale) from methods.painter.diffsketchedit.mask_utils import get_mask_u2net from methods.token2attn.attn_control import AttentionStore, EmptyControl, \ LocalBlend, AttentionReplace, AttentionRefine, AttentionReweight, get_equalizer from methods.token2attn.ptp_utils import view_images, get_word_inds from methods.diffusers_warp import init_diffusion_pipeline, model2res from methods.diffvg_warp import init_diffvg from methods.painter.diffsketchedit.process_svg import remove_low_opacity_paths class DiffSketchEditPipeline(ModelState): def __init__(self, args): super().__init__(args, ignore_log=True) init_diffvg(self.device, True, args.print_timing) if args.model_id == "sdxl": # default LSDSSDXLPipeline scheduler is EulerDiscreteScheduler # when LSDSSDXLPipeline calls, scheduler.timesteps will change in step 4 # which causes problem in sds add_noise() function # because the random t may not in scheduler.timesteps custom_pipeline = Token2AttnMixinASDSSDXLPipeline custom_scheduler = diffusers.DPMSolverMultistepScheduler self.args.cross_attn_res = self.args.cross_attn_res * 2 elif args.model_id == 'sd21': custom_pipeline = Token2AttnMixinASDSPipeline custom_scheduler = diffusers.DDIMScheduler elif args.model_id == 'sd15': custom_pipeline = Token2AttnMixinASDSPipeline custom_scheduler = diffusers.DDIMScheduler else: # sd14 custom_pipeline = Token2AttnMixinASDSPipeline custom_scheduler = None self.diffusion = init_diffusion_pipeline( self.args.model_id, custom_pipeline=custom_pipeline, custom_scheduler=custom_scheduler, device=self.device, local_files_only=not args.download, force_download=args.force_download, resume_download=args.resume_download, ldm_speed_up=args.ldm_speed_up, enable_xformers=args.enable_xformers, gradient_checkpoint=args.gradient_checkpoint, ) # init clip model and clip score wrapper self.cargs = self.args.clip self.clip_score_fn = CLIPScoreWrapper(self.cargs.model_name, device=self.device, visual_score=True, feats_loss_type=self.cargs.feats_loss_type, feats_loss_weights=self.cargs.feats_loss_weights, fc_loss_weight=self.cargs.fc_loss_weight) def update_info(self, seed, token_ind, prompt_input): prompt_dir_name = prompt_input.split(' ') prompt_dir_name = '_'.join(prompt_dir_name) attn_log_ = f"-tk{token_ind}" logdir_ = f"seed{seed}" \ f"{attn_log_}" \ f"-stage={self.args.run_stage}" logdir_sec_ = f"" self.args.path_svg = "" if self.args.run_stage > 0: logdir_sec_ = f"{logdir_sec_}-local={self.args.vector_local_edit}" last_svg_base = os.path.join(self.args.results_path, self.args.edit_type, prompt_dir_name, logdir_[:-1] + str(self.args.run_stage - 1)) if self.args.run_stage != 1: last_svg_base += logdir_sec_ self.args.path_svg = os.path.join(last_svg_base, "visual_best.svg") self.args.attention_init = False logdir_ = f"{prompt_dir_name}" + f"/" + logdir_ + logdir_sec_ super().__init__(self.args, log_path_suffix=logdir_) # create log dir self.png_logs_dir = self.results_path / "png_logs" self.svg_logs_dir = self.results_path / "svg_logs" self.attn_logs_dir = self.results_path / "attn_logs" if self.accelerator.is_main_process: self.png_logs_dir.mkdir(parents=True, exist_ok=True) self.svg_logs_dir.mkdir(parents=True, exist_ok=True) self.attn_logs_dir.mkdir(parents=True, exist_ok=True) self.g_device = torch.Generator().manual_seed(seed) def load_render(self, target_img, attention_map, mask=None): renderer = Painter(self.args, num_strokes=self.args.num_paths, num_segments=self.args.num_segments, imsize=self.args.image_size, device=self.device, target_im=target_img, attention_map=attention_map, mask=mask) return renderer def attn_map_normalizing(self, cross_attn_map): cross_attn_map = 255 * cross_attn_map / cross_attn_map.max() # [res, res, 3] cross_attn_map = cross_attn_map.unsqueeze(-1).expand(*cross_attn_map.shape, 3) # [3, res, res] cross_attn_map = cross_attn_map.permute(2, 0, 1).unsqueeze(0) # [3, clip_size, clip_size] cross_attn_map = F.interpolate(cross_attn_map, size=self.args.image_size, mode='bicubic') cross_attn_map = torch.clamp(cross_attn_map, min=0, max=255) # rgb to gray cross_attn_map = rgb2gray(cross_attn_map.squeeze(0).permute(1, 2, 0)).astype(np.float32) # torch to numpy if cross_attn_map.shape[-1] != self.args.image_size and cross_attn_map.shape[-2] != self.args.image_size: cross_attn_map = cross_attn_map.reshape(self.args.image_size, self.args.image_size) # to [0, 1] cross_attn_map = (cross_attn_map - cross_attn_map.min()) / (cross_attn_map.max() - cross_attn_map.min()) return cross_attn_map def compute_local_edit_maps(self, cross_attn_maps_src_tar, prompts, words, save_path, threshold=0.3): """ cross_attn_maps_src_tar: [(res, res, 77), (res, res, 77)] """ local_edit_region = np.zeros(shape=(self.args.image_size, self.args.image_size), dtype=np.float32) for i, (prompt, word) in enumerate(zip(prompts, words)): ind = get_word_inds(prompt, word, self.diffusion.tokenizer) # list assert len(ind) == 1 ind = ind[0] cross_attn_map = cross_attn_maps_src_tar[i][:, :, ind] # (res, res) cross_attn_map = self.attn_map_normalizing(cross_attn_map) # (image_size, image_size), [0.0, 1.0] cross_attn_map_bin = cross_attn_map >= threshold local_edit_region += cross_attn_map_bin local_edit_region = (np.clip(local_edit_region, 0, 1) * 255).astype(np.uint8) local_edit_region = Image.fromarray(local_edit_region, 'L') local_edit_region.save(save_path, 'PNG') def extract_ldm_attn(self, prompts, token_ind, changing_region_words, reweight_word, reweight_weight): ######################### Change here for editing methods ######################### ## init controller if not self.args.attention_init: controller = EmptyControl() else: lb = LocalBlend(prompts=prompts, words=changing_region_words, tokenizer=self.diffusion.tokenizer, device=self.device) # changing region # if self.args.edit_type == "none": # controller = AttentionStore() if self.args.edit_type == "replace": controller = AttentionReplace(prompts=prompts, num_steps=self.args.num_inference_steps, cross_replace_steps=0.4, # larger is more similar shape self_replace_steps=0.4, local_blend=lb, tokenizer=self.diffusion.tokenizer, device=self.device) elif self.args.edit_type == "refine": controller = AttentionRefine(prompts=prompts, num_steps=self.args.num_inference_steps, cross_replace_steps=0.8, # larger is more similar shape self_replace_steps=0.4, local_blend=lb, tokenizer=self.diffusion.tokenizer, device=self.device) elif self.args.edit_type == "reweight": equalizer = get_equalizer(self.diffusion.tokenizer, prompts[1:], reweight_word, reweight_weight) controller = AttentionReweight(prompts=prompts, num_steps=self.args.num_inference_steps, cross_replace_steps=0.8, # larger is more similar shape self_replace_steps=0.4, local_blend=lb, equalizer=equalizer, # controller=controller_a, tokenizer=self.diffusion.tokenizer, device=self.device) else: raise Exception('Unknown edit_type:', self.args.edit_type) ######################### Change here for editing methods (end) ######################### height = width = model2res(self.args.model_id) outputs = self.diffusion(prompt=prompts, negative_prompt=[self.args.negative_prompt] * len(prompts), height=height, width=width, controller=controller, num_inference_steps=self.args.num_inference_steps, guidance_scale=self.args.guidance_scale, generator=self.g_device) print('outputs.images', len(outputs.images)) for ii, img in enumerate(outputs.images): if ii == 0: filename = "ldm_generated_image.png" target_file = self.results_path / filename else: filename = "ldm_generated_image" + str(ii) + ".png" target_file_tmp = self.results_path / filename view_images([np.array(img)], save_image=True, fp=target_file_tmp) if self.args.attention_init: """ldm cross-attention map""" cross_attention_maps, tokens = \ self.diffusion.get_cross_attention(prompts, controller, res=self.args.cross_attn_res, from_where=("up", "down"), save_path=self.results_path / "cross_attn.png", select=0) for ii in range(1, len(outputs.images)): cross_attn_png_name = "cross_attn" + str(ii) + ".png" cross_attention_maps_i, tokens_i = \ self.diffusion.get_cross_attention(prompts, controller, res=self.args.cross_attn_res, from_where=("up", "down"), save_path=self.results_path / cross_attn_png_name, select=ii) self.print(f"the length of tokens is {len(tokens)}, select {token_ind}-th token") # [res, res, seq_len] self.print(f"origin cross_attn_map shape: {cross_attention_maps.shape}") # [res, res] cross_attn_map = cross_attention_maps[:, :, token_ind] self.print(f"select cross_attn_map shape: {cross_attn_map.shape}\n") cross_attn_map = self.attn_map_normalizing(cross_attn_map) ######################### ldm cross-attention map (for vector local editing) ######################### cross_attention_maps_local_list = [] for ii in range(len(outputs.images)): cross_attention_maps_local = \ self.diffusion.get_cross_attention2(prompts, controller, res=self.args.vector_local_edit_attn_res, from_where=("up", "down"), select=ii) # (res, res, 77) cross_attention_maps_local_list.append(cross_attention_maps_local) if ii == 0: continue save_name = "cross_attn_local_edit_" + str(self.args.vector_local_edit_attn_res) + "-" + str(ii) + ".png" if self.args.edit_type == "replace": self.compute_local_edit_maps([cross_attention_maps_local_list[ii-1]], [prompts[ii-1]], [changing_region_words[ii][0]], save_path=self.results_path / save_name, threshold=self.args.vector_local_edit_bin_threshold_replace) elif self.args.edit_type == "refine": self.compute_local_edit_maps([cross_attention_maps_local_list[ii]], [prompts[ii]], [changing_region_words[ii][1]], save_path=self.results_path / save_name, threshold=self.args.vector_local_edit_bin_threshold_refine) elif self.args.edit_type == "reweight": self.compute_local_edit_maps([cross_attention_maps_local_list[ii-1]], [prompts[ii-1]], [changing_region_words[ii][0]], save_path=self.results_path / save_name, threshold=self.args.vector_local_edit_bin_threshold_reweight) if self.args.sd_image_only: return target_file.as_posix(), None ######################### ######################### """ldm self-attention map""" self_attention_maps, svd, vh_ = \ self.diffusion.get_self_attention_comp(prompts, controller, res=self.args.self_attn_res, from_where=("up", "down"), img_size=self.args.image_size, max_com=self.args.max_com, save_path=self.results_path) # comp self-attention map if self.args.mean_comp: self_attn = np.mean(vh_, axis=0) self.print(f"use the mean of {self.args.max_com} comps.") else: self_attn = vh_[self.args.comp_idx] self.print(f"select {self.args.comp_idx}-th comp.") # to [0, 1] self_attn = (self_attn - self_attn.min()) / (self_attn.max() - self_attn.min()) # visual final self-attention self_attn_vis = np.copy(self_attn) self_attn_vis = self_attn_vis * 255 self_attn_vis = np.repeat(np.expand_dims(self_attn_vis, axis=2), 3, axis=2).astype(np.uint8) view_images(self_attn_vis, save_image=True, fp=self.results_path / "self-attn-final.png") """attention map fusion""" attn_map = self.args.attn_coeff * cross_attn_map + (1 - self.args.attn_coeff) * self_attn # to [0, 1] attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min()) self.print(f"-> fusion attn_map: {attn_map.shape}") else: attn_map = None return target_file.as_posix(), attn_map @property def clip_norm_(self): return transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)) def clip_pair_augment(self, x: torch.Tensor, y: torch.Tensor, im_res: int, augments: str = "affine_norm", num_aug: int = 4): # init augmentations augment_list = [] if "affine" in augments: augment_list.append( transforms.RandomPerspective(fill=0, p=1.0, distortion_scale=0.5) ) augment_list.append( transforms.RandomResizedCrop(im_res, scale=(0.8, 0.8), ratio=(1.0, 1.0)) ) augment_list.append(self.clip_norm_) # CLIP Normalize # compose augmentations augment_compose = transforms.Compose(augment_list) # make augmentation pairs x_augs, y_augs = [self.clip_score_fn.normalize(x)], [self.clip_score_fn.normalize(y)] # repeat N times for n in range(num_aug): augmented_pair = augment_compose(torch.cat([x, y])) x_augs.append(augmented_pair[0].unsqueeze(0)) y_augs.append(augmented_pair[1].unsqueeze(0)) xs = torch.cat(x_augs, dim=0) ys = torch.cat(y_augs, dim=0) return xs, ys def painterly_rendering(self, prompts, token_ind, changing_region_words, reweight_word, reweight_weight): # log prompts self.print(f"prompts: {prompts}") self.print(f"negative_prompt: {self.args.negative_prompt}") self.print(f"token_ind: {token_ind}") self.print(f"changing_region_words: {changing_region_words}") self.print(f"reweight_word: {reweight_word}") self.print(f"reweight_weight: {reweight_weight}\n") if self.args.negative_prompt is None: self.args.negative_prompt = "" log_path = os.path.join(self.results_path.as_posix(), 'log.txt') with open(log_path, "w") as f: f.write("prompts: " + str(prompts) + "\n") f.write("negative_prompt: " + self.args.negative_prompt + "\n") f.write("token_ind: " + str(token_ind) + "\n") f.write("changing_region_words: " + str(changing_region_words) + "\n") f.write("reweight_word: " + str(reweight_word) + "\n") f.write("reweight_weight: " + str(reweight_weight) + "\n") f.close() # init attention if self.args.run_stage == 0: target_file, attention_map = self.extract_ldm_attn(prompts, token_ind, changing_region_words, reweight_word, reweight_weight) else: results_base = self.results_path.as_posix() target_file = os.path.join(results_base[:results_base.find('stage=' + str(self.args.run_stage))] + 'stage=0', "ldm_generated_image" + str(self.args.run_stage) + ".png") attention_map = None if not self.args.sd_image_only: # timesteps_ = self.diffusion.scheduler.timesteps.cpu().numpy().tolist() # self.print(f"{len(timesteps_)} denoising steps, {timesteps_}") perceptual_loss_fn = None if self.args.perceptual.coeff > 0: if self.args.perceptual.name == "lpips": lpips_loss_fn = LPIPS(net=self.args.perceptual.lpips_net).to(self.device) perceptual_loss_fn = partial(lpips_loss_fn.forward, return_per_layer=False, normalize=False) elif self.args.perceptual.name == "dists": perceptual_loss_fn = DISTS_PIQ() inputs, mask = self.get_target(target_file, self.args.image_size, self.results_path, self.args.u2net_path, self.args.mask_object, self.args.fix_scale, self.device) inputs = inputs.detach() # inputs as GT self.print("inputs shape: ", inputs.shape) # load renderer renderer = Painter(self.args, num_strokes=self.args.num_paths, num_segments=self.args.num_segments, imsize=self.args.image_size, device=self.device, target_im=inputs, attention_map=attention_map, mask=mask, results_base=self.results_path.as_posix()) # init img img = renderer.init_image(stage=0) self.print("init_image shape: ", img.shape) log_tensor_img(img, self.results_path, output_prefix="init_sketch") # load optimizer optimizer = SketchPainterOptimizer(renderer, self.args.lr, self.args.optim_opacity, self.args.optim_rgba, self.args.color_lr, self.args.optim_width, self.args.width_lr) optimizer.init_optimizers() # log params self.print(f"-> Painter points Params: {len(renderer.get_points_params())}") self.print(f"-> Painter width Params: {len(renderer.get_width_parameters())}") self.print(f"-> Painter opacity Params: {len(renderer.get_color_parameters())}") best_visual_loss, best_semantic_loss = 100, 100 best_iter_v, best_iter_s = 0, 0 min_delta = 1e-6 vid_idx = 1 self.print(f"\ntotal optimization steps: {self.args.num_iter}") with tqdm(initial=self.step, total=self.args.num_iter, disable=not self.accelerator.is_main_process) as pbar: while self.step < self.args.num_iter: raster_sketch = renderer.get_image().to(self.device) target_prompt = prompts[self.args.run_stage] # ASDS loss sds_loss, grad = torch.tensor(0), torch.tensor(0) if self.step >= self.args.sds.warmup: grad_scale = self.args.sds.grad_scale if self.step > self.args.sds.warmup else 0 sds_loss, grad = self.diffusion.score_distillation_sampling( raster_sketch, crop_size=self.args.sds.crop_size, augments=self.args.sds.augmentations, prompt=[target_prompt], negative_prompt=[self.args.negative_prompt], guidance_scale=self.args.sds.guidance_scale, grad_scale=grad_scale, t_range=list(self.args.sds.t_range), ) # CLIP data augmentation raster_sketch_aug, inputs_aug = self.clip_pair_augment( raster_sketch, inputs, im_res=224, augments=self.cargs.augmentations, num_aug=self.cargs.num_aug ) # raster_sketch: (1, 3, 224, 224), [0, 1] # inputs: (1, 3, 224, 224), [0, 1] # raster_sketch_aug: (5, 3, 224, 224), [2+, -1.7] # inputs_aug: (5, 3, 224, 224), [2+, -1.7] # clip visual loss total_visual_loss = torch.tensor(0) l_clip_fc, l_clip_conv, clip_conv_loss_sum = torch.tensor(0), [], torch.tensor(0) if self.args.clip.vis_loss > 0: l_clip_fc, l_clip_conv = self.clip_score_fn.compute_visual_distance( raster_sketch_aug, inputs_aug, clip_norm=False ) clip_conv_loss_sum = sum(l_clip_conv) total_visual_loss = self.args.clip.vis_loss * (clip_conv_loss_sum + l_clip_fc) # perceptual loss l_percep = torch.tensor(0.) if perceptual_loss_fn is not None: l_perceptual = perceptual_loss_fn(raster_sketch, inputs).mean() l_percep = l_perceptual * self.args.perceptual.coeff # text-visual loss l_tvd = torch.tensor(0.) if self.cargs.text_visual_coeff > 0: l_tvd = self.clip_score_fn.compute_text_visual_distance( raster_sketch_aug, target_prompt ) * self.cargs.text_visual_coeff # total loss loss = sds_loss + total_visual_loss + l_percep + l_tvd # optimization optimizer.zero_grad_() loss.backward() optimizer.step_() # if self.step % self.args.pruning_freq == 0: # renderer.path_pruning() # update lr if self.args.lr_scheduler: optimizer.update_lr(self.step, self.args.lr, self.args.decay_steps) # records pbar.set_description( f"lr: {optimizer.get_lr():.2f}, " f"l_total: {loss.item():.4f}, " f"l_clip_fc: {l_clip_fc.item():.4f}, " f"l_clip_conv({len(l_clip_conv)}): {clip_conv_loss_sum.item():.4f}, " f"l_tvd: {l_tvd.item():.4f}, " f"l_percep: {l_percep.item():.4f}, " f"sds: {grad.item():.4e}" ) # log video if self.args.make_video and (self.step % self.args.video_frame_freq == 0) \ and self.accelerator.is_main_process: log_tensor_img(raster_sketch, output_dir=self.png_logs_dir, output_prefix=f'frame{vid_idx}', dpi=100) vid_idx += 1 # log raster and svg if self.step % self.args.save_step == 0 and self.accelerator.is_main_process: # log png plt_batch(inputs, raster_sketch, self.step, target_prompt, save_path=self.png_logs_dir.as_posix(), name=f"iter{self.step}") # log svg renderer.save_svg(self.svg_logs_dir.as_posix(), f"svg_iter{self.step}") # log cross attn if self.args.log_cross_attn: controller = AttentionStore() _, _ = self.diffusion.get_cross_attention([target_prompt], controller, res=self.args.cross_attn_res, from_where=("up", "down"), save_path=self.attn_logs_dir / f"iter{self.step}.png") # logging the best raster images and SVG if self.step % self.args.eval_step == 0 and self.accelerator.is_main_process: with torch.no_grad(): # visual metric l_clip_fc, l_clip_conv = self.clip_score_fn.compute_visual_distance( raster_sketch_aug, inputs_aug, clip_norm=False ) loss_eval = sum(l_clip_conv) + l_clip_fc cur_delta = loss_eval.item() - best_visual_loss if abs(cur_delta) > min_delta and cur_delta < 0: best_visual_loss = loss_eval.item() best_iter_v = self.step plt_batch(inputs, raster_sketch, best_iter_v, target_prompt, save_path=self.results_path.as_posix(), name="visual_best") renderer.save_svg(self.results_path.as_posix(), "visual_best") # semantic metric loss_eval = self.clip_score_fn.compute_text_visual_distance( raster_sketch_aug, target_prompt ) cur_delta = loss_eval.item() - best_semantic_loss if abs(cur_delta) > min_delta and cur_delta < 0: best_semantic_loss = loss_eval.item() best_iter_s = self.step plt_batch(inputs, raster_sketch, best_iter_s, target_prompt, save_path=self.results_path.as_posix(), name="semantic_best") renderer.save_svg(self.results_path.as_posix(), "semantic_best") # log attention if self.step == 0 and self.args.attention_init and self.accelerator.is_main_process: plt_attn(renderer.get_attn(), renderer.get_thresh(), inputs, renderer.get_inds(), (self.results_path / "attention_map.jpg").as_posix()) self.step += 1 pbar.update(1) # saving final svg renderer.save_svg(self.svg_logs_dir.as_posix(), "final_svg_tmp") # stroke pruning if self.args.opacity_delta != 0: remove_low_opacity_paths(self.svg_logs_dir / "final_svg_tmp.svg", self.results_path / "final_svg.svg", self.args.opacity_delta) # save raster img final_raster_sketch = renderer.get_image().to(self.device) save_tensor_img(final_raster_sketch, save_path=self.results_path, name='final_render') # convert the intermediate renderings to a video if self.args.make_video: from subprocess import call call([ "ffmpeg", "-framerate", 24, "-i", (self.png_logs_dir / "frame%d.png").as_posix(), "-vb", "20M", (self.results_path / "out.mp4").as_posix() ]) # self.close(msg="painterly rendering complete.") def get_target(self, target_file, image_size, output_dir, u2net_path, mask_object, fix_scale, device): if not is_image_file(target_file): raise TypeError(f"{target_file} is not image file.") target = Image.open(target_file) if target.mode == "RGBA": # Create a white rgba background new_image = Image.new("RGBA", target.size, "WHITE") # Paste the image on the background. new_image.paste(target, (0, 0), target) target = new_image target = target.convert("RGB") # U2Net mask mask = target if mask_object: if pathlib.Path(u2net_path).exists(): masked_im, mask = get_mask_u2net(target, output_dir, u2net_path, device) target = masked_im else: self.print(f"'{u2net_path}' is not exist, disable mask target") if fix_scale: target = fix_image_scale(target) # define image transforms transforms_ = [] if target.size[0] != target.size[1]: transforms_.append(transforms.Resize((image_size, image_size))) else: transforms_.append(transforms.Resize(image_size)) transforms_.append(transforms.CenterCrop(image_size)) transforms_.append(transforms.ToTensor()) # preprocess data_transforms = transforms.Compose(transforms_) target_ = data_transforms(target).unsqueeze(0).to(self.device) return target_, mask