Update app.py

app.py

@@ -7,55 +7,52 @@ import random
 import spaces
 import gradio as gr
 import numpy as np
+
 from PIL import Image, ImageCms
 import torch
 from diffusers import FluxTransformer2DModel
 from diffusers.utils import load_image
 from pipeline_flux_control_removal import FluxControlRemovalPipeline
-
-# 初始化模型部分
 pipe = None
 torch.set_grad_enabled(False)
-
-# 示例数据
-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_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"
+#     ]
     
-]
+# ]
 
-# 模型加载代码（保持不变）
 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(
@@ -71,7 +68,6 @@ with torch.no_grad():
         new_linear.bias.copy_(transformer.x_embedder.bias)
     transformer.x_embedder = new_linear
     transformer.register_to_config(in_channels=initial_input_channels*4)
-
 pipe = FluxControlRemovalPipeline.from_pretrained(
     base_model_path,
     transformer=transformer,
@@ -83,7 +79,6 @@ gr.Info(str(f"Inject LoRA: {lora_path}"))
 pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors")
 gr.Info(str(f"Model loading: {int((100 / 100) * 100)}%"))
 
-# 辅助函数
 @spaces.GPU
 def set_seed(seed):
     torch.manual_seed(seed)
@@ -92,7 +87,6 @@ def set_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
 
-# 主要处理函数
 @spaces.GPU 
 def predict(
     input_image,
@@ -100,14 +94,16 @@ def predict(
     ddim_steps,
     seed,
     scale,
-    image_state,  # 使用State替代全局变量
-    mask_state    # 使用State替代全局变量
+    image_paths,
+    mask_paths
+
 ):
-    if image_state is not None and mask_state is not None:
-        input_image["background"] = load_image(image_state).convert("RGB")
-        input_image["layers"][0] = load_image(mask_state).convert("RGB")
+    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')
     if icc_profile:
@@ -118,17 +114,25 @@ def predict(
         input_image["background"] = ImageCms.profileToProfile(input_image["background"], src_profile, srgb_profile)
         input_image["background"].info.pop('icc_profile', None)
 
-    # ... 保持原有图像尺寸调整逻辑不变
+    if size1 < size2:
+        input_image["background"] = input_image["background"].convert("RGB").resize((1024, int(size2 / size1 * 1024)))
+    else:
+        input_image["background"] = input_image["background"].convert("RGB").resize((int(size1 / size2 * 1024), 1024))
+
+    img = np.array(input_image["background"].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))
 
-    # 保持原有seed处理逻辑
     if seed == -1:
         seed = random.randint(1, 2147483647)
-        set_seed(seed)
+        set_seed(random.randint(1, 2147483647))
     else:
         set_seed(seed)
-
-    # 保持原有mask处理逻辑
-    if image_state is None:
+    if image_paths is None:
         img=input_image["layers"][0]
         img_data = np.array(img)
         alpha_channel = img_data[:, :, 3]
@@ -138,8 +142,8 @@ def predict(
         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"),
@@ -152,7 +156,6 @@ 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
     red[:, :, 0] = 180.0
@@ -166,19 +169,42 @@ def predict(
     )
 
     dict_res = [input_image["background"], input_image["layers"][0], result_m, result]
+
     dict_out = [result]
-    
+    image_path = None
+    mask_path = None
     return dict_out, dict_res
+   
+
+def infer(
+    input_image,
+    ddim_steps,
+    seed,
+    scale,
+    removal_prompt,
+
+):
+    img_path = image_path
+    msk_path = mask_path
+    return predict(input_image, 
+                   removal_prompt, 
+                   ddim_steps, 
+                   seed,
+                   scale,
+                   img_path,
+                   msk_path
+    )
 
-# 示例处理函数
 def process_example(image_paths, mask_paths):
+    global image_path, mask_path
     image = Image.open(image_paths).convert("RGB")
     mask = Image.open(mask_paths).convert("L") 
     black_background = Image.new("RGB", image.size, (0, 0, 0))
     masked_image = Image.composite(black_background, image, mask)
-    return masked_image, image_paths, mask_paths  # 返回路径到State
-
-# 界面布局（保持原有CSS和布局逻辑）
+    
+    image_path = image_paths
+    mask_path = mask_paths
+    return masked_image
 custom_css = """
 
 .contain { max-width: 1200px !important; }
@@ -235,7 +261,6 @@ custom_css = """
 .panel { height: 100%; }
 """
 
-
 with gr.Blocks(
     css=custom_css,
     theme=gr.themes.Soft(
@@ -245,20 +270,23 @@ with gr.Blocks(
     ),
     title="Omnieraser"
 ) as demo:
-    # 添加状态存储
-    image_state = gr.State()
-    mask_state = gr.State()
 
-    # 保持原有组件声明
+
     ddim_steps = gr.Slider(visible=False, value=28)
     scale = gr.Slider(visible=False, value=3.5)
     seed = gr.Slider(visible=False, value=-1)
     removal_prompt = gr.Textbox(visible=False, value="There is nothing here.")
 
-    # 保持原有界面布局
+    gr.Markdown("""
+    <div align="center">
+        <h1 style="font-size: 2.5em; margin-bottom: 0.5em;">🪄 Omnieraser</h1>
+    </div>
+    """)
+
     with gr.Row(equal_height=False):
         with gr.Column(scale=1, variant="panel"):
             gr.Markdown("## 📥 Input Panel")
+            
             with gr.Group():
                 input_image = gr.Sketchpad(
                     sources=["upload"],
@@ -268,7 +296,11 @@ with gr.Blocks(
                     interactive=True
                 )
             with gr.Row(variant="compact"):
-                run_button = gr.Button("🚀 Start Processing", variant="primary", size="lg")
+                run_button = gr.Button(
+                    "🚀 Start Processing",
+                    variant="primary",
+                    size="lg"
+                )
             with gr.Group():
                 gr.Markdown("### ⚙️ Control Parameters")
                 seed = gr.Slider(
@@ -279,21 +311,21 @@ 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, image_state, mask_state],  # 更新状态输出
-                    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=[
+            #             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(scale=1, variant="panel"):
             gr.Markdown("## 📤 Output Panel")
@@ -306,6 +338,7 @@ with gr.Blocks(
                         preview=True,
                         object_fit="contain"
                     )
+
                 with gr.Tab("Visualization Steps"):
                     gallery = gr.Gallery(
                         label="Workflow Steps",
@@ -314,19 +347,16 @@ with gr.Blocks(
                         object_fit="contain"
                     )
 
-    # 更新按钮点击事件
     run_button.click(
-        fn=lambda i, d, s, sc, rp, img, msk: predict(i, rp, d, s, sc, img, msk),
+        fn=infer,
         inputs=[
             input_image,
             ddim_steps,
             seed,
             scale,
             removal_prompt,
-            image_state,  # 添加状态输入
-            mask_state    # 添加状态输入
         ],
         outputs=[inpaint_result, gallery]
     )
-
-demo.launch()
+    
+demo.launch()