main.py

import os
import spaces
import torch
import numpy as np
import argparse
from peft import LoraConfig
from pipeline_dedit_sdxl import DEditSDXLPipeline
from pipeline_dedit_sd import DEditSDPipeline
from utils import load_image, load_mask, load_mask_edit
from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys

@spaces.GPU(duration=45)
    
def run_main(
    name="example_tmp",
    name_2=None,
    mask_np_list=None, 
    mask_label_list=None,
    image_gt=None,
    dpm="sd",
    resolution=512,
    seed=42,
    embedding_learning_rate=1e-4,
    max_emb_train_steps=200,
    diffusion_model_learning_rate=5e-5,
    max_diffusion_train_steps=200,
    train_batch_size=1,
    gradient_accumulation_steps=1,
    num_tokens=1,

    load_trained=False ,
    num_sampling_steps=50,
    guidance_scale= 3 ,
    strength=0.8,

    train_full_lora=False ,
    lora_rank=4,
    lora_alpha=4,

    prompt_auxin_list = None,
    prompt_auxin_idx_list= None,

    load_edited_mask=False,
    load_edited_processed_mask=False,
    edge_thickness=20,
    num_imgs= 1 ,
    active_mask_list = None,
    tgt_index=None,

    recon=False ,
    recon_an_item=False,
    recon_prompt=None,

    text=False,
    tgt_prompt=None,

    image=False ,
    src_index=None,
    tgt_name=None,

    move_resize=False ,
    tgt_indices_list=None,
    delta_x_list=None,
    delta_y_list=None,
    priority_list=None,
    force_mask_remain=None,
    resize_list=None,

    remove=False,
    load_edited_removemask=False
):

    torch.cuda.manual_seed_all(seed)
    torch.manual_seed(seed)  
    base_input_folder = "."
    base_output_folder  = "."

    input_folder = os.path.join(base_input_folder, name)
    mask_list = []
    for mask_np in mask_np_list:
        mask = torch.from_numpy(mask_np.astype(np.uint8))
        mask_list.append(mask)
    
    #mask_list, mask_label_list = load_mask(input_folder)
    assert mask_list[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution)
    #try:
    #    image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(resolution) ), size = resolution)
    #except:
    #    image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(resolution) ), size = resolution)

    if image:
        input_folder_2 = os.path.join(base_input_folder, name_2)
        mask_list_2, mask_label_list_2 = load_mask(input_folder_2)
        assert mask_list_2[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution)
        try:
            image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.png".format(resolution) ), size = resolution)
        except:
            image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.jpg".format(resolution) ), size = resolution)
        output_dir = os.path.join(base_output_folder, name + "_" + name_2)
        os.makedirs(output_dir, exist_ok = True)
    else:
        output_dir = os.path.join(base_output_folder, name)
        os.makedirs(output_dir, exist_ok = True)

    if dpm == "sd":
        if image:
            pipe = DEditSDPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = resolution, num_tokens = num_tokens)
        else:
            pipe = DEditSDPipeline(mask_list, mask_label_list, resolution = resolution, num_tokens = num_tokens)
            
    elif dpm == "sdxl":
        if image:
            pipe = DEditSDXLPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = resolution, num_tokens = num_tokens)
        else:
            pipe = DEditSDXLPipeline(mask_list, mask_label_list, resolution = resolution, num_tokens = num_tokens)

    else:
        raise NotImplementedError

    set_string_list = pipe.set_string_list
    if prompt_auxin_list is not None:
        for auxin_idx, auxin_prompt in zip(prompt_auxin_idx_list, prompt_auxin_list):
            set_string_list[auxin_idx] = auxin_prompt.replace("*", set_string_list[auxin_idx] )
    print(set_string_list)

    if image: 
        set_string_list_2 = pipe.set_string_list_2
        print(set_string_list_2)

    if load_trained:
        unet_save_path = os.path.join(output_dir, "unet.pt")
        unet_state_dict = torch.load(unet_save_path)
        text_encoder1_save_path = os.path.join(output_dir, "text_encoder1.pt")
        text_encoder1_state_dict = torch.load(text_encoder1_save_path)
        if dpm == "sdxl":
            text_encoder2_save_path = os.path.join(output_dir, "text_encoder2.pt")
            text_encoder2_state_dict = torch.load(text_encoder2_save_path)

        if 'lora' in ''.join(unet_state_dict.keys()):
            unet_lora_config = LoraConfig(
                    r=lora_rank,
                    lora_alpha=lora_alpha,
                    init_lora_weights="gaussian",
                    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
                )
            pipe.unet.add_adapter(unet_lora_config) 
        
        pipe.unet.load_state_dict(unet_state_dict)
        pipe.text_encoder.load_state_dict(text_encoder1_state_dict)
        if dpm == "sdxl":
            pipe.text_encoder_2.load_state_dict(text_encoder2_state_dict)
    else:
        if image:
            pipe.mask_list = [m.cuda() for m in pipe.mask_list]
            pipe.mask_list_2 = [m.cuda() for m in pipe.mask_list_2] 
            pipe.train_emb_2imgs(
                image_gt,
                image_gt_2, 
                set_string_list,
                set_string_list_2,
                gradient_accumulation_steps = gradient_accumulation_steps,
                embedding_learning_rate = embedding_learning_rate,
                max_emb_train_steps = max_emb_train_steps,
                train_batch_size = train_batch_size,
            )
            
            pipe.train_model_2imgs(
                image_gt,
                image_gt_2, 
                set_string_list,
                set_string_list_2,
                gradient_accumulation_steps = gradient_accumulation_steps,
                max_diffusion_train_steps = max_diffusion_train_steps,
                diffusion_model_learning_rate = diffusion_model_learning_rate ,
                train_batch_size =train_batch_size,
                train_full_lora = train_full_lora,
                lora_rank = lora_rank, 
                lora_alpha = lora_alpha
            )
            
        else:
            pipe.mask_list = [m.cuda() for m in pipe.mask_list] 
            pipe.train_emb(
                image_gt,
                set_string_list,
                gradient_accumulation_steps = gradient_accumulation_steps,
                embedding_learning_rate = embedding_learning_rate,
                max_emb_train_steps = max_emb_train_steps,
                train_batch_size = train_batch_size,
            )

            pipe.train_model(
                image_gt,
                set_string_list,
                gradient_accumulation_steps = gradient_accumulation_steps,
                max_diffusion_train_steps = max_diffusion_train_steps,
                diffusion_model_learning_rate = diffusion_model_learning_rate ,
                train_batch_size = train_batch_size,
                train_full_lora = train_full_lora,
                lora_rank = lora_rank, 
                lora_alpha = lora_alpha
            )

        
        unet_save_path = os.path.join(output_dir, "unet.pt")
        torch.save(pipe.unet.state_dict(),unet_save_path )
        text_encoder1_save_path = os.path.join(output_dir, "text_encoder1.pt")
        torch.save(pipe.text_encoder.state_dict(), text_encoder1_save_path)
        if dpm == "sdxl":
            text_encoder2_save_path = os.path.join(output_dir, "text_encoder2.pt")
            torch.save(pipe.text_encoder_2.state_dict(), text_encoder2_save_path )
        

    if recon:
        output_dir = os.path.join(output_dir, "recon")
        os.makedirs(output_dir, exist_ok = True)
        if recon_an_item:
            mask_list = [torch.from_numpy(np.ones_like(mask_list[0].numpy()))]
            tgt_string = set_string_list[tgt_index]
            tgt_string = recon_prompt.replace("*", tgt_string)
            set_string_list = [tgt_string]
        print(set_string_list)
        save_path = os.path.join(output_dir, "out_recon.png")
        x_np = pipe.inference_with_mask(
            save_path,
            guidance_scale = guidance_scale,
            num_sampling_steps = num_sampling_steps,
            seed = seed,
            num_imgs = num_imgs,
            set_string_list = set_string_list,
            mask_list = mask_list
        )
    
    if text:
        print("*** Text-guided editing ")
        output_dir = os.path.join(output_dir, "text")
        os.makedirs(output_dir, exist_ok = True)
        save_path = os.path.join(output_dir, "out_text.png")
        set_string_list[tgt_index] = tgt_prompt
        mask_active = torch.zeros_like(mask_list[0])
        mask_active = mask_union_torch(mask_active, mask_list[tgt_index])
        
        if active_mask_list is not None:
            for midx in active_mask_list:
                mask_active = mask_union_torch(mask_active, mask_list[midx])

        if load_edited_mask:
            mask_list_edited, mask_label_list_edited = load_mask_edit(input_folder)
            mask_diff = get_mask_difference_torch(mask_list_edited,  mask_list)
            mask_active = mask_union_torch(mask_active, mask_diff)
            mask_list = mask_list_edited
            save_path = os.path.join(output_dir, "out_textEdited.png")
        
        mask_hard = mask_substract_torch(torch.ones_like(mask_list[0]), mask_active)
        mask_soft = create_outer_edge_mask_torch(mask_active, edge_thickness = edge_thickness)
        mask_hard = mask_substract_torch(mask_hard, mask_soft)

        pipe.inference_with_mask(
            save_path,
            orig_image = image_gt,
            set_string_list = set_string_list,
            guidance_scale = guidance_scale,
            strength = strength,
            num_imgs = num_imgs,
            mask_hard= mask_hard,
            mask_soft = mask_soft,
            mask_list = mask_list,
            seed = seed,
            num_sampling_steps = num_sampling_steps
        )

    if remove:
        output_dir = os.path.join(output_dir, "remove")
        save_path = os.path.join(output_dir, "out_remove.png")
        os.makedirs(output_dir, exist_ok = True)
        mask_active = torch.zeros_like(mask_list[0])
        
        if load_edited_mask:
            mask_list_edited, _ = load_mask_edit(input_folder)
            mask_diff = get_mask_difference_torch(mask_list_edited,  mask_list)
            mask_active = mask_union_torch(mask_active, mask_diff)
            mask_list = mask_list_edited
            
        if load_edited_processed_mask:
            # manually edit or draw masks after removing one index, then load
            mask_list_processed, _ = load_mask_edit(output_dir)
            mask_remain = get_mask_difference_torch(mask_list_processed, mask_list)
        else:
            # generate masks after removing one index, using nearest neighbor algorithm
            mask_list_processed, mask_remain = process_mask_remove_torch(mask_list, tgt_index)
            save_mask_list_to_npys(output_dir, mask_list_processed, mask_label_list, name = "mask")
            visualize_mask_list(mask_list_processed, os.path.join(output_dir, "seg_removed.png"))
        check_cover_all_torch(*mask_list_processed)
        mask_active = mask_union_torch(mask_active, mask_remain)
        
        if active_mask_list is not None:
            for midx in active_mask_list:
                mask_active = mask_union_torch(mask_active, mask_list[midx])

        mask_hard = 1 - mask_active
        mask_soft = create_outer_edge_mask_torch(mask_remain, edge_thickness = edge_thickness)
        mask_hard = mask_substract_torch(mask_hard, mask_soft)    

        pipe.inference_with_mask(
            save_path, 
            orig_image = image_gt,
            guidance_scale = guidance_scale,
            strength = strength,
            num_imgs = num_imgs,
            mask_hard= mask_hard,
            mask_soft = mask_soft,
            mask_list = mask_list_processed, 
            seed = seed,
            num_sampling_steps = num_sampling_steps
        )

    if image:
        output_dir = os.path.join(output_dir, "image")
        save_path = os.path.join(output_dir, "out_image.png")
        os.makedirs(output_dir, exist_ok = True)
        mask_active = torch.zeros_like(mask_list[0])
        
        if None not in (tgt_name, src_index, tgt_index):
            if tgt_name == name:
                set_string_list_tgt = set_string_list
                set_string_list_src = set_string_list_2
                image_tgt = image_gt
                if load_edited_mask:
                    mask_list_edited, _ = load_mask_edit(input_folder)
                    mask_diff = get_mask_difference_torch(mask_list_edited,  mask_list)
                    mask_active = mask_union_torch(mask_active, mask_diff)
                    mask_list = mask_list_edited
                    save_path = os.path.join(output_dir, "out_imageEdited.png")
                mask_list_tgt = mask_list
                
            elif tgt_name == name_2:
                set_string_list_tgt = set_string_list_2
                set_string_list_src = set_string_list
                image_tgt = image_gt_2
                if load_edited_mask:
                    mask_list_2_edited, _ = load_mask_edit(input_folder_2)
                    mask_diff = get_mask_difference_torch(mask_list_2_edited,  mask_list_2)
                    mask_active = mask_union_torch(mask_active, mask_diff)
                    mask_list_2 = mask_list_2_edited
                    save_path = os.path.join(output_dir, "out_imageEdited.png")
                mask_list_tgt = mask_list_2
            else:
                exit("tgt_name should be either name or name_2")
                
            set_string_list_tgt[tgt_index] = set_string_list_src[src_index]
            
            mask_active = mask_list_tgt[tgt_index]
            mask_frozen = (1-mask_active.float()).to(mask_active.device)
            mask_soft = create_outer_edge_mask_torch(mask_active.cpu(), edge_thickness = edge_thickness)
            mask_hard = mask_substract_torch(mask_frozen.cpu(), mask_soft.cpu())
            
            mask_list_tgt = [m.cuda() for m in mask_list_tgt]

            pipe.inference_with_mask(
                save_path,
                set_string_list = set_string_list_tgt,
                mask_list = mask_list_tgt, 
                guidance_scale = guidance_scale,
                num_sampling_steps = num_sampling_steps,
                mask_hard = mask_hard.cuda(),
                mask_soft = mask_soft.cuda(), 
                num_imgs = num_imgs,
                orig_image = image_tgt,
                strength = strength,
            )

    if move_resize:
        output_dir = os.path.join(output_dir, "move_resize")
        os.makedirs(output_dir, exist_ok = True)
        save_path = os.path.join(output_dir, "out_moveresize.png")
        mask_active = torch.zeros_like(mask_list[0])
        
        if load_edited_mask:
            mask_list_edited, _ = load_mask_edit(input_folder)
            mask_diff = get_mask_difference_torch(mask_list_edited,  mask_list)
            mask_active = mask_union_torch(mask_active, mask_diff)
            mask_list = mask_list_edited
            # save_path = os.path.join(output_dir, "out_moveresizeEdited.png")
            
        if load_edited_processed_mask:
            mask_list_processed, _ = load_mask_edit(output_dir)
            mask_remain = get_mask_difference_torch(mask_list_processed, mask_list)
        else:
            mask_list_processed, mask_remain = process_mask_move_torch(
                mask_list,
                tgt_indices_list, 
                delta_x_list,
                delta_y_list, priority_list, 
                force_mask_remain = force_mask_remain,
                resize_list = resize_list
            )
            save_mask_list_to_npys(output_dir, mask_list_processed, mask_label_list, name = "mask")
            visualize_mask_list(mask_list_processed, os.path.join(output_dir, "seg_move_resize.png"))
        active_idxs = tgt_indices_list
        
        mask_active = mask_union_torch(mask_active, *[m for midx, m in enumerate(mask_list_processed) if midx in active_idxs])
        mask_active = mask_union_torch(mask_remain, mask_active)
        if active_mask_list is not None:
            for midx in active_mask_list:
                mask_active = mask_union_torch(mask_active, mask_list_processed[midx])

        mask_frozen =(1 - mask_active.float())
        mask_soft = create_outer_edge_mask_torch(mask_active, edge_thickness = edge_thickness)
        mask_hard = mask_substract_torch(mask_frozen, mask_soft)

        check_mask_overlap_torch(mask_hard, mask_soft)

        pipe.inference_with_mask(
            save_path,
            strength = strength,
            orig_image = image_gt, 
            guidance_scale = guidance_scale,
            num_sampling_steps =  num_sampling_steps,
            num_imgs = num_imgs,
            mask_hard= mask_hard,
            mask_soft = mask_soft,
            mask_list = mask_list_processed,
            seed = seed
        )