update gradio

app.py

@@ -56,7 +56,7 @@ add_sampling_metadata = True
     
 @spaces.GPU(duration=120)
 @torch.inference_mode()
-def edit(init_image, brush_canvas,
+def edit(brush_canvas,
             source_prompt, target_prompt, 
             inversion_num_steps, denoise_num_steps, 
             skip_step, 
@@ -66,26 +66,20 @@ def edit(init_image, brush_canvas,
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.cuda.empty_cache()
     
+    rgba_init_image = brush_canvas["background"]
+    init_image = rgba_init_image[:,:,:3]
     shape = init_image.shape        
     height = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16
     width = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16
-
     init_image = init_image[:height, :width, :]
-    brush_canvas = brush_canvas["composite"][:,:,:3][:height, :width, :]
-    # 如果brush_Canvas是三通道黑白图，说明就是输入的mask
+    rgba_init_image = rgba_init_image[:height, :width, :]
+
+    rgba_mask = brush_canvas["layers"][0][:height, :width, :]
+    mask = rgba_mask[:,:,3]/255
+    mask = mask.astype(int)
     
-    if np.all(brush_canvas[:,:,0] == brush_canvas[:,:,1]) and np.all(brush_canvas[:,:,1] == brush_canvas[:,:,2]):
-        mask = brush_canvas[:,:,0]/255
-        mask = mask.astype(int)
-    else:
-        mask = np.any(init_image != brush_canvas, axis=-1)  # 得到一个二维的布尔数组
-        mask = mask.astype(int)
-    mask_array = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
-    mask_array[:,:,0] = mask * 255  # R
-    mask_array[:,:,3] = mask * 128  # A (半透明，128表示50%透明度)
-    mask_image = Image.fromarray(mask_array, 'RGBA')     
-    original_image = Image.fromarray(np.concatenate((init_image, np.full((height, width, 1), 255, dtype=np.uint8)), axis=2), 'RGBA')
-    masked_image = Image.alpha_composite(original_image, mask_image)
+    rgba_mask[:,:,3] = rgba_mask[:,:,3]//2
+    masked_image = Image.alpha_composite(Image.fromarray(rgba_init_image, 'RGBA'), Image.fromarray(rgba_mask, 'RGBA'))
     mask = torch.from_numpy(mask).unsqueeze(0).unsqueeze(0).to(torch.bfloat16).to(device)
     
     init_image = encode(init_image, device).to(device)
@@ -180,8 +174,8 @@ def create_demo(model_name: str):
     
         🔔🔔[<b>Important</b>] Editing steps:<br>
         1️⃣ Upload your image that needs to be edited (The resolution is expected be less than 1360*768, or the memory of GPU may be not enough.) <br>
-        2️⃣ Re-upload the original image and use the brush tool to draw your mask area. <br>
-        3️⃣ Fill in your source prompt and target prompt, then adjust the hyperparameters. <br>
+        2️⃣ Fill in your source prompt and use the brush tool to draw your mask area. <br>
+        3️⃣ Fill in your target prompt, then adjust the hyperparameters. <br>
         4️⃣ Click the "Edit" button to generate your edited image! <br>
         """
     article = r"""
@@ -204,15 +198,12 @@ def create_demo(model_name: str):
                 inversion_num_steps = gr.Slider(1, 50, 28, step=1, label="Number of inversion steps")
                 target_prompt = gr.Textbox(label="Target Prompt", value='In a cluttered wooden cabin, a workbench holds a green neon sign that reads "I love iccv"' )
                 denoise_num_steps = gr.Slider(1, 50, 28, step=1, label="Number of denoise steps")
-                init_image = gr.Image(label="Input Image", visible=True)
                 brush_canvas = gr.ImageEditor(label="Brush Canvas",
                                                 sources=('upload'), 
-                                                brush=gr.Brush(default_size=10,
-                                                               default_color="#000000"),
+                                                brush=gr.Brush(colors=["#ff0000"],color_mode='fixed'),
                                                 interactive=True,
-                                                container=True,
                                                 transforms=[],
-                                                height="auto",
+                                                container=True,
                                                 format='png',scale=1)
                 
                 edit_btn = gr.Button("edit")
@@ -220,9 +211,8 @@ def create_demo(model_name: str):
                 
             with gr.Column():
                 with gr.Accordion("Advanced Options", open=True):
-                    # num_steps = gr.Slider(1, 30, 25, step=1, label="Number of steps")
 
-                    skip_step = gr.Slider(0, 30, 4, step=1, label="Number of inject steps")
+                    skip_step = gr.Slider(0, 30, 4, step=1, label="Number of skip steps")
                     inversion_guidance = gr.Slider(1.0, 10.0, 1.5, step=0.1, label="inversion Guidance", interactive=not is_schnell)
                     denoise_guidance = gr.Slider(1.0, 10.0, 5.5, step=0.1, label="denoise Guidance", interactive=not is_schnell)
                     seed = gr.Textbox('0', label="Seed (-1 for random)", visible=True)
@@ -234,7 +224,7 @@ def create_demo(model_name: str):
                 output_image = gr.Image(label="Generated Image")
         edit_btn.click(
             fn=edit,
-            inputs=[init_image, brush_canvas,
+            inputs=[brush_canvas,
                     source_prompt, target_prompt, 
                     inversion_num_steps, denoise_num_steps, 
                     skip_step,