second

app copy.py

@@ -317,7 +317,7 @@ with gr.Blocks() as demo:
             canvas_text_edit =  gr.State() # store mask
             with gr.Row():
                 with gr.Column():
-                    canvas_text_edit = gr.Image(value = None, label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
+                    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():
@@ -342,8 +342,9 @@ with gr.Blocks() as demo:
                             tgt_idx,
                             guidance_scale
                         ], 
-                        outputs = [canvas_text_edit]
+                        outputs = []
                     )
 
 
+
 demo.queue().launch(share=True, debug=True)

app.py

@@ -10,7 +10,8 @@ from utils_mask import process_mask_to_follow_priority, mask_union, visualize_ma
 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
 
@@ -32,7 +33,7 @@ def create_segmentation(mask_np_list):
     segmentation = Image.fromarray(np.uint8(segmentation*255))
     return segmentation
 
-def load_mask_ui(input_folder,load_edit = False):
+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:
@@ -44,28 +45,29 @@ def load_mask_ui(input_folder,load_edit = False):
 
     return mask_np_list, mask_label_list
 
-def load_image_ui(input_folder, load_edit):
+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.png", "img_1024.png", "img_512.png"]:
+            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)  
+        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(
-        input_folder,
         num_tokens,
         num_sampling_steps,
         strength,
         edge_thickness,
         tgt_prompt,
         tgt_idx,
-        guidance_scale
+        guidance_scale,
+        input_folder="example_tmp"
     ):
     subprocess.run(["python", 
                     "main.py" ,
@@ -89,14 +91,14 @@ def run_edit_text(
 
 
 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
+        gradient_accumulation_steps,
+        input_folder = "example_tmp"
     ):
     subprocess.run(["python", 
                     "main.py" ,
@@ -124,6 +126,7 @@ def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
 
     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) 
@@ -159,6 +162,7 @@ def edit_mask_add(canvas,  image, idx, mask_np_list):
     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:
@@ -168,7 +172,7 @@ def slider_release(index, image,  mask_np_list_updated, mask_label_list):
         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):
+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:
@@ -182,7 +186,7 @@ def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder):
     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):
+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:
@@ -195,6 +199,10 @@ def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder):
     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
@@ -213,8 +221,7 @@ with gr.Blocks() as demo:
     with gr.Tab(label="1 Edit mask"):
         with gr.Row():
             with gr.Column():
-                canvas = gr.Image(value = None, type="pil",  label="Draw Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-                input_folder = gr.Textbox(value="example1", label="input folder", interactive= True, )
+                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, 
@@ -223,23 +230,22 @@ with gr.Blocks() as demo:
 
                 text_button  = gr.Button("1.2 Load original masks")
                 text_button.click(load_image_ui, 
-                        [input_folder, false] , 
+                        [ 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] , 
+                        [ 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)
-
+ 
+            # 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)
@@ -256,17 +262,17 @@ with gr.Blocks() as demo:
 
                 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] , 
+                        [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, input_folder] , 
+                        [mask_np_list_updated,  mask_label_list] , 
                         [] )  
                 
                 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")    
@@ -274,70 +280,88 @@ with gr.Blocks() as demo:
                         [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.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.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 )
+            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="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 )
+                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.Textbox(value="5", label="Batch size", interactive= True )
-    #             gradient_accumulation_steps=gr.Textbox(value="5", label="Gradient accumulation", 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")
-    #             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, label="Editing results", show_label=True, height=LENGTH, width=LENGTH)
-    #                 # canvas_text_edit = gr.Gallery(label = "Edited results")
+                add_button  = gr.Button("Run optimization")
+
+                run_optimization = partial(
+                    run_main,                 
+                    num_tokens=int(num_tokens.value),
+                    embedding_learning_rate = float(embedding_learning_rate.value), 
+                    max_emb_train_steps = int(max_emb_train_steps.value), 
+                    diffusion_model_learning_rate= float(diffusion_model_learning_rate.value), 
+                    max_diffusion_train_steps = int(max_diffusion_train_steps.value),
+                    train_batch_size=int(train_batch_size.value),
+                    gradient_accumulation_steps=int(gradient_accumulation_steps.value)
+                )
+                add_button.click(run_optimization, 
+                        inputs = [], 
+                        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>""")
+                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]
+                    )
                     
-    #                 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 )
+                    def load_pil_img():
+                        from PIL import Image
+                        return Image.open("example_tmp/text/out_text_0.png")
                     
-    #                 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 = [canvas_text_edit]
-    #                 )
+                    load_button  = gr.Button("Load results")
+                    load_button.click(load_pil_img, 
+                        inputs = [], 
+                        outputs = [canvas_text_edit]
+                    )
+
+
 
 
 demo.queue().launch(share=True, debug=True)

main copy.py

@@ -0,0 +1,480 @@
+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

@@ -9,416 +9,406 @@ 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()
-
-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)
+def run_main(
+    name="example_tmp",
+    name_2=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, 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_2 = load_image(os.path.join(input_folder_2, "img_{}.png".format(args.resolution) ), size = args.resolution)
+        image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(resolution) ), size = 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)
+        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:
-        pipe = DEditSDPipeline(mask_list, mask_label_list, resolution = args.resolution, num_tokens = args.num_tokens)
+        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) 
         
-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)
+        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:
-        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"],
+        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.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 = 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
+            )
+
         
-        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
-        )
+        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 )
         
-    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,
+
+    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.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
+        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
         )
 
-    
-    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 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_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
+        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
             
-        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
+        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:
-            exit("tgt_name should be either name or name_2")
-            
-        set_string_list_tgt[args.tgt_index] = set_string_list_src[args.src_index]
+            # 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)
         
-        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]
+        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,
-            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,
+            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 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 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 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
+        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
         )
-        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
-    )

pipeline_dedit_sd.py

@@ -27,11 +27,11 @@ class DEditSDPipeline:
         mask_label_list, 
         mask_list_2 = None,
         mask_label_list_2 = None, 
-        resolution = 1024,
+        resolution = 512,
         num_tokens = 1
     ):
         super().__init__()
-        model_id = "./stable-diffusion-v1-5"
+        model_id = "CompVis/stable-diffusion-v1-4"
         self.model_id = model_id
         self.tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", use_fast=False)
         text_encoder_cls_one = import_model_class_from_model_name_or_path(model_id, subfolder = "text_encoder")
@@ -810,4 +810,5 @@ class DEditSDPipeline:
             seed = seed
         )
         save_images(x0, save_path)
-        return x0
+        # from PIL import Image
+        # return Image.open("example_tmp/text/out_text_0.png")

segment.py

@@ -102,7 +102,8 @@ def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
     #     image = load_image(os.path.join(input_folder, "img.png" ), size = size)
     # except:
     #     image = load_image(os.path.join(input_folder, "img.jpg" ), size = size)
-    # image =Image.fromarray(image)
+    image =Image.fromarray(image)
+    image = image.resize((size, size))
     os.makedirs(name, exist_ok=True)
     image.save(os.path.join(name,"img_{}.png".format(size)))
     inputs = processor(image, return_tensors="pt")