Update app.py

app.py

@@ -86,7 +86,7 @@ QUALITY_PROMPT = ", high quality, detailed, vibrant, professional lighting"
 # --- Main Inference Function ---
 @spaces.GPU(duration=40)
 def infer(
-    images,
+    image,
     prompt,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -107,19 +107,19 @@ def infer(
     # Set up the generator for reproducibility
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Load input images into PIL Images
-    pil_images = []
-    if images is not None:
-        for item in images:
-            try:
-                if isinstance(item[0], Image.Image):
-                    pil_images.append(item[0].convert("RGB"))
-                elif isinstance(item[0], str):
-                    pil_images.append(Image.open(item[0]).convert("RGB"))
-                elif hasattr(item, "name"):
-                    pil_images.append(Image.open(item.name).convert("RGB"))
-            except Exception:
-                continue
+    # Load input image into PIL Image
+    pil_image = None
+    if image is not None:
+        try:
+            if isinstance(image, Image.Image):
+                pil_image = image.convert("RGB")
+            elif isinstance(image, str):
+                pil_image = Image.open(image).convert("RGB")
+            elif hasattr(image, "name"):
+                pil_image = Image.open(image.name).convert("RGB")
+        except Exception as e:
+            print(f"Error loading image: {e}")
+            raise gr.Error("Failed to load input image.")
 
     if height == 256 and width == 256:
         height, width = None, None
@@ -128,8 +128,8 @@ def infer(
     print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
 
     # Generate the image
-    image = pipe(
-        image=pil_images if len(pil_images) > 0 else None,
+    output = pipe(
+        image=[pil_image] if pil_image is not None else None,
         prompt=prompt_final,
         height=height,
         width=width,
@@ -140,7 +140,7 @@ def infer(
         num_images_per_prompt=num_images_per_prompt,
     ).images
 
-    return image
+    return output[0] if output else None
 
 # --- Gradio User Interface ---
 def create_demo():
@@ -168,7 +168,6 @@ def create_demo():
         }
         #general_items {
             width: 100%;
-            max-width: 100%;
             margin: 2rem 0;
             display: flex;
             flex-direction: column;
@@ -181,7 +180,6 @@ def create_demo():
             padding: 1rem;
             box-shadow: 0 0 8px rgba(255, 255, 255, 0.2);
             width: 100%;
-            max-width: 100%;
             box-sizing: border-box;
         }
         h1 {
@@ -204,16 +202,16 @@ def create_demo():
             border: none;
             margin: 0.75rem 0;
             width: 100%;
-            max-width: 100%;
         }
-        .gr-gallery {
+        .image-container {
+            aspect-ratio: 1/1;
             width: 100%;
-            max-width: 100%;
+            height: auto;
             border: 1px solid #FFFFFF;
             border-radius: 4px;
             box-sizing: border-box;
         }
-        .gr-gallery img {
+        .image-container img {
             width: 100%;
             height: auto;
             box-sizing: border-box;
@@ -225,7 +223,6 @@ def create_demo():
             border-radius: 4px;
             padding: 0.5rem;
             width: 100%;
-            max-width: 100%;
             box-sizing: border-box;
         }
         input:hover, textarea:hover {
@@ -243,7 +240,6 @@ def create_demo():
             box-shadow: 0 0 8px rgba(255, 255, 255, 0.3);
             transition: box-shadow 0.3s, transform 0.3s;
             width: 100%;
-            max-width: 100%;
             min-height: 48px;
             cursor: pointer;
         }
@@ -251,21 +247,7 @@ def create_demo():
             box-shadow: 0 0 12px rgba(255, 255, 255, 0.5);
             transform: scale(1.05);
         }
-        button[aria-label="Download"] {
-            transform: scale(3);
-            transform-origin: top right;
-            background: #000000 !important;
-            color: #FFFFFF !important;
-            border: 1px solid #FFFFFF !important;
-            border-radius: 4px;
-            padding: 0.4rem !important;
-            margin: 0.5rem !important;
-            box-shadow: 0 0 8px rgba(255, 255, 255, 0.3);
-            transition: box-shadow 0.3s;
-        }
-        button[aria-label="Download"]:hover {
-            box-shadow: 0 0 12px rgba(255, 255, 255, 0.5);
-        }
+        button[aria-label="Download"], button[aria-label="Download"]:hover,
         button[aria-label="Fullscreen"], button[aria-label="Fullscreen"]:hover,
         button[aria-label="Share"], button[aria-label="Share"]:hover {
             display: none !important;
@@ -295,18 +277,21 @@ def create_demo():
             gr.Markdown("# Qwen Image Editor")
             gr.Markdown("Edit your images with precise instructions", elem_id="subtitle")
             with gr.Column(elem_id="input_column"):
-                input_images = gr.Gallery(
-                    label="Input Images",
-                    show_label=True,
+                input_image = gr.Image(
+                    label="Input Image",
                     type="pil",
                     interactive=True,
-                    elem_classes=["gradio-component", "gr-gallery"]
+                    show_download_button=False,
+                    show_share_button=False,
+                    elem_classes=["gradio-component", "image-container"]
                 )
-                result = gr.Gallery(
-                    label="Result",
-                    show_label=True,
+                result_image = gr.Image(
+                    label="Result Image",
                     type="pil",
-                    elem_classes=["gradio-component", "gr-gallery"]
+                    interactive=False,
+                    show_download_button=False,
+                    show_share_button=False,
+                    elem_classes=["gradio-component", "image-container"]
                 )
                 prompt = gr.Textbox(
                     label="Prompt",
@@ -323,8 +308,8 @@ def create_demo():
         gr.on(
             triggers=[run_button.click, prompt.submit],
             fn=infer,
-            inputs=[input_images, prompt],
-            outputs=[result],
+            inputs=[input_image, prompt],
+            outputs=[result_image],
         )
 
     return demo