demo interface changes

.gitattributes

@@ -1 +0,0 @@
-img2.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -3,25 +3,6 @@ title: "D-Edit"
 colorFrom: "red"
 colorTo: "blue"
 sdk: "gradio"
-sdk_version: 5.0.2
-python_version: "3.10"
 app_file: app.py
 pinned: false
-models:
-  - CompVis/stable-diffusion-v1-4
-tags:
-  - image-editing
-  - diffusion
-  - research
 ---
-
-# D-Edit
-
-Public research demo for **D-Edit: A Generic Framework for Image Editing**.
-
-- [Paper on arXiv](https://arxiv.org/abs/2403.04880)
-- [Hugging Face paper page](https://huggingface.co/papers/2403.04880)
-
-The demo first segments an input image and then optimizes object embeddings and the diffusion model for text-guided editing. Segmentation runs on CPU; the optimization and editing stage requires CUDA.
-
-The public configuration defaults to a training batch size of 1, limits queue concurrency to one request, and clears cached GPU memory between optimization and editing. These safeguards are intended to fit the workflow on a 16 GB GPU such as an Nvidia T4.

app.py

@@ -1,9 +1,6 @@
+
 import os
 import copy
-import gc
-import spaces
-import torch
-from main import run_main
 from PIL import Image
 import matplotlib 
 import numpy as np
@@ -13,19 +10,9 @@ from utils_mask import process_mask_to_follow_priority, mask_union, visualize_ma
 from pathlib import Path
 from PIL import Image
 from functools import partial
-import time
-
+from main import run_main
 LENGTH=512 #length of the square area displaying/editing images
 TRANSPARENCY = 150 # transparency of the mask in display
-DEFAULT_TRAIN_BATCH_SIZE = 1
-MAX_TRAIN_BATCH_SIZE = 4
-
-
-def clear_cuda_cache():
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
 
 def add_mask(mask_np_list_updated, mask_label_list):
     mask_new = np.zeros_like(mask_np_list_updated[0])
@@ -45,24 +32,89 @@ def create_segmentation(mask_np_list):
     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())
 
-def run_segmentation_wrapper(image):
+    return mask_np_list, mask_label_list
+
+def load_image_ui(load_edit, input_folder="example_tmp"):
     try:
-        print(image.shape)
-        image, mask_np_list,mask_label_list = run_segmentation(image)
-        #image = image.convert('RGB')
+        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))  
-        max_val = len(mask_np_list)-1
-        sliderup = gr.Slider(value = 0, minimum=0, maximum=max_val, step=1, visible=True)
-        gr.Info('Segmentation finish. Select mask id and move to the next step.')
-        return image, segmentation, mask_np_list, mask_label_list, image, sliderup, sliderup , 'Segmentation finish. Select mask id and move to the next step.'
-    except Exception as e:
-        print(e)
-        sliderup = gr.Slider(value = 0, minimum=0, maximum=1, step=1, visible=False)
-        gr.Warning('Please upload an image before proceeding.')
-        return None,None,None,None,None, sliderup, sliderup , 'Please upload an image before proceeding.'
-        
+        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=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"))
+
+
+# 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)
@@ -74,8 +126,11 @@ def transparent_paste_with_mask(backimg, foreimg, mask_np,transparency = 128):
     bimg_np = np.array(bimg)
     mask_np = mask_np[:,:,np.newaxis]
 
-    new_img_np = bimg_np*mask_np + (1-mask_np)* backimg_solid_np
-    return Image.fromarray(np.uint8(new_img_np))
+    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:
@@ -106,32 +161,17 @@ 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)-1:
-        return image, "out of range", ""
+
+    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]
-        index = mask_label.rfind('-')
-        mask_label = mask_label[:index]
-        if mask_label == 'handbag':
-            mask_prompt = "white handbag"
-        elif mask_label == 'person':
-            mask_prompt = "little boy"
-        elif mask_label == 'wall-other-merged':
-            mask_prompt = "white wall"
-        elif mask_label == 'table-merged':
-            mask_prompt = "table"
-        else:
-            mask_prompt = mask_label
         segmentation = create_segmentation(mask_np_list_updated)
         new_image = transparent_paste_with_mask(image, segmentation, mask_np, transparency = TRANSPARENCY)
