src/demo/model.py

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