Update app.py

app.py

@@ -15,44 +15,44 @@ from diffusers.utils import load_image
 from pipeline_flux_control_removal import FluxControlRemovalPipeline
 pipe = None
 torch.set_grad_enabled(False)
-image_path = mask_path =None
-# image_examples = [
-#     [
-#         "example/image/3c43156c-2b44-4ebf-9c47-7707ec60b166.png",
-#         "example/mask/3c43156c-2b44-4ebf-9c47-7707ec60b166.png"
-#     ],
-#     [
-#         "example/image/0e5124d8-fe43-4b5c-819f-7212f23a6d2a.png",
-#         "example/mask/0e5124d8-fe43-4b5c-819f-7212f23a6d2a.png"
-#     ],
-#     [
-#         "example/image/0f900fe8-6eab-4f85-8121-29cac9509b94.png",
-#         "example/mask/0f900fe8-6eab-4f85-8121-29cac9509b94.png"
-#     ],
-#     [
-#         "example/image/3ed1ee18-33b0-4964-b679-0e214a0d8848.png",
-#         "example/mask/3ed1ee18-33b0-4964-b679-0e214a0d8848.png"
-#     ],
-#     [
-#         "example/image/9a3b6af9-c733-46a4-88d4-d77604194102.png",
-#         "example/mask/9a3b6af9-c733-46a4-88d4-d77604194102.png"
-#     ],
-#     [
-#         "example/image/87cdf3e2-0fa1-4d80-a228-cbb4aba3f44f.png",
-#         "example/mask/87cdf3e2-0fa1-4d80-a228-cbb4aba3f44f.png"
-#     ],
-#     [
-#         "example/image/55dd199b-d99b-47a2-a691-edfd92233a6b.png",
-#         "example/mask/55dd199b-d99b-47a2-a691-edfd92233a6b.png"
-#     ]
+
+image_examples = [
+    [
+        "example/image/3c43156c-2b44-4ebf-9c47-7707ec60b166.png",
+        "example/mask/3c43156c-2b44-4ebf-9c47-7707ec60b166.png"
+    ],
+    [
+        "example/image/0e5124d8-fe43-4b5c-819f-7212f23a6d2a.png",
+        "example/mask/0e5124d8-fe43-4b5c-819f-7212f23a6d2a.png"
+    ],
+    [
+        "example/image/0f900fe8-6eab-4f85-8121-29cac9509b94.png",
+        "example/mask/0f900fe8-6eab-4f85-8121-29cac9509b94.png"
+    ],
+    [
+        "example/image/3ed1ee18-33b0-4964-b679-0e214a0d8848.png",
+        "example/mask/3ed1ee18-33b0-4964-b679-0e214a0d8848.png"
+    ],
+    [
+        "example/image/9a3b6af9-c733-46a4-88d4-d77604194102.png",
+        "example/mask/9a3b6af9-c733-46a4-88d4-d77604194102.png"
+    ],
+    [
+        "example/image/87cdf3e2-0fa1-4d80-a228-cbb4aba3f44f.png",
+        "example/mask/87cdf3e2-0fa1-4d80-a228-cbb4aba3f44f.png"
+    ],
+    [
+        "example/image/55dd199b-d99b-47a2-a691-edfd92233a6b.png",
+        "example/mask/55dd199b-d99b-47a2-a691-edfd92233a6b.png"
+    ]
     
-# ]
+]
+
 
 base_model_path = 'black-forest-labs/FLUX.1-dev'
 lora_path = 'theSure/Omnieraser'  
 transformer = FluxTransformer2DModel.from_pretrained(base_model_path, subfolder='transformer', torch_dtype=torch.bfloat16)
 gr.Info(str(f"Model loading: {int((40 / 100) * 100)}%"))
