Update app.py

app.py

@@ -248,11 +248,11 @@ def generate_single_image(label_str):
     img_pil = Image.fromarray((img_np * 255).astype(np.uint8))
     
     return img_pil
-
-def generate_batch_images(label_str, num_images, progress=gr.Progress()):
+def generate_images(label_str, num_images, progress=gr.Progress()):
     global loaded_model
     cancel_event.clear()
     
+    # Input validation
     if num_images < 1 or num_images > 10:
         raise gr.Error("Number of images must be between 1 and 10")
     
@@ -282,8 +282,9 @@ def generate_batch_images(label_str, num_images, progress=gr.Progress()):
             )
         
         if images is None:
-            return None
+            return None, None
             
+        # Process all generated images
         processed_images = []
         for img in images:
             img_np = img.cpu().permute(1, 2, 0).numpy()
@@ -291,7 +292,11 @@ def generate_batch_images(label_str, num_images, progress=gr.Progress()):
             pil_img = Image.fromarray((img_np * 255).astype(np.uint8))
             processed_images.append(pil_img)
         
-        return processed_images
+        # Return both single image and gallery based on count
+        if num_images == 1:
+            return processed_images[0], processed_images
+        else:
+            return None, processed_images
 
     except torch.cuda.OutOfMemoryError:
         torch.cuda.empty_cache()
@@ -300,7 +305,7 @@ def generate_batch_images(label_str, num_images, progress=gr.Progress()):
         traceback.print_exc()
         if str(e) != "Generation was cancelled by user":
             raise gr.Error(f"Generation failed: {str(e)}")
-        return None
+        return None, None
     finally:
         torch.cuda.empty_cache()
 
@@ -312,7 +317,7 @@ print("Loading model...")
 loaded_model = load_model(model_path, device)
 print("Model loaded successfully!")
 
-# --- Gradio UI (from first file with modifications) ---
+# Unified Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(
     primary_hue="violet",
     neutral_hue="slate",
@@ -351,32 +356,44 @@ with gr.Blocks(theme=gr.themes.Soft(
             """)
         
         with gr.Column(scale=2):
-            with gr.Tab("Single Image"):
-                single_image = gr.Image(
-                    type="pil",
-                    label="Generated X-ray",
-                    height=400
-                )
-            with gr.Tab("Batch Images"):
-                gallery = gr.Gallery(
-                    label="Generated X-rays",
-                    columns=3,
-                    height="auto",
-                    object_fit="contain"
-                )
+            # Unified output display that adapts to single/batch
+            with gr.Tabs():
+                with gr.TabItem("Output", id="output_tab"):
+                    single_image = gr.Image(
+                        label="Generated X-ray",
+                        height=400,
+                        visible=True
+                    )
+                    gallery = gr.Gallery(
+                        label="Generated X-rays",
+                        columns=3,
+                        height="auto",
+                        object_fit="contain",
+                        visible=False
+                    )
+    
+    def update_ui_based_on_count(num_images):
+        if num_images == 1:
+            return {
+                single_image: gr.update(visible=True),
+                gallery: gr.update(visible=False)
+            }
+        else:
+            return {
+                single_image: gr.update(visible=False),
+                gallery: gr.update(visible=True)
+            }
     
-    # Single image generation
-    condition.change(
-        fn=generate_single_image,
-        inputs=condition,
-        outputs=single_image
+    num_images.change(
+        fn=update_ui_based_on_count,
+        inputs=num_images,
+        outputs=[single_image, gallery]
     )
     
-    # Batch image generation
     submit_btn.click(
-        fn=generate_batch_images,
+        fn=generate_images,
         inputs=[condition, num_images],
-        outputs=gallery
+        outputs=[single_image, gallery]
     )
     
     cancel_btn.click(