-    gr.Info('Edit '+ mask_label)
-    return new_image, mask_label, mask_prompt
-def image_change():
-    return gr.Slider(value = 0, minimum=0, maximum=1, step=1, visible=False)
+    return new_image, mask_label
 
 def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder="example_tmp"):
-    print(mask_np_list_updated)
     try: 
         assert np.all(sum(mask_np_list_updated)==1)
     except:
@@ -146,7 +186,6 @@ def save_as_orig_mask(mask_np_list_updated, mask_label_list, input_folder="examp
     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"):
-    print(mask_np_list_updated)
     try: 
         assert np.all(sum(mask_np_list_updated)==1)
     except:
@@ -158,30 +197,12 @@ def save_as_edit_mask(mask_np_list_updated, mask_label_list, input_folder="examp
     savepath = os.path.join(input_folder, "seg_edited.png")
     visualize_mask_list_clean(mask_np_list_updated, savepath)
   
-
-
-def button_clickable(is_clickable):
-    return gr.Button(interactive=is_clickable)
-
-
-
-def load_pil_img():
-    from PIL import Image
-    return Image.open("example_tmp/text/out_text_0.png")
-
-def change_image(img):
-    return None
-                
-
+  
 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()
@@ -197,227 +218,203 @@ with gr.Blocks() as demo:
     with gr.Row():
         gr.Markdown("""# D-Edit""")
 
-
-    with gr.Row():
-        with gr.Column():
-            canvas = gr.Image(value = None, type="numpy",  label="Show Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-            example_inps = [['./img.png'],['./img2.png'],['./img3.png'],['./img4.png']]
-            gr.Examples(examples=example_inps, inputs=[canvas],
-                        label='examples', cache_examples='lazy', outputs=[],
-                        fn=change_image)
-            gr.Markdown(f"Each image must first undergo segmentation. Afterwards, you can modify the \n mask ID and the prompt for image editing, then proceed with the editing process. \n The link of D-edit paper: [https://arxiv.org/abs/2403.04880v2](https://arxiv.org/abs/2403.04880v2), [https://huggingface.co/papers/2403.04880](https://huggingface.co/papers/2403.04880)")
-
-        with gr.Column():
-            result_info0 = gr.Text(label="Response")
-            segment_button  = gr.Button("Step 1. Run segmentation")
-            flag = gr.State(False)
-
+    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, block_flag] , 
+                        [block_flag] )
+                
+                text_button  = gr.Button("Waiting 1.1 to complete")
+                text_button.click(load_image_ui, 
+                        [ false] , 
+                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
+
+                load_edit_button = gr.Button("Waiting 1.1 to complete")    
+                load_edit_button.click(load_image_ui, 
+                        [ true] , 
+                        [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas] )
+                
+                show_segment = gr.Checkbox(label = "Waiting 1.1 to complete")
+                flag = gr.State(False)
+                show_segment.select(show_segmentation,
+                                    [image_loaded, segmentation, flag], 
+                                    [canvas, flag])
+                def show_more_buttons():
+                    return gr.Button("1.2 Load original masks"), gr.Button("1.2 Load edited masks")   , gr.Checkbox(label = "Show Segmentation") 
+                block_flag.change(show_more_buttons, [], [text_button,load_edit_button,show_segment ])
+                
+                
             # mask_np_list_updated.value = copy.deepcopy(mask_np_list.value) #!!
             mask_np_list_updated = mask_np_list
-            gr.Markdown("""<p style="text-align: center; font-size: 20px">Edit Mask (Do not change it during the editing process)</p>""")
-            slider =  gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
-            label = gr.Text(label='label')
-
-        
-            
+            with gr.Column():
+                gr.Markdown("""<p style="text-align: center; font-size: 20px">Edit Mask (Optional)</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]
+                    )
 
-            result_info = gr.Text(label="Response")
-            
-            opt_flag = gr.State(0)
-            gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings</p>""")
-            with gr.Accordion(label="Advanced settings", open=False):
+                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():
+                
+                txt_box = gr.Textbox("Click to start optimization...", interactive = False)
+                
+                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.00025", label="Embedding optimization: Learning rate", interactive= True )
-                max_emb_train_steps =  gr.Number(value="6", label="embedding optimization: Training steps", 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.0002", label="UNet Optimization: Learning rate", interactive= True )
-                max_diffusion_train_steps = gr.Number(value="28", 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.Number(
-                    value=DEFAULT_TRAIN_BATCH_SIZE,
-                    minimum=1,
-                    maximum=MAX_TRAIN_BATCH_SIZE,
-                    precision=0,
-                    label="Batch size",
-                    info="Keep this at 1 for T4 and other 16 GB GPUs.",
-                    interactive=True,
-                )
-                gradient_accumulation_steps=gr.Number(value=2, minimum=1, precision=0, label="Gradient accumulation", interactive= True )
-            
-            def run_optimization_wrapper (
-                    mask_np_list,
-                    mask_label_list,
-                    image,
-                    opt_flag,                     
-                    num_tokens,
-                    embedding_learning_rate , 
-                    max_emb_train_steps , 
-                    diffusion_model_learning_rate , 
-                    max_diffusion_train_steps,
-                    train_batch_size,
-                    gradient_accumulation_steps,
-            ):
-                try:
+                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 (
+                        opt_flag,                     
+                        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,  
-                        mask_np_list=mask_np_list, 
-                        mask_label_list=mask_label_list,
-                        image_gt=np.array(image),
+                        run_main,                 
                         num_tokens=int(num_tokens),
                         embedding_learning_rate = float(embedding_learning_rate), 
-                        max_emb_train_steps = min(int(max_emb_train_steps),50), 
+                        max_emb_train_steps = int(max_emb_train_steps), 
                         diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
-                        max_diffusion_train_steps = min(int(max_diffusion_train_steps),100),
-                        train_batch_size=max(1, min(int(train_batch_size), MAX_TRAIN_BATCH_SIZE)),
-                        gradient_accumulation_steps=max(1, int(gradient_accumulation_steps))
+                        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()
-                    gr.Info("Optimization Finished! Move to the next step.")
-                    return "Optimization finished! Move to the next step."#,gr.Button("Step 3. Run Editing",interactive = True)
-                except torch.cuda.OutOfMemoryError as exc:
-                    print(f"CUDA out of memory during optimization: {exc}")
-                    clear_cuda_cache()
-                    message = "GPU memory is full. Keep Batch size at 1 and try again."
-                    gr.Error(message)
-                    return f"Error: {message}"
-                except Exception as exc:
-                    print(f"Optimization failed: {exc}")
-                    message = "Optimization failed. Check the image and mask, then try again."
-                    gr.Error(message)
-                    return f"Error: {message}"
-
-
-
-    if 1:
-        with gr.Row():
-            with gr.Column():
-                canvas_text_edit = gr.Image(value = None, type = "pil", label="Editing results", show_label=True,visible = True)
-                # canvas_text_edit = gr.Gallery(label = "Edited results")
-            
-            with gr.Column():
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting</p>""")
-                tgt_prompt =  gr.Textbox(value="text prompt", label="Editing: Text prompt", interactive= True )
-                with gr.Accordion(label="Advanced settings", open=False):
-                    slider2 = gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
-                    guidance_scale = gr.Textbox(value="5", label="Editing: CFG guidance scale", interactive= True )
-                    num_sampling_steps = gr.Number(value="20", label="Editing: Sampling steps", interactive= True )
+                    return opt_flag+1
+                    
+                add_button.click(run_optimization_wrapper, 
+                        inputs = [
+                            opt_flag,
+                            num_tokens,
+                            embedding_learning_rate , 
+                            max_emb_train_steps , 
+                            diffusion_model_learning_rate , 
+                            max_diffusion_train_steps,
+                            train_batch_size,
+                            gradient_accumulation_steps
+                        ], 
+                        outputs = [opt_flag]
+                )
+                
+                def change_text(txt_box):
+                    return gr.Textbox("Optimization Finished!", interactive = False)
+                def change_text2(txt_box):
+                    return gr.Textbox("Start optimization, check logs for progress...", interactive = False)
+                add_button.click(change_text2, txt_box, txt_box)
+                opt_flag.change(change_text, txt_box, txt_box)
+
+    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("Step 2. Run Editing",interactive = True)
-                def run_edit_text_wrapper(
-                        mask_np_list,
-                        mask_label_list,
-                        image,
-                        num_tokens,
-                        guidance_scale,
-                        num_sampling_steps ,
-                        strength ,
-                        edge_thickness,
-                        tgt_prompt ,
-                        tgt_index
-                ):
+                    
+                    add_button  = gr.Button("Run Editing")
+                    def run_edit_text_wrapper(
+                            num_tokens,
+                            guidance_scale,
+                            num_sampling_steps ,
+                            strength ,
+                            edge_thickness,
+                            tgt_prompt ,
+                            tgt_index
+                    ):
+                            
+                        run_edit_text = partial(
+                            run_main,
+                            load_trained=True,
+                            text=True,
+                            num_tokens = int(num_tokens_global.value),
+                            guidance_scale = float(guidance_scale),
+                            num_sampling_steps = int(num_sampling_steps),
+                            strength = float(strength),
+                            edge_thickness = int(edge_thickness),
+                            num_imgs = 1,
+                            tgt_prompt = tgt_prompt,
+                            tgt_index = int(tgt_index)
+                        )
+                        return run_edit_text()
                         
-                    run_edit_text = partial(
-                        run_main,
-                        mask_np_list=mask_np_list, 
-                        mask_label_list=mask_label_list,
-                        image_gt=np.array(image),
-                        load_trained=True,
-                        text=True,
-                        num_tokens = int(num_tokens),
-                        guidance_scale = float(guidance_scale),
-                        num_sampling_steps = int(num_sampling_steps),
-                        strength = float(strength),
-                        edge_thickness = int(edge_thickness),
-                        num_imgs = 1,
-                        tgt_prompt = tgt_prompt,
-                        tgt_index = int(tgt_index)
+                    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]
+                    )
+                    
+                    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]
                     )
-                    run_edit_text()
-                    gr.Info('Image editing completed.')
-                    return load_pil_img()
-
-
-
-            @spaces.GPU(duration=120)
-            def run_total_wrapper(mask_np_list, mask_label_list, image_loaded, opt_flag, num_tokens, embedding_learning_rate, max_emb_train_steps, diffusion_model_learning_rate, max_diffusion_train_steps, train_batch_size, gradient_accumulation_steps, num_tokens_global, guidance_scale, num_sampling_steps, strength, edge_thickness, tgt_prompt, slider2):
-                if not torch.cuda.is_available():
-                    message = "GPU hardware is required for Step 2. Please try again after the Space receives a GPU."
-                    gr.Error(message)
-                    return f"Error: {message}", None
-                if image_loaded is None or not mask_np_list:
-                    message = "Run Step 1 segmentation before starting the edit."
-                    gr.Error(message)
-                    return f"Error: {message}", None
-
-                try:
-                    result_info = run_optimization_wrapper(mask_np_list, mask_label_list, image_loaded, opt_flag, num_tokens, embedding_learning_rate, max_emb_train_steps, diffusion_model_learning_rate, max_diffusion_train_steps, train_batch_size, gradient_accumulation_steps)
-                    if not result_info.startswith("Optimization finished"):
-                        return result_info, None
-                    clear_cuda_cache()
-                    canvas_text_edit = run_edit_text_wrapper(mask_np_list, mask_label_list, image_loaded, num_tokens_global, guidance_scale, num_sampling_steps, strength, edge_thickness, tgt_prompt, slider2)
-                    return result_info, canvas_text_edit
-                except torch.cuda.OutOfMemoryError as exc:
-                    print(f"CUDA out of memory during editing: {exc}")
-                    message = "GPU memory is full. Keep Batch size at 1 and reduce the training steps."
-                    gr.Error(message)
-                    return f"Error: {message}", None
-                except Exception as exc:
-                    print(f"Editing failed: {exc}")
-                    message = "Editing failed. Please retry with the default settings."
-                    gr.Error(message)
-                    return f"Error: {message}", None
-                finally:
-                    clear_cuda_cache()
-
-
-            add_button.click(
-                run_total_wrapper, 
-                inputs=[
-                    mask_np_list,
-                    mask_label_list,
-                    image_loaded,
-                    opt_flag,
-                    num_tokens,
-                    embedding_learning_rate,
-                    max_emb_train_steps,
-                    diffusion_model_learning_rate,
-                    max_diffusion_train_steps,
-                    train_batch_size,
-                    gradient_accumulation_steps,
-                    num_tokens_global,
-                    guidance_scale,
-                    num_sampling_steps,
-                    strength,
-                    edge_thickness,
-                    tgt_prompt,
-                    slider2
-                ],
-                outputs=[result_info, canvas_text_edit],
-            )
-
-
-
-
-        canvas.upload(image_change, inputs=[], outputs=[slider])       
-
-        slider.release(slider_release, 
-                        inputs = [slider, image_loaded,   mask_np_list_updated, mask_label_list], 
-                        outputs= [canvas, label,tgt_prompt])
-
-        slider.change(
-            lambda x: x,
-            inputs=[slider],
-            outputs=[slider2]
-        )
-
 
-        segment_button.click(run_segmentation_wrapper, 
-                [canvas] ,
-                [image_loaded, segmentation,  mask_np_list, mask_label_list, canvas, slider, slider2, result_info0] )
 
 
 
-demo.queue(max_size=4, default_concurrency_limit=1).launch()
+demo.queue().launch(share=True, debug=True)

main.py

@@ -12,9 +12,6 @@ from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualiz
 def run_main(
     name="example_tmp",
     name_2=None,
-    mask_np_list=None, 
-    mask_label_list=None,
-    image_gt=None,
     dpm="sd",
     resolution=512,
     seed=42,
@@ -74,17 +71,13 @@ def run_main(
     base_output_folder  = "."
 
     input_folder = os.path.join(base_input_folder, name)
-    mask_list = []
-    for mask_np in mask_np_list:
-        mask = torch.from_numpy(mask_np.astype(np.uint8))
-        mask_list.append(mask)
-    
-    #mask_list, mask_label_list = load_mask(input_folder)
+
+    mask_list, mask_label_list = load_mask(input_folder)
     assert mask_list[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution)
-    #try:
-    #    image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(resolution) ), size = resolution)
-    #except:
-    #    image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(resolution) ), size = resolution)
+    try:
+        image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(resolution) ), size = resolution)
+    except:
+        image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(resolution) ), size = resolution)
 
     if image:
         input_folder_2 = os.path.join(base_input_folder, name_2)

requirements.txt

@@ -1,16 +1,12 @@
-huggingface_hub
-matplotlib==3.9.4
-
-accelerate==0.34.2
-diffusers==0.30.2
-numpy==1.26.4
-torch==2.8.0
-torchvision==0.23.0
+torch==2.2.0
+torchvision==0.17.0
 transformers==4.37.2
-xformers==0.0.32.post2
+accelerate==0.23.0
+gradio==3.41.1
+xformers==0.0.24
+diffusers==0.26.3
 scipy
-setuptools
 tqdm
+numpy
 safetensors
-peft==0.13.2
-pydantic==2.10.6
+peft

segment.py

@@ -3,7 +3,6 @@ from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
 from PIL import Image
 import torch
 from collections import defaultdict
-from functools import lru_cache
 import matplotlib.pyplot as plt
 from matplotlib import cm
 import matplotlib.patches as mpatches
@@ -11,17 +10,6 @@ import os
 import numpy as np
 import argparse
 import matplotlib
-import gradio as gr
-
-
-@lru_cache(maxsize=1)
-def get_segmentation_model():
-    model_id = "facebook/mask2former-swin-base-coco-panoptic"
-    processor = AutoImageProcessor.from_pretrained(model_id)
-    model = Mask2FormerForUniversalSegmentation.from_pretrained(model_id)
-    model.eval()
-    return processor, model
-
 
 def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):
     if type(image_path) is str:
@@ -56,18 +44,14 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
     instances_counter = defaultdict(int)
     handles = []
     label_list = []
-
-    mask_np_list = []
-    
     if not noseg:
         if torch.min(segmentation) == 0: 
             mask = segmentation==0
             mask = mask.cpu().detach().numpy()   # [512,512]   bool
-            print(mask.shape)
             segment_label = "rest"
+            np.save( os.path.join(save_folder, "mask{}_{}.npy".format(0,"rest")) , mask)
             color = viridis(0)
             label = f"{segment_label}-{0}"
-            mask_np_list.append(mask)  
             handles.append(mpatches.Patch(color=color, label=label))
             label_list.append(label)
        
@@ -77,11 +61,10 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             if torch.min(segmentation) != 0: 
                 segment_id -= 1
             mask = mask.cpu().detach().numpy()   # [512,512] bool
-            print(mask.shape)
-            mask_np_list.append(mask) 
+            
             segment_label = model.config.id2label[segment['label_id']]
             instances_counter[segment['label_id']] += 1
-
+            np.save( os.path.join(save_folder, "mask{}_{}.npy".format(segment_id,segment_label)) , mask)
             color = viridis(segment_id)
             
             label = f"{segment_label}-{segment_id}"
@@ -89,10 +72,8 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             label_list.append(label)
     else:
         mask = np.full(segmentation.shape, True)
-        print(mask.shape)
-            
         segment_label = "all"
-        mask_np_list.append(mask) 
+        np.save( os.path.join(save_folder, "mask{}_{}.npy".format(0,"all")) , mask)
         color = viridis(0)
         label = f"{segment_label}-{0}"
         handles.append(mpatches.Patch(color=color, label=label))
@@ -103,16 +84,17 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
     # plt.savefig(os.path.join(save_folder, 'mask_clear.png'), dpi=500)
     ax.legend(handles=handles)
     plt.savefig(os.path.join(save_folder, 'seg_init.png'), dpi=500 )
-    plt.close(fig)
     print("; ".join(label_list))
-    return mask_np_list,label_list
 
 
-def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
+
+
+def run_segmentation(image, block_flag, name="example_tmp", size = 512, noseg=False):
 
     base_folder_path = "."
 
-    processor, model = get_segmentation_model()
+    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
 
 
     # input_folder = os.path.join(base_folder_path, name )
@@ -123,15 +105,15 @@ def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
     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)))
+    image.save(os.path.join(name,"img_{}.png".format(size)))
     inputs = processor(image, return_tensors="pt")
-    with torch.inference_mode():
+    with torch.no_grad():
         outputs = model(**inputs)
 
     panoptic_segmentation = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
     save_folder = os.path.join(base_folder_path, name)
     os.makedirs(save_folder, exist_ok=True)
-    mask_list,label_list = draw_panoptic_segmentation(**panoptic_segmentation, save_folder = save_folder, noseg = noseg, model = model)
+    draw_panoptic_segmentation(**panoptic_segmentation, save_folder = save_folder, noseg = noseg, model = model)
     print("Finish segment")
-    #block_flag += 1
-    return  image,mask_list,label_list#, gr.Button.update("1.2 Load edited masks",visible = True), gr.Checkbox.update(label = "Show Segmentation",visible = True) 
+    block_flag += 1
+    return  block_flag

utils.py

@@ -249,6 +249,7 @@ def load_mask (input_folder):
     except:
         print("please check mask")
         # plt.imsave( "out_mask.png", mask_list_edit[0]) 
+        import pdb; pdb.set_trace()
     return mask_list, mask_label_list
 
 def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):