-# enable image inputs
 with torch.no_grad():
     initial_input_channels = transformer.config.in_channels
     new_linear = torch.nn.Linear(
@@ -90,64 +90,48 @@ def set_seed(seed):
 @spaces.GPU 
 def predict(
     input_image,
+    uploaded_mask
     prompt,
     ddim_steps,
     seed,
     scale,
-    image_paths,
-    mask_paths
 
 ):
-    global image_path, mask_path
-    gr.Info(str(f"Set seed = {seed}"))
-    if image_paths is not None:
-        input_image["background"] = load_image(image_paths).convert("RGB")
-        input_image["layers"][0] = load_image(mask_paths).convert("RGB")
-        
-    size1, size2 = input_image["background"].convert("RGB").size
-    icc_profile = input_image["background"].info.get('icc_profile')
+    gr.Info(str(f"Set seed = {seed}"))  
+    size1, size2 = input_image.convert("RGB").size
+    icc_profile = input_image.info.get('icc_profile')
     if icc_profile:
         gr.Info(str(f"Image detected to contain ICC profile, converting color space to sRGB..."))
         srgb_profile = ImageCms.createProfile("sRGB")
         io_handle = io.BytesIO(icc_profile)   
         src_profile = ImageCms.ImageCmsProfile(io_handle)  
-        input_image["background"] = ImageCms.profileToProfile(input_image["background"], src_profile, srgb_profile)
-        input_image["background"].info.pop('icc_profile', None)
+        input_image = ImageCms.profileToProfile(input_image["background"], src_profile, srgb_profile)
+        input_image.info.pop('icc_profile', None)
 
     if size1 < size2:
-        input_image["background"] = input_image["background"].convert("RGB").resize((1024, int(size2 / size1 * 1024)))
+        input_image = input_image.convert("RGB").resize((1024, int(size2 / size1 * 1024)))
     else:
-        input_image["background"] = input_image["background"].convert("RGB").resize((int(size1 / size2 * 1024), 1024))
+        input_image = input_image.convert("RGB").resize((int(size1 / size2 * 1024), 1024))
 
-    img = np.array(input_image["background"].convert("RGB"))
+    img = np.array(input_image.convert("RGB"))
 
     W = int(np.shape(img)[0] - np.shape(img)[0] % 8)
     H = int(np.shape(img)[1] - np.shape(img)[1] % 8)
 
-    input_image["background"] = input_image["background"].resize((H, W))
-    input_image["layers"][0] = input_image["layers"][0].resize((H, W))
+    input_image = input_imageresize((H, W))
+    uploaded_mask = uploaded_mask.resize((H, W))
 
     if seed == -1:
         seed = random.randint(1, 2147483647)
         set_seed(random.randint(1, 2147483647))
     else:
         set_seed(seed)
-    if image_paths is None:
-        img=input_image["layers"][0]
-        img_data = np.array(img)
-        alpha_channel = img_data[:, :, 3]
-        white_background = np.ones_like(alpha_channel) * 255
-        gray_image = white_background.copy()
-        gray_image[alpha_channel == 0] = 0
-        gray_image_pil = Image.fromarray(gray_image).convert('L')
-    else:
-        gray_image_pil = input_image["layers"][0]
     base_model_path = 'black-forest-labs/FLUX.1-dev'
     lora_path = 'theSure/Omnieraser'
     result = pipe(
         prompt=prompt,
-        control_image=input_image["background"].convert("RGB"),
-        control_mask=gray_image_pil.convert("RGB"),
+        control_image=input_image.convert("RGB"),
+        control_mask=uploaded_mask.convert("RGB"),
         width=H,
         height=W,
         num_inference_steps=ddim_steps,
@@ -156,19 +140,19 @@ def predict(
         max_sequence_length=512,
     ).images[0]
 
-    mask_np = np.array(input_image["layers"][0].convert("RGB"))
-    red = np.array(input_image["background"]).astype("float") * 1
+    mask_np = np.array(uploaded_mask.convert("RGB"))
+    red = np.array(input_image).astype("float") * 1
     red[:, :, 0] = 180.0
     red[:, :, 2] = 0
     red[:, :, 1] = 0
-    result_m = np.array(input_image["background"])
+    result_m = np.array(input_image)
     result_m = Image.fromarray(
         (
             result_m.astype("float") * (1 - mask_np.astype("float") / 512.0) + mask_np.astype("float") / 512.0 * red
         ).astype("uint8")
     )
 
-    dict_res = [input_image["background"], input_image["layers"][0], result_m, result]
+    dict_res = [input_image, uploaded_mask, result_m, result]
 
     dict_out = [result]
     image_path = None
@@ -178,21 +162,19 @@ def predict(
 
 def infer(
     input_image,
+    uploaded_mask，
     ddim_steps,
     seed,
     scale,
     removal_prompt,
 
 ):
-    img_path = image_path
-    msk_path = mask_path
     return predict(input_image, 
+                   uploaded_mask
                    removal_prompt, 
                    ddim_steps, 
                    seed,
                    scale,
-                   img_path,
-                   msk_path
     )
 
 def process_example(image_paths, mask_paths):
@@ -288,13 +270,8 @@ with gr.Blocks(
             gr.Markdown("## 📥 Input Panel")
             
             with gr.Group():
-                input_image = gr.Sketchpad(
-                    sources=["upload"],
-                    type="pil",
-                    label="Upload & Annotate",
-                    elem_id="custom-image",
-                    interactive=True
-                )
+                image_input = gr.Image(label="Upload Image", type="pil", image_mode="RGB")
+                uploaded_mask = gr.Image(label="Upload Mask", type="pil", image_mode="L") 
             with gr.Row(variant="compact"):
                 run_button = gr.Button(
                     "🚀 Start Processing",
@@ -311,21 +288,18 @@ with gr.Blocks(
                     step=1,
                     info="-1 for random generation"
                 )
-            # with gr.Column(variant="panel"):
-            #     gr.Markdown("### 🖼️ Example Gallery", elem_classes=["example-title"])
-            #     example = gr.Examples(
-            #         examples=image_examples,
-            #         inputs=[
-            #             gr.Image(label="Image", type="filepath",visible=False),
-            #             gr.Image(label="Mask", type="filepath",visible=False)
-            #         ],
-            #         outputs=[input_image],
-            #         fn=process_example,
-            #         run_on_click=True,
-            #         examples_per_page=10,
-            #         label="Click any example to load",
-            #         elem_id="inline-examples"
-            #     )
+            with gr.Column(variant="panel"):
+                gr.Markdown("### 🖼️ Example Gallery", elem_classes=["example-title"])
+                example = gr.Examples(
+                    examples=image_examples,
+                    inputs=[
+                        image_input, uploaded_mask
+                    ],
+                    outputs=[inpaint_result, gallery],
+                    examples_per_page=10,
+                    label="Click any example to load",
+                    elem_id="inline-examples"
+                )
 
         with gr.Column(scale=1, variant="panel"):
             gr.Markdown("## 📤 Output Panel")
@@ -351,6 +325,7 @@ with gr.Blocks(
         fn=infer,
         inputs=[
             input_image,
+            uploaded_mask
             ddim_steps,
             seed,
             scale,