update

.gitignore

@@ -5,6 +5,7 @@ example1_example2_1024/
 example1/
 old/
 example_tmp/
+z_*
 
 out_active.png
 out_mask.png

app copy 2.py

@@ -1,385 +0,0 @@
-
-import os
-import copy
-from PIL import Image
-import matplotlib 
-import numpy as np
-import gradio as gr
-from utils import load_mask, load_mask_edit
-from utils_mask import process_mask_to_follow_priority, mask_union, visualize_mask_list_clean
-from pathlib import Path
-import subprocess
-from PIL import Image
-from functools import partial
-from main import run_main
-LENGTH=512 #length of the square area displaying/editing images
-TRANSPARENCY = 150 # transparency of the mask in display
-
-def add_mask(mask_np_list_updated, mask_label_list):
-    mask_new = np.zeros_like(mask_np_list_updated[0])
-    mask_np_list_updated.append(mask_new)
-    mask_label_list.append("new")
-    return mask_np_list_updated, mask_label_list
-
-def create_segmentation(mask_np_list):
-    viridis = matplotlib.pyplot.get_cmap(name = 'viridis', lut = len(mask_np_list))
-    segmentation = 0
-    for i, m  in enumerate(mask_np_list):
-        color = matplotlib.colors.to_rgb(viridis(i))
-        color_mat = np.ones_like(m)                                                                            
-        color_mat = np.stack([color_mat*color[0], color_mat*color[1],color_mat*color[2] ], axis = 2)
-        color_mat = color_mat * m[:,:,np.newaxis]
-        segmentation += color_mat
-    segmentation = Image.fromarray(np.uint8(segmentation*255))
-    return segmentation
-
-def load_mask_ui(input_folder="example_tmp",load_edit = False):
-    if not load_edit:
-        mask_list, mask_label_list = load_mask(input_folder)
-    else:
-        mask_list, mask_label_list = load_mask_edit(input_folder) 
-        
-    mask_np_list = []
-    for  m  in mask_list:
-        mask_np_list. append( m.cpu().numpy())
-
-    return mask_np_list, mask_label_list
-
-def load_image_ui(load_edit, input_folder="example_tmp"):
-    try:
-        for img_path in Path(input_folder).iterdir():
-            if img_path.name in ["img_512.png"]:
-                image = Image.open(img_path)
-        mask_np_list, mask_label_list = load_mask_ui(input_folder, load_edit = load_edit)
-        image = image.convert('RGB')
-        segmentation = create_segmentation(mask_np_list)
-        print("!!", len(mask_np_list))  
-        return image, segmentation, mask_np_list, mask_label_list, image
-    except:
-        print("Image folder invalid: The folder should contain image.png")
-        return None, None, None, None, None
-
-def run_edit_text(
-        num_tokens,
-        num_sampling_steps,
-        strength,
-        edge_thickness,
-        tgt_prompt,
-        tgt_idx,
-        guidance_scale,
-        input_folder="example_tmp"
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--text",
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--load_trained",
-                    "--num_tokens={}".format(num_tokens),
-                    "--seed={}".format(2024),
-                    "--guidance_scale={}".format(guidance_scale),
-                    "--num_sampling_step={}".format(num_sampling_steps),
-                    "--strength={}".format(strength),
-                    "--edge_thickness={}".format(edge_thickness),
-                    "--num_imgs={}".format(2),
-                    "--tgt_prompt={}".format(tgt_prompt) ,
-                    "--tgt_index={}".format(tgt_idx)            
-    ])
-    
-    return Image.open(os.path.join(input_folder, "text", "out_text_0.png"))
-
-
-def run_optimization(
-        num_tokens,
-        embedding_learning_rate, 
-        max_emb_train_steps, 
-        diffusion_model_learning_rate, 
-        max_diffusion_train_steps,
-        train_batch_size,
-        gradient_accumulation_steps,
-        input_folder = "example_tmp"
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--num_tokens={}".format(num_tokens),
-                    "--embedding_learning_rate={}".format(embedding_learning_rate),
-                    "--diffusion_model_learning_rate={}".format(diffusion_model_learning_rate),
-                    "--max_emb_train_steps={}".format(max_emb_train_steps),
-                    "--max_diffusion_train_steps={}".format(max_diffusion_train_steps),
-                    "--train_batch_size={}".format(train_batch_size),
-                    "--gradient_accumulation_steps={}".format(gradient_accumulation_steps)
-                    
-    ])
-    return 
-
-
-def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
-    backimg_solid_np =  np.array(backimg)
-    bimg = backimg.copy()
-    fimg = foreimg.copy()
-    fimg.putalpha(transparency)
-    bimg.paste(fimg, (0,0), fimg)
-
-    bimg_np = np.array(bimg)
-    mask_np = mask_np[:,:,np.newaxis]
-
-    try:
-        new_img_np = bimg_np*mask_np + (1-mask_np)* backimg_solid_np
-        return Image.fromarray(new_img_np) 
-    except:
-        import pdb; pdb.set_trace()
-
-def show_segmentation(image, segmentation, flag):
-    if flag is False:
-        flag = True
-        mask_np = np.ones([image.size[0],image.size[1]]).astype(np.uint8)
-        image_edit = transparent_paste_with_mask(image, segmentation, mask_np ,transparency = TRANSPARENCY)
-        return image_edit, flag
-    else:
-        flag = False
-        return image,flag
-
-def edit_mask_add(canvas,  image, idx, mask_np_list):
-    mask_sel = mask_np_list[idx]
-    mask_new = np.uint8(canvas["mask"][:, :, 0]/ 255.)
-    mask_np_list_updated = []
-    for midx, m  in enumerate(mask_np_list):
-        if midx == idx:
-            mask_np_list_updated.append(mask_union(mask_sel, mask_new))
-        else:
-            mask_np_list_updated.append(m)
-    
-    priority_list = [0 for _ in range(len(mask_np_list_updated))]
-    priority_list[idx] = 1
-    mask_np_list_updated = process_mask_to_follow_priority(mask_np_list_updated, priority_list)
-    mask_ones = np.ones([mask_sel.shape[0], mask_sel.shape[1]]).astype(np.uint8)
-    segmentation = create_segmentation(mask_np_list_updated)
-    image_edit = transparent_paste_with_mask(image, segmentation, mask_ones ,transparency = TRANSPARENCY)
-    return mask_np_list_updated, image_edit
-
-def slider_release(index, image,  mask_np_list_updated, mask_label_list):
-
-    if index > len(mask_np_list_updated):
-        return image, "out of range"
-    else:
-        mask_np = mask_np_list_updated[index]
-        mask_label = mask_label_list[index]
-        segmentation = create_segmentation(mask_np_list_updated)
-        new_image = transparent_paste_with_mask(image, segmentation, mask_np, transparency = TRANSPARENCY)
-    return new_image, mask_label
-
-def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder="example_tmp"):
-    try: 
-        assert np.all(sum(mask_np_list_updated)==1)
-    except:
-        print("please check mask")
-        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
-        import pdb; pdb.set_trace()
-        
-    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
-        # np.save(os.path.join(input_folder, "maskEDIT{}_{}.npy".format(midx, mask_label)),mask )
-        np.save(os.path.join(input_folder, "mask{}_{}.npy".format(midx, mask_label)),mask )
-    savepath = os.path.join(input_folder, "seg_current.png")
-    visualize_mask_list_clean(mask_np_list_updated, savepath)
-    
-def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder="example_tmp"):
-    try: 
-        assert np.all(sum(mask_np_list_updated)==1)
-    except:
-        print("please check mask")
-        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
-        import pdb; pdb.set_trace()
-    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
-        np.save(os.path.join(input_folder, "maskEdited{}_{}.npy".format(midx, mask_label)), mask)
-    savepath = os.path.join(input_folder, "seg_edited.png")
-    visualize_mask_list_clean(mask_np_list_updated, savepath)
-  
-  
-import shutil
-if os.path.isdir("./example_tmp"):
-    shutil.rmtree("./example_tmp")
-
-from segment import run_segmentation
-with gr.Blocks() as demo:
-    image = gr.State() # store mask
-    image_loaded = gr.State()
-    segmentation    = gr.State()
-
-    mask_np_list    = gr.State([])
-    mask_label_list = gr.State([])
-    mask_np_list_updated = gr.State([])
-    true    = gr.State(True)
-    false    = gr.State(False)
-
-    with gr.Row():
-        gr.Markdown("""# D-Edit""")
-
-    with gr.Tab(label="1 Edit mask"):
-        with gr.Row():
-            with gr.Column():
-                canvas = gr.Image(value = "./img.png", type="numpy",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-                
-                segment_button  = gr.Button("1.1 Run segmentation")
-                segment_button.click(run_segmentation, 
-                        [canvas] , 
-                        [] )
-
-                text_button  = gr.Button("1.2 Load original masks")
-                text_button.click(load_image_ui, 
-                        [ false] , 
-                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
-
-                load_edit_button = gr.Button("1.2 Load edited masks")    
-                load_edit_button.click(load_image_ui, 
-                        [ true] , 
-                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
-                
-                show_segment = gr.Checkbox(label = "Show Segmentation")
-                flag = gr.State(False)
-                show_segment.select(show_segmentation,
-                                    [image_loaded, segmentation, flag], 
-                                    [canvas, flag])
- 
-            # mask_np_list_updated.value = copy.deepcopy(mask_np_list.value) #!!
-            mask_np_list_updated = mask_np_list
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Draw Mask</p>""")
-                slider =  gr.Slider(0, 20, step=1,  interactive=True)
-                label = gr.Textbox()
-                slider.release(slider_release, 
-                        inputs = [slider, image_loaded,   mask_np_list_updated, mask_label_list], 
-                        outputs= [canvas, label]
-                    )
-                add_button  = gr.Button("Add")
-                add_button.click( edit_mask_add, 
-                        [canvas, image_loaded, slider, mask_np_list_updated] , 
-                        [mask_np_list_updated, canvas]
-                    )
-
-                save_button2  = gr.Button("Set and Save as edited masks")
-                save_button2.click( save_as_edit_mask, 
-                        [mask_np_list_updated,  mask_label_list] , 
-                        [] )  
-                
-                save_button  = gr.Button("Set and Save as original masks")
-                save_button.click( save_as_orig_mask, 
-                        [mask_np_list_updated,  mask_label_list] , 
-                        [] )  
-                
-                back_button  = gr.Button("Back to current seg")
-                back_button.click( load_mask_ui, 
-                                [] , 
-                                [ mask_np_list_updated,mask_label_list] )
-
-                add_mask_button = gr.Button("Add new empty mask")    
-                add_mask_button.click(add_mask, 
-                        [mask_np_list_updated, mask_label_list] , 
-                        [mask_np_list_updated, mask_label_list] )
-                
-    with gr.Tab(label="2 Optimization"):
-        with gr.Row():
-            
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
-                num_tokens = gr.Number(value="5", label="num tokens to represent each object", interactive= True)
-                embedding_learning_rate = gr.Textbox(value="0.0001", label="Embedding optimization: Learning rate", interactive= True )
-                max_emb_train_steps =  gr.Number(value="200", label="embedding optimization: Training steps", interactive= True )
-                
-                diffusion_model_learning_rate = gr.Textbox(value="0.00005", label="UNet Optimization: Learning rate", interactive= True )
-                max_diffusion_train_steps = gr.Number(value="200", label="UNet Optimization: Learning rate: Training steps", interactive= True )
-                
-                train_batch_size = gr.Number(value="5", label="Batch size", interactive= True )
-                gradient_accumulation_steps=gr.Number(value="5", label="Gradient accumulation", interactive= True )
-                
-                add_button  = gr.Button("Run optimization")
-                def run_optimization_wrapper (                        
-                        num_tokens,
-                        embedding_learning_rate , 
-                        max_emb_train_steps , 
-                        diffusion_model_learning_rate , 
-                        max_diffusion_train_steps,
-                        train_batch_size,
-                        gradient_accumulation_steps
-                ):
-                    run_optimization = partial(
-                        run_main,                 
-                        num_tokens=int(num_tokens),
-                        embedding_learning_rate = float(embedding_learning_rate), 
-                        max_emb_train_steps = int(max_emb_train_steps), 
-                        diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
-                        max_diffusion_train_steps = int(max_diffusion_train_steps),
-                        train_batch_size=int(train_batch_size),
-                        gradient_accumulation_steps=int(gradient_accumulation_steps)
-                    )
-                    run_optimization()
-                    
-                add_button.click(run_optimization_wrapper, 
-                        inputs = [
-                            num_tokens,
-                            embedding_learning_rate , 
-                            max_emb_train_steps , 
-                            diffusion_model_learning_rate , 
-                            max_diffusion_train_steps,
-                            train_batch_size,
-                            gradient_accumulation_steps
-                        ], 
-                        outputs = []
-                )
-
-
-    with gr.Tab(label="3 Editing"):
-        with gr.Tab(label="3.1 Text-based editing"):
-        
-            with gr.Row():
-                with gr.Column():
-                    canvas_text_edit = gr.Image(value = None, type = "pil", label="Editing results", show_label=True)
-                    # canvas_text_edit = gr.Gallery(label = "Edited results")
-                
-                with gr.Column():
-                    gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting (SD)</p>""")
-                    
-                    tgt_prompt =  gr.Textbox(value="White bag", label="Editing: Text prompt", interactive= True )
-                    tgt_index = gr.Number(value="0", label="Editing: Object index", interactive= True )
-                    guidance_scale = gr.Textbox(value="6", label="Editing: CFG guidance scale", interactive= True )
-                    num_sampling_steps = gr.Number(value="50", label="Editing: Sampling steps", interactive= True )
-                    edge_thickness = gr.Number(value="10", label="Editing: Edge thickness", interactive= True )
-                    strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
-                    
-                    add_button  = gr.Button("Run Editing")
-                    run_edit_text = partial(
-                        run_main,
-                        load_trained=True,
-                        text=True,
-                        num_tokens = int(num_tokens.value),
-                        guidance_scale = float(guidance_scale.value),
-                        num_sampling_steps = int(num_sampling_steps.value),
-                        strength = float(strength.value),
-                        edge_thickness = int(edge_thickness.value),
-                        num_imgs = 1,
-                        tgt_prompt = tgt_prompt.value,
-                        tgt_index = int(tgt_index.value)
-                    )
-                        
-                    add_button.click(run_edit_text, 
-                        inputs = [], 
-                        outputs = [canvas_text_edit]
-                    )
-                    
-                    def load_pil_img():
-                        from PIL import Image
-                        return Image.open("example_tmp/text/out_text_0.png")
-                    
-                    load_button  = gr.Button("Load results")
-                    load_button.click(load_pil_img, 
-                        inputs = [], 
-                        outputs = [canvas_text_edit]
-                    )
-
-
-
-
-demo.queue().launch(share=True, debug=True)

app copy.py

@@ -1,350 +0,0 @@
-
-import os
-import copy
-from PIL import Image
-import matplotlib 
-import numpy as np
-import gradio as gr
-from utils import load_mask, load_mask_edit
-from utils_mask import process_mask_to_follow_priority, mask_union, visualize_mask_list_clean
-from pathlib import Path
-import subprocess
-from PIL import Image
-
-LENGTH=512 #length of the square area displaying/editing images
-TRANSPARENCY = 150 # transparency of the mask in display
-
-def add_mask(mask_np_list_updated, mask_label_list):
-    mask_new = np.zeros_like(mask_np_list_updated[0])
-    mask_np_list_updated.append(mask_new)
-    mask_label_list.append("new")
-    return mask_np_list_updated, mask_label_list
-
-def create_segmentation(mask_np_list):
-    viridis = matplotlib.pyplot.get_cmap(name = 'viridis', lut = len(mask_np_list))
-    segmentation = 0
-    for i, m  in enumerate(mask_np_list):
-        color = matplotlib.colors.to_rgb(viridis(i))
-        color_mat = np.ones_like(m)                                                                            
-        color_mat = np.stack([color_mat*color[0], color_mat*color[1],color_mat*color[2] ], axis = 2)
-        color_mat = color_mat * m[:,:,np.newaxis]
-        segmentation += color_mat
-    segmentation = Image.fromarray(np.uint8(segmentation*255))
-    return segmentation
-
-def load_mask_ui(input_folder,load_edit = False):
-    if not load_edit:
-        mask_list, mask_label_list = load_mask(input_folder)
-    else:
-        mask_list, mask_label_list = load_mask_edit(input_folder) 
-        
-    mask_np_list = []
-    for  m  in mask_list:
-        mask_np_list. append( m.cpu().numpy())
-
-    return mask_np_list, mask_label_list
-
-def load_image_ui(input_folder, load_edit):
-    try:
-        for img_path in Path(input_folder).iterdir():
-            if img_path.name in ["img.png", "img_1024.png", "img_512.png"]:
-                image = Image.open(img_path)
-        mask_np_list, mask_label_list = load_mask_ui(input_folder, load_edit = load_edit)
-        image = image.convert('RGB')
-        segmentation = create_segmentation(mask_np_list)  
-        return image, segmentation, mask_np_list, mask_label_list, image
-    except:
-        print("Image folder invalid: The folder should contain image.png")
-        return None, None, None, None, None
-
-def run_segmentation(input_folder):
-    subprocess.run(["python", "segment.py" , "--name={}".format(input_folder)])
-    return 
-
-
-
-def run_edit_text(
-        input_folder,
-        num_tokens,
-        num_sampling_steps,
-        strength,
-        edge_thickness,
-        tgt_prompt,
-        tgt_idx,
-        guidance_scale
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--text",
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--load_trained",
-                    "--num_tokens={}".format(num_tokens),
-                    "--seed={}".format(2024),
-                    "--guidance_scale={}".format(guidance_scale),
-                    "--num_sampling_step={}".format(num_sampling_steps),
-                    "--strength={}".format(strength),
-                    "--edge_thickness={}".format(edge_thickness),
-                    "--num_imgs={}".format(2),
-                    "--tgt_prompt={}".format(tgt_prompt) ,
-                    "--tgt_index={}".format(tgt_idx)            
-    ])
-    
-    return Image.open(os.path.join(input_folder, "text", "out_text_0.png"))
-
-
-def run_optimization(
-        input_folder,
-        num_tokens,
-        embedding_learning_rate, 
-        max_emb_train_steps, 
-        diffusion_model_learning_rate, 
-        max_diffusion_train_steps,
-        train_batch_size,
-        gradient_accumulation_steps
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--num_tokens={}".format(num_tokens),
-                    "--embedding_learning_rate={}".format(embedding_learning_rate),
-                    "--diffusion_model_learning_rate={}".format(diffusion_model_learning_rate),
-                    "--max_emb_train_steps={}".format(max_emb_train_steps),
-                    "--max_diffusion_train_steps={}".format(max_diffusion_train_steps),
-                    "--train_batch_size={}".format(train_batch_size),
-                    "--gradient_accumulation_steps={}".format(gradient_accumulation_steps)
-                    
-    ])
-    return 
-
-
-def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
-    backimg_solid_np =  np.array(backimg)
-    bimg = backimg.copy()
-    fimg = foreimg.copy()
-    fimg.putalpha(transparency)
-    bimg.paste(fimg, (0,0), fimg)
-
-    bimg_np = np.array(bimg)
-    mask_np = mask_np[:,:,np.newaxis]
-    try:
-        new_img_np = bimg_np*mask_np + (1-mask_np)* backimg_solid_np
-        return Image.fromarray(new_img_np) 
-    except:
-        import pdb; pdb.set_trace()
-
-def show_segmentation(image, segmentation, flag):
-    if flag is False:
-        flag = True
-        mask_np = np.ones([image.size[0],image.size[1]]).astype(np.uint8)
-        image_edit = transparent_paste_with_mask(image, segmentation, mask_np ,transparency = TRANSPARENCY)
-        return image_edit, flag
-    else:
-        flag = False
-        return image,flag
-
-def edit_mask_add(canvas,  image, idx, mask_np_list):
-    mask_sel = mask_np_list[idx]
-    mask_new = np.uint8(canvas["mask"][:, :, 0]/ 255.)
-    mask_np_list_updated = []
-    for midx, m  in enumerate(mask_np_list):
-        if midx == idx:
-            mask_np_list_updated.append(mask_union(mask_sel, mask_new))
-        else:
-            mask_np_list_updated.append(m)
-    
-    priority_list = [0 for _ in range(len(mask_np_list_updated))]
-    priority_list[idx] = 1
-    mask_np_list_updated = process_mask_to_follow_priority(mask_np_list_updated, priority_list)
-    mask_ones = np.ones([mask_sel.shape[0], mask_sel.shape[1]]).astype(np.uint8)
-    segmentation = create_segmentation(mask_np_list_updated)
-    image_edit = transparent_paste_with_mask(image, segmentation, mask_ones ,transparency = TRANSPARENCY)
-    return mask_np_list_updated, image_edit
-
-def slider_release(index, image,  mask_np_list_updated, mask_label_list):
-    if index > len(mask_np_list_updated):
-        return image, "out of range"
-    else:
-        mask_np = mask_np_list_updated[index]
-        mask_label = mask_label_list[index]
-        segmentation = create_segmentation(mask_np_list_updated)
-        new_image = transparent_paste_with_mask(image, segmentation, mask_np, transparency = TRANSPARENCY)
-    return new_image, mask_label
-
-def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder):
-    try: 
-        assert np.all(sum(mask_np_list_updated)==1)
-    except:
-        print("please check mask")
-        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
-        import pdb; pdb.set_trace()
-        
-    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
-        # np.save(os.path.join(input_folder, "maskEDIT{}_{}.npy".format(midx, mask_label)),mask )
-        np.save(os.path.join(input_folder, "mask{}_{}.npy".format(midx, mask_label)),mask )
-    savepath = os.path.join(input_folder, "seg_current.png")
-    visualize_mask_list_clean(mask_np_list_updated, savepath)
-    
-def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder):
-    try: 
-        assert np.all(sum(mask_np_list_updated)==1)
-    except:
-        print("please check mask")
-        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
-        import pdb; pdb.set_trace()
-    for midx, (mask, mask_label) in enumerate(zip(mask_np_list_updated, mask_label_list)):
-        np.save(os.path.join(input_folder, "maskEdited{}_{}.npy".format(midx, mask_label)), mask)
-    savepath = os.path.join(input_folder, "seg_edited.png")
-    visualize_mask_list_clean(mask_np_list_updated, savepath)
-  
-with gr.Blocks() as demo:
-    image = gr.State() # store mask
-    image_loaded = gr.State()
-    segmentation    = gr.State()
-
-    mask_np_list    = gr.State([])
-    mask_label_list = gr.State([])
-    mask_np_list_updated = gr.State([])
-    true    = gr.State(True)
-    false    = gr.State(False)
-
-
-    with gr.Row():
-        gr.Markdown("""# D-Edit""")
-
-    with gr.Tab(label="1 Edit mask"):
-        with gr.Row():
-            with gr.Column():
-                canvas = gr.Image(value = None, type="numpy",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-                input_folder = gr.Textbox(value="example1", label="input folder", interactive= True, )
-                
-                segment_button  = gr.Button("1.1 Run segmentation")
-                segment_button.click(run_segmentation, 
-                        [input_folder] , 
-                        [] )
-             
-                
-                text_button  = gr.Button("1.2 Load original masks")
-                text_button.click(load_image_ui, 
-                        [input_folder, false] , 
-                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
-            
-                load_edit_button = gr.Button("1.2 Load edited masks")    
-                load_edit_button.click(load_image_ui, 
-                        [input_folder, true] , 
-                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
-                
-                show_segment = gr.Checkbox(label = "Show Segmentation")
-       
-                flag = gr.State(False)
-                show_segment.select(show_segmentation,
-                                    [image_loaded, segmentation, flag], 
-                                    [canvas, flag])
-       
-            mask_np_list_updated = copy.deepcopy(mask_np_list)
-
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Draw Mask</p>""")
-                slider =  gr.Slider(0, 20, step=1,  interactive=True)
-                label = gr.Textbox()
-                slider.release(slider_release, 
-                        inputs = [slider, image_loaded,   mask_np_list_updated, mask_label_list], 
-                        outputs= [canvas, label]
-                    )
-                add_button  = gr.Button("Add")
-                add_button.click( edit_mask_add, 
-                        [canvas, image_loaded, slider, mask_np_list_updated] , 
-                        [mask_np_list_updated, canvas]
-                    )
-
-                save_button2  = gr.Button("Set and Save as edited masks")
-                save_button2.click( save_as_edit_mask, 
-                        [mask_np_list_updated,  mask_label_list, input_folder] , 
-                        [] )  
-                
-                save_button  = gr.Button("Set and Save as original masks")
-                save_button.click( save_as_orig_mask, 
-                        [mask_np_list_updated,  mask_label_list, input_folder] , 
-                        [] )  
-                
-                back_button  = gr.Button("Back to current seg")
-                back_button.click( load_mask_ui, 
-                                [input_folder] , 
-                                [ mask_np_list_updated,mask_label_list] )
-
-                add_mask_button = gr.Button("Add new empty mask")    
-                add_mask_button.click(add_mask, 
-                        [mask_np_list_updated, mask_label_list] , 
-                        [mask_np_list_updated, mask_label_list] )
-                
-    with gr.Tab(label="2 Optimization"):
-        with gr.Row():
-            with gr.Column():
-                canvas_opt = gr.Image(value = canvas.value, type="pil",  label="Loaded Image", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-            
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
-                num_tokens = gr.Textbox(value="5", label="num tokens to represent each object", interactive= True)
-                embedding_learning_rate = gr.Textbox(value="1e-4", label="Embedding optimization: Learning rate", interactive= True )
-                max_emb_train_steps =  gr.Textbox(value="500", label="embedding optimization: Training steps", interactive= True )
-                
-                diffusion_model_learning_rate = gr.Textbox(value="5e-5", label="UNet Optimization: Learning rate", interactive= True )
-                max_diffusion_train_steps = gr.Textbox(value="500", label="UNet Optimization: Learning rate: Training steps", interactive= True )
-                
-                train_batch_size = gr.Textbox(value="5", label="Batch size", interactive= True )
-                gradient_accumulation_steps=gr.Textbox(value="5", label="Gradient accumulation", interactive= True )
-                
-                add_button  = gr.Button("Run optimization")
-                add_button.click(run_optimization, 
-                        inputs = [
-                            input_folder,
-                            num_tokens,
-                            embedding_learning_rate, 
-                            max_emb_train_steps, 
-                            diffusion_model_learning_rate, 
-                            max_diffusion_train_steps,
-                            train_batch_size,gradient_accumulation_steps
-                        ], 
-                        outputs = []
-                )
-
-
-    with gr.Tab(label="3 Editing"):
-        with gr.Tab(label="3.1 Text-based editing"):
-            canvas_text_edit =  gr.State() # store mask
-            with gr.Row():
-                with gr.Column():
-                    canvas_text_edit = gr.Image(value = None, type="pil", label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
-                    # canvas_text_edit = gr.Gallery(label = "Edited results")
-                
-                with gr.Column():
-                    gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting (SD)</p>""")
-                    
-                    tgt_prompt =  gr.Textbox(value="Dog", label="Editing: Text prompt", interactive= True )
-                    tgt_idx = gr.Textbox(value="0", label="Editing: Object index", interactive= True )
-                    guidance_scale = gr.Textbox(value="6", label="Editing: CFG guidance scale", interactive= True )
-                    num_sampling_steps = gr.Textbox(value="50", label="Editing: Sampling steps", interactive= True )
-                    edge_thickness = gr.Textbox(value="10", label="Editing: Edge thickness", interactive= True )
-                    strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
-                    
-                    add_button  = gr.Button("Run Editing")
-                    add_button.click(run_edit_text, 
-                        inputs = [
-                            input_folder,
-                            num_tokens,
-                            num_sampling_steps,
-                            strength,
-                            edge_thickness,
-                            tgt_prompt,
-                            tgt_idx,
-                            guidance_scale
-                        ], 
-                        outputs = []
-                    )
-
-
-
-demo.queue().launch(share=True, debug=True)

app.py

@@ -59,62 +59,62 @@ def load_image_ui(load_edit, input_folder="example_tmp"):
         print("Image folder invalid: The folder should contain image.png")
         return None, None, None, None, None
 
-def run_edit_text(
-        num_tokens,
-        num_sampling_steps,
-        strength,
-        edge_thickness,
-        tgt_prompt,
-        tgt_idx,
-        guidance_scale,
-        input_folder="example_tmp"
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--text",
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--load_trained",
-                    "--num_tokens={}".format(num_tokens),
-                    "--seed={}".format(2024),
-                    "--guidance_scale={}".format(guidance_scale),
-                    "--num_sampling_step={}".format(num_sampling_steps),
-                    "--strength={}".format(strength),
-                    "--edge_thickness={}".format(edge_thickness),
-                    "--num_imgs={}".format(2),
-                    "--tgt_prompt={}".format(tgt_prompt) ,
-                    "--tgt_index={}".format(tgt_idx)            
-    ])
+# def run_edit_text(
+#         num_tokens,
+#         num_sampling_steps,
+#         strength,
+#         edge_thickness,
+#         tgt_prompt,
+#         tgt_idx,
+#         guidance_scale,
+#         input_folder="example_tmp"
+#     ):
+#     subprocess.run(["python", 
+#                     "main.py" ,
+#                     "--text=True",
+#                     "--name={}".format(input_folder),
+#                     "--dpm={}".format("sd"),
+#                     "--resolution={}".format(512),
+#                     "--load_trained",
+#                     "--num_tokens={}".format(num_tokens),
+#                     "--seed={}".format(2024),
+#                     "--guidance_scale={}".format(guidance_scale),
+#                     "--num_sampling_step={}".format(num_sampling_steps),
+#                     "--strength={}".format(strength),
+#                     "--edge_thickness={}".format(edge_thickness),
+#                     "--num_imgs={}".format(2),
+#                     "--tgt_prompt={}".format(tgt_prompt) ,
+#                     "--tgt_index={}".format(tgt_idx)            
+#     ])
     
-    return Image.open(os.path.join(input_folder, "text", "out_text_0.png"))
+#     return Image.open(os.path.join(input_folder, "text", "out_text_0.png"))
 
 
-def run_optimization(
-        num_tokens,
-        embedding_learning_rate, 
-        max_emb_train_steps, 
-        diffusion_model_learning_rate, 
-        max_diffusion_train_steps,
-        train_batch_size,
-        gradient_accumulation_steps,
-        input_folder = "example_tmp"
-    ):
-    subprocess.run(["python", 
-                    "main.py" ,
-                    "--name={}".format(input_folder),
-                    "--dpm={}".format("sd"),
-                    "--resolution={}".format(512),
-                    "--num_tokens={}".format(num_tokens),
-                    "--embedding_learning_rate={}".format(embedding_learning_rate),
-                    "--diffusion_model_learning_rate={}".format(diffusion_model_learning_rate),
-                    "--max_emb_train_steps={}".format(max_emb_train_steps),
-                    "--max_diffusion_train_steps={}".format(max_diffusion_train_steps),
-                    "--train_batch_size={}".format(train_batch_size),
-                    "--gradient_accumulation_steps={}".format(gradient_accumulation_steps)
+# def run_optimization(
+#         num_tokens,
+#         embedding_learning_rate, 
+#         max_emb_train_steps, 
+#         diffusion_model_learning_rate, 
+#         max_diffusion_train_steps,
+#         train_batch_size,
+#         gradient_accumulation_steps,
+#         input_folder = "example_tmp"
+#     ):
+#     subprocess.run(["python", 
+#                     "main.py" ,
+#                     "--name={}".format(input_folder),
+#                     "--dpm={}".format("sd"),
+#                     "--resolution={}".format(512),
+#                     "--num_tokens={}".format(num_tokens),
+#                     "--embedding_learning_rate={}".format(embedding_learning_rate),
+#                     "--diffusion_model_learning_rate={}".format(diffusion_model_learning_rate),
+#                     "--max_emb_train_steps={}".format(max_emb_train_steps),
+#                     "--max_diffusion_train_steps={}".format(max_diffusion_train_steps),
+#                     "--train_batch_size={}".format(train_batch_size),
+#                     "--gradient_accumulation_steps={}".format(gradient_accumulation_steps)
                     
-    ])
-    return 
+#     ])
+#     return 
 
 
 def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
@@ -215,6 +215,7 @@ with gr.Blocks() as demo:
     true    = gr.State(True)
     false    = gr.State(False)
     block_flag = gr.State(0)
+    num_tokens_global = gr.State(5)
     with gr.Row():
         gr.Markdown("""# D-Edit""")
 
@@ -293,6 +294,7 @@ with gr.Blocks() as demo:
                 opt_flag = gr.State(0)
                 gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings (SD)</p>""")
                 num_tokens = gr.Number(value="5", label="num tokens to represent each object", interactive= True)
+                num_tokens_global = num_tokens
                 embedding_learning_rate = gr.Textbox(value="0.0001", label="Embedding optimization: Learning rate", interactive= True )
                 max_emb_train_steps =  gr.Number(value="200", label="embedding optimization: Training steps", interactive= True )
                 
@@ -380,7 +382,7 @@ with gr.Blocks() as demo:
                             run_main,
                             load_trained=True,
                             text=True,
-                            num_tokens = int(num_tokens),
+                            num_tokens = int(num_tokens_global.value),
                             guidance_scale = float(guidance_scale),
                             num_sampling_steps = int(num_sampling_steps),
                             strength = float(strength),
@@ -391,8 +393,15 @@ with gr.Blocks() as demo:
                         )
                         return run_edit_text()
                         
-                    add_button.click(run_edit_text, 
-                        inputs = [], 
+                    add_button.click(run_edit_text_wrapper, 
+                        inputs = [num_tokens_global,
+                                    guidance_scale,
+                                    num_sampling_steps,
+                                    strength ,
+                                    edge_thickness,
+                                    tgt_prompt ,
+                                    tgt_index
+                                ],        
                         outputs = [canvas_text_edit]
                     )
                     

main copy.py

@@ -1,480 +0,0 @@
-import os
-import torch
-import numpy as np
-import argparse
-from peft import LoraConfig
-from old.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
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--name",  type=str,required=True, default=None)
-parser.add_argument("--name_2", type=str,required=False, default=None)
-parser.add_argument("--dpm",   type=str,required=True, default="sd")
-parser.add_argument("--resolution",  type=int, default=1024)
-parser.add_argument("--seed",  type=int, default=42)
-parser.add_argument("--embedding_learning_rate",  type=float, default=1e-4)
-parser.add_argument("--max_emb_train_steps",  type=int, default=200)
-parser.add_argument("--diffusion_model_learning_rate", type=float, default=5e-5)
-parser.add_argument("--max_diffusion_train_steps", type=int, default=200)
-parser.add_argument("--train_batch_size",  type=int, default=1)
-parser.add_argument("--gradient_accumulation_steps",  type=int, default=1)
-parser.add_argument("--num_tokens",  type=int, default=1)
-
-
-parser.add_argument("--load_trained", default=False, action="store_true" )
-parser.add_argument("--num_sampling_steps",  type=int, default=50)
-parser.add_argument("--guidance_scale", type=float, default = 3 )
-parser.add_argument("--strength",  type=float, default=0.8)
-
-parser.add_argument("--train_full_lora", default=False, action="store_true" )
-parser.add_argument("--lora_rank",  type=int, default=4)
-parser.add_argument("--lora_alpha",  type=int, default=4)
-
-parser.add_argument("--prompt_auxin_list", nargs="+", type=str, default = None)
-parser.add_argument("--prompt_auxin_idx_list", nargs="+", type=int, default = None)
-
-# general editing configs
-parser.add_argument("--load_edited_mask", default=False, action="store_true")
-parser.add_argument("--load_edited_processed_mask", default=False, action="store_true")
-parser.add_argument("--edge_thickness", type=int, default=20)
-parser.add_argument("--num_imgs", type=int, default = 1 )
-parser.add_argument('--active_mask_list', nargs="+", type=int)
-parser.add_argument("--tgt_index",  type=int, default=None)
-
-# recon
-parser.add_argument("--recon", default=False, action="store_true" )
-parser.add_argument("--recon_an_item", default=False, action="store_true" )
-parser.add_argument("--recon_prompt",  type=str, default=None)
-
-# text-based editing
-parser.add_argument("--text", default=False, action="store_true")
-parser.add_argument("--tgt_prompt",  type=str, default=None)
-
-# image-based editing
-parser.add_argument("--image", default=False, action="store_true" )
-parser.add_argument("--src_index",  type=int, default=None)
-parser.add_argument("--tgt_name",   type=str, default=None)
-
-# mask-based move
-parser.add_argument("--move_resize", default=False, action="store_true" )
-parser.add_argument('--tgt_indices_list', nargs="+", type=int)
-parser.add_argument("--delta_x_list", nargs="+", type=int)
-parser.add_argument("--delta_y_list", nargs="+", type=int)
-parser.add_argument("--priority_list", nargs="+", type=int)
-parser.add_argument("--force_mask_remain", type=int, default=None)
-parser.add_argument("--resize_list", nargs="+", type=float)
-
-# remove
-parser.add_argument("--remove", default=False, action="store_true" )
-parser.add_argument("--load_edited_removemask", default=False, action="store_true")
-
-args = parser.parse_args()
-
-
-def run_main(
-    name=None,
-    name_2=None,
-    dpm="sd",
-    resolution=1024,
-    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="store_true" ,
-    num_sampling_steps=50,
-    guidance_scale= 3 ,
-    strength=0.8,
-
-    train_full_lora="store_true" ,
-    lora_rank=4,
-    lora_alpha=4,
-
-    prompt_auxin_list = None,
-    prompt_auxin_idx_list= None,
-
-    load_edited_mask="store_true",
-    load_edited_processed_mask="store_true",
-    edge_thickness=20,
-    num_imgs= 1 ,
-    active_mask_list = None,
-    tgt_index=None,
-
-    recon=False ,
-    recon_an_item=False,
-    recon_prompt=None,
-
-    text="store_true",
-    tgt_prompt=None,
-
-    image="store_true" ,
-    src_index=None,
-    tgt_name=None,
-
-    move_resize="store_true" ,
-    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(args.seed)
-    torch.manual_seed(args.seed)  
-    base_input_folder = "."
-    base_output_folder  = "."
-
-    input_folder = os.path.join(base_input_folder, args.name)
-
-
-    mask_list, mask_label_list = load_mask(input_folder)
-    assert mask_list[0].shape[0] == args.resolution, "Segmentation should be done on size {}".format(args.resolution)
-    try:
-        image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(args.resolution) ), size = args.resolution)
-    except:
-        image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(args.resolution) ), size = args.resolution)
-
-    if args.image:
-        input_folder_2 = os.path.join(base_input_folder, args.name_2)
-        mask_list_2, mask_label_list_2 = load_mask(input_folder_2)
-        assert mask_list_2[0].shape[0] == args.resolution, "Segmentation should be done on size {}".format(args.resolution)
-        try:
-            image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.png".format(args.resolution) ), size = args.resolution)
-        except:
-            image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.jpg".format(args.resolution) ), size = args.resolution)
-        output_dir = os.path.join(base_output_folder, args.name + "_" + args.name_2)
-        os.makedirs(output_dir, exist_ok = True)
-    else:
-        output_dir = os.path.join(base_output_folder, args.name)
-        os.makedirs(output_dir, exist_ok = True)
-
-    if args.dpm == "sd":
-        if args.image:
-            pipe = DEditSDPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = args.resolution, num_tokens = args.num_tokens)
-        else:
-            pipe = DEditSDPipeline(mask_list, mask_label_list, resolution = args.resolution, num_tokens = args.num_tokens)
-            
-    elif args.dpm == "sdxl":
-        if args.image:
-            pipe = DEditSDXLPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = args.resolution, num_tokens = args.num_tokens)
-        else:
-            pipe = DEditSDXLPipeline(mask_list, mask_label_list, resolution = args.resolution, num_tokens = args.num_tokens)
-
-    else:
-        raise NotImplementedError
-
-    set_string_list = pipe.set_string_list
-    if args.prompt_auxin_list is not None:
-        for auxin_idx, auxin_prompt in zip(args.prompt_auxin_idx_list, args.prompt_auxin_list):
-            set_string_list[auxin_idx] = auxin_prompt.replace("*", set_string_list[auxin_idx] )
-    print(set_string_list)
-
-    if args.image: 
-        set_string_list_2 = pipe.set_string_list_2
-        print(set_string_list_2)
-
-    if args.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 args.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=args.lora_rank,
-                    lora_alpha=args.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 args.dpm == "sdxl":
-            pipe.text_encoder_2.load_state_dict(text_encoder2_state_dict)
-    else:
-        if args.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 = args.gradient_accumulation_steps,
-                embedding_learning_rate = args.embedding_learning_rate,
-                max_emb_train_steps = args.max_emb_train_steps,
-                train_batch_size = args.train_batch_size,
-            )
-            
-            pipe.train_model_2imgs(
-                image_gt,
-                image_gt_2, 
-                set_string_list,
-                set_string_list_2,
-                gradient_accumulation_steps = args.gradient_accumulation_steps,
-                max_diffusion_train_steps = args.max_diffusion_train_steps,
-                diffusion_model_learning_rate = args.diffusion_model_learning_rate ,
-                train_batch_size =args.train_batch_size,
-                train_full_lora = args.train_full_lora,
-                lora_rank = args.lora_rank, 
-                lora_alpha = args.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 = args.gradient_accumulation_steps,
-                embedding_learning_rate = args.embedding_learning_rate,
-                max_emb_train_steps = args.max_emb_train_steps,
-                train_batch_size = args.train_batch_size,
-            )
-
-            pipe.train_model(
-                image_gt,
-                set_string_list,
-                gradient_accumulation_steps = args.gradient_accumulation_steps,
-                max_diffusion_train_steps = args.max_diffusion_train_steps,
-                diffusion_model_learning_rate = args.diffusion_model_learning_rate ,
-                train_batch_size = args.train_batch_size,
-                train_full_lora = args.train_full_lora,
-                lora_rank = args.lora_rank, 
-                lora_alpha = args.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 args.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 args.recon:
-        output_dir = os.path.join(output_dir, "recon")
-        os.makedirs(output_dir, exist_ok = True)
-        if args.recon_an_item:
-            mask_list = [torch.from_numpy(np.ones_like(mask_list[0].numpy()))]
-            tgt_string = set_string_list[args.tgt_index]
-            tgt_string = args.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 = args.guidance_scale,
-            num_sampling_steps = args.num_sampling_steps,
-            seed = args.seed,
-            num_imgs = args.num_imgs,
-            set_string_list = set_string_list,
-            mask_list = mask_list
-        )
-        
-    if args.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[args.tgt_index] = args.tgt_prompt
-        mask_active = torch.zeros_like(mask_list[0])
-        mask_active = mask_union_torch(mask_active, mask_list[args.tgt_index])
-        
-        if args.active_mask_list is not None:
-            for midx in args.active_mask_list:
-                mask_active = mask_union_torch(mask_active, mask_list[midx])
-
-        if args.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 = args.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 = args.guidance_scale,
-            strength = args.strength,
-            num_imgs = args.num_imgs,
-            mask_hard= mask_hard,
-            mask_soft = mask_soft,
-            mask_list = mask_list,
-            seed = args.seed,
-            num_sampling_steps = args.num_sampling_steps
-        )
-
-    if args.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 args.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 args.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, args.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 args.active_mask_list is not None:
-            for midx in args.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 = args.edge_thickness)
-        mask_hard = mask_substract_torch(mask_hard, mask_soft)    
-
-        pipe.inference_with_mask(
-            save_path, 
-            orig_image = image_gt,
-            guidance_scale = args.guidance_scale,
-            strength = args.strength,
-            num_imgs = args.num_imgs,
-            mask_hard= mask_hard,
-            mask_soft = mask_soft,
-            mask_list = mask_list_processed, 
-            seed = args.seed,
-            num_sampling_steps = args.num_sampling_steps
-        )
-
-    if args.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 (args.tgt_name, args.src_index, args.tgt_index):
-            if args.tgt_name == args.name:
-                set_string_list_tgt = set_string_list
-                set_string_list_src = set_string_list_2
-                image_tgt = image_gt
-                if args.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 args.tgt_name == args.name_2:
-                set_string_list_tgt = set_string_list_2
-                set_string_list_src = set_string_list
-                image_tgt = image_gt_2
-                if args.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[args.tgt_index] = set_string_list_src[args.src_index]
-            
-            mask_active = mask_list_tgt[args.tgt_index]
-            mask_frozen = (1-mask_active.float()).to(mask_active.device)
-            mask_soft = create_outer_edge_mask_torch(mask_active.cpu(), edge_thickness = args.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 = args.guidance_scale,
-                num_sampling_steps = args.num_sampling_steps,
-                mask_hard = mask_hard.cuda(),
-                mask_soft = mask_soft.cuda(), 
-                num_imgs = args.num_imgs,
-                orig_image = image_tgt,
-                strength = args.strength,
-            )
-
-    if args.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 args.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 args.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,
-                args.tgt_indices_list, 
-                args.delta_x_list,
-                args.delta_y_list, args.priority_list, 
-                force_mask_remain = args.force_mask_remain,
-                resize_list = args.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 = args.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 args.active_mask_list is not None:
-            for midx in args.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 = args.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 = args.strength,
-            orig_image = image_gt, 
-            guidance_scale = args.guidance_scale,
-            num_sampling_steps =  args.num_sampling_steps,
-            num_imgs = args.num_imgs,
-            mask_hard= mask_hard,
-            mask_soft = mask_soft,
-            mask_list = mask_list_processed,
-            seed = args.seed
-        )

main.py

@@ -64,6 +64,7 @@ def run_main(
     remove=False,
     load_edited_removemask=False
 ):
+
     torch.cuda.manual_seed_all(seed)
     torch.manual_seed(seed)  
     base_input_folder = "."
@@ -220,9 +221,9 @@ def run_main(
             set_string_list = set_string_list,
             mask_list = mask_list
         )
-        
+    
     if text:
-        print("Text-guided editing ")
+        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")

pipeline_dedit_sd.py

@@ -810,5 +810,5 @@ class DEditSDPipeline:
             seed = seed
         )
         save_images(x0, save_path)
-        # from PIL import Image
-        # return Image.open("example_tmp/text/out_text_0.png")
+        from PIL import Image
+        return Image.open("example_tmp/text/out_text_0.png")