import numpy as np import torch from diffusers import DDIMScheduler import cv2 from utils.sdxl import sdxl from utils.inversion import Inversion import math import torch.nn.functional as F import utils.utils as utils import os import matplotlib.pyplot as plt from PIL import Image, ImageDraw, ImageFont import spaces MAX_NUM_WORDS = 77 class LayerFusion: def get_mask(self, maps, alpha, use_pool,x_t): k = 1 maps = (maps * alpha).sum(-1).mean(1) if use_pool: maps = F.max_pool2d(maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k)) mask = F.interpolate(maps, size=(x_t.shape[2:])) #[2, 1, 128, 128] mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0] mask=(mask - mask.min ()) / (mask.max () - mask.min ()) mask = mask.gt(self.mask_threshold) self.mask=mask mask = mask[:1] + mask return mask def get_one_mask(self, maps, use_pool, x_t, idx_lst, i=None, sav_img=False): k=1 if sav_img is False: mask_tot = 0 for obj in idx_lst: mask = maps[0, :, :, :, obj].mean(0).reshape(1, 1, 32, 32) if use_pool: mask = F.max_pool2d(mask, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k)) mask = F.interpolate(mask, size=(x_t.shape[2:])) mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0] mask=(mask - mask.min ()) / (mask.max () - mask.min ()) mask = mask.gt(self.mask_threshold[int(self.counter/10)]) mask_tot |= mask mask = mask_tot return mask else: for obj in idx_lst: mask = maps[0, :, :, :, obj].mean(0).reshape(1, 1, 32, 32) if use_pool: mask = F.max_pool2d(mask, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k)) mask = F.interpolate(mask, size=(1024, 1024))#[1, 1, 1024, 1024] mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0] mask=(mask - mask.min ()) / (mask.max () - mask.min ()) mask = mask.gt(0.6) mask = np.array(mask[0][0].clone().cpu()).astype(np.uint8)*255 cv2.imwrite(f'./img/sam_mask/{self.blend_list[i][0]}_{self.counter}.jpg', mask) return mask def mv_op(self, mp, op, scale=0.2, ones=False, flip=None): _, b, H, W = mp.shape if ones == False: new_mp = torch.zeros_like(mp) else: new_mp = torch.ones_like(mp) K = int(scale*W) if op == 'right': new_mp[:, :, :, K:] = mp[:, :, :, 0:W-K] elif op == 'left': new_mp[:, :, :, 0:W-K] = mp[:, :, :, K:] elif op == 'down': new_mp[:, :, K:, :] = mp[:, :, 0:W-K, :] elif op == 'up': new_mp[:, :, 0:W-K, :] = mp[:, :, K:, :] if flip is not None: new_mp = torch.flip(new_mp, dims=flip) return new_mp def mv_layer(self, x_t, bg_id, fg_id, op_id): bg_img = x_t[bg_id:(bg_id+1)].clone() fg_img = x_t[fg_id:(fg_id+1)].clone() fg_mask = self.fg_mask_list[fg_id-3] op_list = self.op_list[fg_id-3] for item in op_list: op, scale = item[0], item[1] if scale != 0: fg_img = self.mv_op(fg_img, op=op, scale=scale) fg_mask = self.mv_op(fg_mask, op=op, scale=scale) x_t[op_id:(op_id+1)] = bg_img*(1-fg_mask) + fg_img*fg_mask def __call__(self, x_t): self.counter += 1 # inpainting if self.blend_time[0] <= self.counter <= self.blend_time[1]: x_t[1:2] = x_t[1:2]*self.remove_mask + x_t[0:1]*(1-self.remove_mask) if self.counter == self.blend_time[1] + 1 and self.mode != "removal": b = x_t.shape[0] bg_id = 1 #bg_layer op_id = 2 #canvas for fg_id in range(3, b): #fg_layer self.mv_layer(x_t, bg_id=bg_id, fg_id=fg_id, op_id=op_id) bg_id = op_id return x_t def __init__(self, remove_mask, fg_mask_list, refine_mask=None, blend_time=[0, 40], mode="removal", op_list=None): self.counter = 0 self.mode = mode self.op_list = op_list self.blend_time = blend_time self.remove_mask = remove_mask self.refine_mask = refine_mask if self.refine_mask is not None: self.new_mask = self.remove_mask + self.refine_mask self.new_mask[self.new_mask>0] = 1 else: self.new_mask = None self.fg_mask_list = fg_mask_list class Control(): def step_callback(self, x_t): if self.layer_fusion is not None: x_t = self.layer_fusion(x_t) return x_t def __init__(self, layer_fusion): self.layer_fusion = layer_fusion def register_attention_control(model, controller, mask_time=[0, 40], refine_time=[0, 25]): def ca_forward(self, place_in_unet): to_out = self.to_out if type(to_out) is torch.nn.modules.container.ModuleList: to_out = self.to_out[0] else: to_out = self.to_out self.counter = 0 #time def forward(hidden_states, encoder_hidden_states=None, attention_mask=None): #self_attention x = hidden_states.clone() context = encoder_hidden_states is_cross = context is not None if is_cross is False: if controller.layer_fusion is not None and (mask_time[0] < self.counter < mask_time[1]): b, i, j = x.shape H = W = int(math.sqrt(i)) x_old = x.clone() x = x.reshape(b, H, W, j) new_mask = controller.layer_fusion.remove_mask if new_mask is not None: new_mask[new_mask>0] = 1 new_mask = F.interpolate(new_mask.to(dtype=torch.float32).clone(), size=(H, W), mode='bilinear').cuda() new_mask = (1 - new_mask).reshape(1, H, W).unsqueeze(-1) if (refine_time[0] < self.counter <= refine_time[1]) and controller.layer_fusion.refine_mask is not None: new_mask = controller.layer_fusion.new_mask new_mask = F.interpolate(new_mask.to(dtype=torch.float32).clone(), size=(H, W), mode='bilinear').cuda() new_mask = (1 - new_mask).reshape(1, H, W).unsqueeze(-1) idx = 1 #inpaiint_idx:bg x[int(b/2)+idx, :, :] = (x[int(b/2)+idx, :, :]*new_mask[0]) x = x.reshape(b, i, j) if is_cross: q = self.to_q(x) k = self.to_k(context) v = self.to_v(context) else: context = x q = self.to_q(hidden_states) k = self.to_k(x) v = self.to_v(hidden_states) q = self.head_to_batch_dim(q) k = self.head_to_batch_dim(k) v = self.head_to_batch_dim(v) if hasattr(controller, 'count_layers'): controller.count_layers(place_in_unet,is_cross) sim = torch.einsum("b i d, b j d -> b i j", q.clone(), k.clone()) * self.scale attn = sim.softmax(dim=-1) out = torch.einsum("b i j, b j d -> b i d", attn, v) out = self.batch_to_head_dim(out) global global_cnt self.counter += 1 return to_out(out) return forward def register_recr(net_, count, place_in_unet): if net_.__class__.__name__ == 'Attention': net_.forward = ca_forward(net_, place_in_unet) return count + 1 elif hasattr(net_, 'children'): for net__ in net_.children(): count = register_recr(net__, count, place_in_unet) return count cross_att_count = 0 sub_nets = model.unet.named_children() for net in sub_nets: if "down" in net[0]: cross_att_count += register_recr(net[1], 0, "down") elif "up" in net[0]: cross_att_count += register_recr(net[1], 0, "up") elif "mid" in net[0]: cross_att_count += register_recr(net[1], 0, "mid") controller.num_att_layers = cross_att_count class DesignEdit(): def __init__(self, pretrained_model_path="/home/jyr/model/stable-diffusion-xl-base-1.0"): self.model_dtype = "fp16" self.pretrained_model_path=pretrained_model_path self.num_ddim_steps = 50 self.mask_time = [0, 40] self.op_list = {} self.attend_scale = {} scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False) if self.model_dtype == "fp16": torch_dtype = torch.float16 elif self.model_dtype == "fp32": torch_dtype = torch.float32 self.pipe = sdxl.from_pretrained(self.pretrained_model_path, torch_dtype=torch_dtype, use_safetensors=True, variant=self.model_dtype,scheduler=scheduler) @spaces.GPU def init_model(self, num_ddim_steps=50): device = torch.device('cuda:0') self.pipe.to(device) inversion = Inversion(self.pipe,num_ddim_steps) return self.pipe, inversion @spaces.GPU(duration=120, enable_queue=True) def run_remove(self, original_image=None, mask_1=None, mask_2=None, mask_3=None, refine_mask=None, ori_1=None, ori_2=None, ori_3=None, prompt="", save_dir="./tmp", mode='removal',): # 01-1: self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) if original_image is None: original_image = ori_1 if ori_1 is not None else ori_2 if ori_2 is not None else ori_3 op_list = None attend_scale = 20 sample_ref_match={0 : 0, 1 : 0} ori_shape = original_image.shape # 01-2: prepare: image_gt, remove_mask, fg_mask_list, refine_mask image_gt = Image.fromarray(original_image).resize((1024, 1024)) image_gt = np.stack([np.array(image_gt)]) mask_list = [mask_1, mask_2, mask_3] remove_mask = utils.attend_mask(utils.add_masks_resized(mask_list), attend_scale=attend_scale) # numpy to tensor fg_mask_list = None refine_mask = utils.attend_mask(utils.convert_and_resize_mask(refine_mask)) if refine_mask is not None else None # 01-3: prepare: prompts, blend_time, refine_time prompts = len(sample_ref_match)*[prompt] # 2 blend_time = [0, 41] refine_time = [0, 25] # 02: invert _, x_t, x_stars, prompt_embeds, pooled_prompt_embeds = self.inversion.invert(image_gt, prompts, inv_batch_size=1) # 03: init layer_fusion and controller lb = LayerFusion(remove_mask=remove_mask, fg_mask_list=fg_mask_list, refine_mask=refine_mask, blend_time=blend_time, mode=mode, op_list=op_list) controller = Control(layer_fusion=lb) register_attention_control(model=self.ldm_model, controller=controller, mask_time=self.mask_time, refine_time=refine_time) # 04: generate images images = self.ldm_model(controller=controller, prompt=prompts, latents=x_t, x_stars=x_stars, negative_prompt_embeds=prompt_embeds, negative_pooled_prompt_embeds=pooled_prompt_embeds, sample_ref_match=sample_ref_match) folder = None utils.view_images(images, folder=folder) return [cv2.resize(images[1], (ori_shape[1], ori_shape[0]))] @spaces.GPU(duration=120, enable_queue=True) def run_zooming(self, original_image, width_scale=1, height_scale=1, prompt="", save_dir="./tmp", mode='removal'): self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) # 01-1: op_list = {0: ['zooming', [height_scale, width_scale]]} ori_shape = original_image.shape attend_scale = 30 sample_ref_match = {0 : 0, 1 : 0} # 01-2: prepare: image_gt, remove_mask, fg_mask_list, refine_mask img_new, mask = utils.zooming(original_image, [height_scale, width_scale]) img_new_copy = img_new.copy() mask_copy = mask.copy() image_gt = Image.fromarray(img_new).resize((1024, 1024)) image_gt = np.stack([np.array(image_gt)]) remove_mask = utils.attend_mask(utils.convert_and_resize_mask(mask), attend_scale=attend_scale) # numpy to tensor fg_mask_list = None refine_mask = None # 01-3: prepare: prompts, blend_time, refine_time prompts = len(sample_ref_match)*[prompt] # 2 blend_time = [0, 41] refine_time = [0, 25] # 02: invert _, x_t, x_stars, prompt_embeds, pooled_prompt_embeds = self.inversion.invert(image_gt, prompts, inv_batch_size=1) # 03: init layer_fusion and controller lb = LayerFusion(remove_mask=remove_mask, fg_mask_list=fg_mask_list, blend_time=blend_time, mode=mode, op_list=op_list) controller = Control(layer_fusion=lb) register_attention_control(model=self.ldm_model, controller=controller, mask_time=self.mask_time, refine_time=refine_time) # 04: generate images images = self.ldm_model(controller=controller, prompt=prompts, latents=x_t, x_stars=x_stars, negative_prompt_embeds=prompt_embeds, negative_pooled_prompt_embeds=pooled_prompt_embeds, sample_ref_match=sample_ref_match) folder = None utils.view_images(images, folder=folder) resized_img = cv2.resize(images[1], (ori_shape[1], ori_shape[0])) return [resized_img], [img_new_copy], [mask_copy] @spaces.GPU(duration=120, enable_queue=True) def run_panning(self, original_image, w_direction, w_scale, h_direction, h_scale, prompt="", save_dir="./tmp", mode='removal'): # 01-1: prepare: op_list, attend_scale, sample_ref_match self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) ori_shape = original_image.shape attend_scale = 30 sample_ref_match = {0 : 0, 1 : 0} # 01-2: prepare: image_gt, remove_mask, fg_mask_list, refine_mask op_list = [[w_direction, w_scale], [h_direction, h_scale]] img_new, mask = utils.panning(original_image, op_list=op_list) img_new_copy = img_new.copy() mask_copy = mask.copy() image_gt = Image.fromarray(img_new).resize((1024, 1024)) image_gt = np.stack([np.array(image_gt)]) remove_mask = utils.attend_mask(utils.convert_and_resize_mask(mask), attend_scale=attend_scale) # numpy to tensor fg_mask_list = None refine_mask = None # 01-3: prepare: prompts, blend_time, refine_time prompts = len(sample_ref_match)*[prompt] # 2 blend_time = [0, 41] refine_time = [0, 25] # 02: invert _, x_t, x_stars, prompt_embeds, pooled_prompt_embeds = self.inversion.invert(image_gt, prompts, inv_batch_size=1) # 03: init layer_fusion and controller lb = LayerFusion(remove_mask=remove_mask, fg_mask_list=fg_mask_list, blend_time=blend_time, mode=mode, op_list=op_list) controller = Control(layer_fusion=lb) register_attention_control(model=self.ldm_model, controller=controller, mask_time=self.mask_time, refine_time=refine_time) # 04: generate images images = self.ldm_model(controller=controller, prompt=prompts, latents=x_t, x_stars=x_stars, negative_prompt_embeds=prompt_embeds, negative_pooled_prompt_embeds=pooled_prompt_embeds, sample_ref_match=sample_ref_match) folder = None utils.view_images(images, folder=folder) resized_img = cv2.resize(images[1], (ori_shape[1], ori_shape[0])) return [resized_img], [img_new_copy], [mask_copy] # layer-wise multi-object editing def process_layer_states(self, layer_states): self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) image_paths = [] mask_paths = [] op_list = [] for state in layer_states: img, mask, dx, dy, resize, w_flip, h_flip = state if img is not None: img = cv2.resize(img, (1024, 1024)) mask = utils.convert_and_resize_mask(mask) dx_command = ['right', dx] if dx > 0 else ['left', -dx] dy_command = ['up', dy] if dy > 0 else ['down', -dy] flip_code = None if w_flip == "left/right" and h_flip == "down/up": flip_code = -1 elif w_flip == "left/right": flip_code = 1 # 或者其他默认值，根据您的需要设置 elif h_flip == "down/up": flip_code = 0 op_list.append([dx_command, dy_command]) img, mask, _ = utils.resize_image_with_mask(img, mask, resize) img, mask, _ = utils.flip_image_with_mask(img, mask, flip_code=flip_code) image_paths.append(img) mask_paths.append(utils.attend_mask(mask)) sample_ref_match = {0: 0, 1: 0, 2: 0, 3: 1, 4: 2, 5: 3} required_length = len(image_paths) + 3 truncated_sample_ref_match = {k: sample_ref_match[k] for k in sorted(sample_ref_match.keys())[:required_length]} return image_paths, mask_paths, op_list, truncated_sample_ref_match @spaces.GPU(duration=200) def run_layer(self, bg_img, l1_img, l1_dx, l1_dy, l1_resize, l1_w_flip, l1_h_flip, l2_img, l2_dx, l2_dy, l2_resize, l2_w_flip, l2_h_flip, l3_img, l3_dx, l3_dy, l3_resize, l3_w_flip, l3_h_flip, bg_mask, l1_mask, l2_mask, l3_mask, bg_ori=None, l1_ori=None, l2_ori=None, l3_ori=None, prompt="", save_dir="./tmp", mode='layerwise'): self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) # 00： prepare: layer-wise states bg_img = bg_ori if bg_ori is not None else bg_img l1_img = l1_ori if l1_ori is not None else l1_img l2_img = l2_ori if l2_ori is not None else l2_img l3_img = l3_ori if l3_ori is not None else l3_img for mask in [bg_mask, l1_mask, l2_mask, l3_mask]: if mask is None: mask = np.zeros((1024, 1024), dtype=np.uint8) else: mask = utils.convert_and_resize_mask(mask) l1_state = [l1_img, l1_mask, l1_dx, l1_dy, l1_resize, l1_w_flip, l1_h_flip] l2_state = [l2_img, l2_mask, l2_dx, l2_dy, l2_resize, l2_w_flip, l2_h_flip] l3_state = [l3_img, l3_mask, l3_dx, l3_dy, l3_resize, l3_w_flip, l3_h_flip] ori_shape = bg_img.shape image_paths, fg_mask_list, op_list, sample_ref_match = self.process_layer_states([l1_state, l2_state, l3_state]) if image_paths == []: mode = "removal" # 01-1: prepare: image_gt, remove_mask, fg_mask_list, refine_mask attend_scale = 20 image_gt = [bg_img] + image_paths image_gt = [Image.fromarray(img).resize((1024, 1024)) for img in image_gt] image_gt = np.stack(image_gt) remove_mask = utils.attend_mask(bg_mask, attend_scale=attend_scale) refine_mask = None # 01-2: prepare: promptrun_masks, blend_time, refine_time prompts = len(sample_ref_match)*[prompt] # 2 blend_time = [0, 41] refine_time = [0, 25] attend_scale = [] # 02: invert _, x_t, x_stars, prompt_embeds, pooled_prompt_embeds = self.inversion.invert(image_gt, prompts, inv_batch_size=len(image_gt)) # 03: init layer_fusion and controller lb = LayerFusion(remove_mask=remove_mask, fg_mask_list=fg_mask_list, blend_time=blend_time, refine_mask=refine_mask, mode=mode, op_list=op_list) controller = Control(layer_fusion=lb) register_attention_control(model=self.ldm_model, controller=controller, mask_time=self.mask_time, refine_time=refine_time) # 04: generate images images = self.ldm_model(controller=controller, prompt=prompts, latents=x_t, x_stars=x_stars, negative_prompt_embeds=prompt_embeds, negative_pooled_prompt_embeds=pooled_prompt_embeds, sample_ref_match=sample_ref_match) folder = None utils.view_images(images, folder=folder) if mode == 'removal': resized_img = cv2.resize(images[1], (ori_shape[1], ori_shape[0])) else: resized_img = cv2.resize(images[2], (ori_shape[1], ori_shape[0])) return [resized_img] @spaces.GPU(duration=120, enable_queue=True) def run_moving(self, bg_img, bg_ori, bg_mask, l1_dx, l1_dy, l1_resize, l1_w_flip=None, l1_h_flip=None, selected_points=None, prompt="", save_dir="./tmp", mode='layerwise'): self.ldm_model, self.inversion= self.init_model(num_ddim_steps=self.num_ddim_steps) # 00： prepare: layer-wise states bg_img = bg_ori if bg_ori is not None else bg_img l1_img = bg_img if bg_mask is None: bg_mask = np.zeros((1024, 1024), dtype=np.uint8) else: bg_mask = utils.convert_and_resize_mask(bg_mask) l1_mask = bg_mask l1_state = [l1_img, l1_mask, l1_dx, l1_dy, l1_resize, l1_w_flip, l1_h_flip] ori_shape = bg_img.shape image_paths, fg_mask_list, op_list, sample_ref_match = self.process_layer_states([l1_state]) # 01-1: prepare: image_gt, remove_mask, fg_mask_list, refine_mask attend_scale = 20 image_gt = [bg_img] + image_paths image_gt = [Image.fromarray(img).resize((1024, 1024)) for img in image_gt] image_gt = np.stack(image_gt) remove_mask = utils.attend_mask(bg_mask, attend_scale=attend_scale) refine_mask = None # 01-2: prepare: promptrun_masks, blend_time, refine_time prompts = len(sample_ref_match)*[prompt] # 2 blend_time = [0, 41] refine_time = [0, 25] attend_scale = [] # 02: invert _, x_t, x_stars, prompt_embeds, pooled_prompt_embeds = self.inversion.invert(image_gt, prompts, inv_batch_size=len(image_gt)) # 03: init layer_fusion and controller lb = LayerFusion(remove_mask=remove_mask, fg_mask_list=fg_mask_list, blend_time=blend_time, refine_mask=refine_mask, mode=mode, op_list=op_list) controller = Control(layer_fusion=lb) register_attention_control(model=self.ldm_model, controller=controller, mask_time=self.mask_time, refine_time=refine_time) # 04: generate images images = self.ldm_model(controller=controller, prompt=prompts, latents=x_t, x_stars=x_stars, negative_prompt_embeds=prompt_embeds, negative_pooled_prompt_embeds=pooled_prompt_embeds, sample_ref_match=sample_ref_match) folder = None utils.view_images(images, folder=folder) resized_img = cv2.resize(images[2], (ori_shape[1], ori_shape[0])) return [resized_img] # turn mask to 1024x1024 unit-8 def run_mask(self, mask_1, mask_2, mask_3, mask_4): mask_list = [mask_1, mask_2, mask_3, mask_4] final_mask = utils.add_masks_resized(mask_list) return final_mask