Update app.py

app.py

@@ -752,6 +752,13 @@ class ImageGenerator:
         return result.images[0]
 
 
+def pack_image(filename):
+    if filename is None:
+        return Image.new("P",size=(512, 512))
+    print("这不是none.")
+    image = Image.open(filename)
+    return image
+
 # 实例化生成器
 generator = ImageGenerator()
 
@@ -771,28 +778,53 @@ examples = [
 
 ]
 
+# examples = [
+
+# ["", pack_image("assets/mask2face/handou_seg.png"), None, "FaRL", 42, 28],
+
+# ["", pack_image("assets/mask2face/black_seg.png"), None, "FaRL", 42, 28],
+
+# ["She has red hair", pack_image("assets/multimodal/liuyifei_seg.png"), None, "FaRL", 42, 28],
+
+# ["He is old", pack_image("assets/multimodal/musk_seg.png"), None, "FaRL", 42, 28],
+
+# ["Curly-haired woman with glasses", pack_image(None), None, "FaRL", 42, 28],
+
+# ["Man with beard and tie", pack_image(None), None, "FaRL", 42, 28],
+
+# ]
+
+
 # Gradio 界面（使用 Blocks）
 with gr.Blocks(title="Controllable Face Generation with MoGLE") as demo:
-    gr.Markdown("## 🎭 Controllable Face Generation via Prompt + Face Parsing")
+    gr.Markdown("<center><h1>Face-MoGLE: Mixture of Global and Local Experts with Diffusion Transformer for Controllable Face Generation</h1></center>")
 
     with gr.Row():
         prompt = gr.Textbox(label="Text Prompt", placeholder="Describe the face you'd like to generate...")
 
     with gr.Row():
         with gr.Column():
-            mask_image = gr.Image(type="pil", label="🧩 Segmantic Mask (Optional)")
-            rgb_image = gr.Image(type="pil", label="🖼️ Facial Image (Optional)")
+            
+            mask_image = gr.Image(source="upload", type="pil", label="Semantic Mask (Optional)", tool="color-sketch", image_mode="P", interactive=True, height=512, width=512)
+
+            rgb_image = gr.Image(type="pil", label="Facial Image (Optional)")
             model_choice = gr.Radio(["FaRL", "SegFace"], label="Face Parsing Model", value="FaRL")
             seed = gr.Slider(minimum=0, maximum=100000, step=1, value=42, label="Random Seed")
             num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=28, label="Sampling Step")
             submit_btn = gr.Button("Generate")
 
         with gr.Column():
-            gr.Markdown("### 🧠 Parsed Mask Preview")
-            preview_mask = gr.Image(label="Parsed Mask (from RGB)", interactive=False)
-            output_image = gr.Image(label="🎨 Generated Image")
+            preview_mask = gr.Image(label="Parsed Semantic Mask (From the Facial Image)", interactive=False)
+            output_image = gr.Image(label="Generated Image")
 
     def generate_wrapper(prompt, mask_image, rgb_image, model_choice, seed,num_inference_steps):
+
+        if mask_image is not None:
+            if isinstance(mask_image, Image.Image):
+                mask_image = mask_image.resize((512, 512), Image.BICUBIC)
+            if isinstance(mask_image, str):
+                mask_image = Image.open(mask_image).resize((512, 512), Image.BICUBIC)
+
         if mask_image is None and rgb_image is not None:
             if model_choice == "FaRL":
                 mask_image = generator.parse_face_with_farl(rgb_image)
@@ -809,13 +841,14 @@ with gr.Blocks(title="Controllable Face Generation with MoGLE") as demo:
         outputs=[preview_mask, output_image]
     )
     gr.Examples(
-        examples=examples,
-        inputs=[prompt, mask_image, rgb_image, model_choice, seed, num_inference_steps],
-        outputs=[preview_mask, output_image],
-        fn=lambda *args: generate_wrapper(*args),  # 直接引用已定义的函数
-        cache_examples=False,
-        label="Click any example below to try:"
-    )
+    examples=examples,
+    inputs=[prompt, mask_image, rgb_image, model_choice, seed, num_inference_steps],
+    outputs=[preview_mask, output_image],
+    fn=generate_wrapper,
+    cache_examples=False,
+    label="Click any example below to try:"
+ )
+
 
 if __name__ == "__main__":
     demo.launch()