Update app.py

app.py

@@ -149,49 +149,46 @@ def upscale_image(image, model_selection: str, progress=gr.Progress(track_tqdm=T
         print(f"An error occurred: {traceback.format_exc()}")
         raise gr.Error(f"An error occurred during processing: {e}")
 
+def clear_outputs():
+    return None, None
+
+# --- Batch Upscaling Function ---
 @spaces.GPU
 def upscale_with_all_models(image, progress=gr.Progress(track_tqdm=True)):
-    """Upscale image with all available models and return results as HTML."""
     if image is None:
         raise gr.Error("No image uploaded. Please upload an image to upscale.")
-    
+
     try:
+        progress(0, desc="Loading image...")
         original = load_image(image)
-        all_results = []
+        
+        results = []
         total_models = len(MODELS)
         
-        for idx, (model_name, _) in enumerate(MODELS.items()):
-            progress(idx / total_models, desc=f"Upscaling with {model_name}... ({idx + 1}/{total_models})")
-            print(f"Processing with model: {model_name}")
-            upscaler = get_upscaler(model_name)
-            upscaled = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)
+        for idx, (model_name, model_path) in enumerate(MODELS.items()):
+            progress_val = idx / total_models
+            progress(progress_val, desc=f"Processing with {model_name} ({idx+1}/{total_models})...")
             
-            # Save to temp file
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
-                upscaled.save(temp_file.name, "PNG")
-                all_results.append((temp_file.name, model_name))
+            try:
+                upscaler = get_upscaler(model_name)
+                upscaled_pil_image = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)
+                
+                # Save result with model name label
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
+                    upscaled_pil_image.save(temp_file.name, "PNG")
+                    results.append((temp_file.name, model_name))
+                    
+            except Exception as e:
+                print(f"Error processing {model_name}: {e}")
+                # Continue with next model even if one fails
+                continue
         
-        progress(1.0, desc="All upscaling complete!")
-        
-        # Build HTML output with images
-        html = '<div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 20px;">'
-        for img_path, model_name in all_results:
-            html += f'''
-            <div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px;">
-                <img src="file={img_path}" style="width: 100%; height: auto; border-radius: 4px;">
-                <p style="text-align: center; margin-top: 8px; font-weight: bold;">{model_name}</p>
-            </div>
-            '''
-        html += '</div>'
-        
-        return html
-    
+        progress(1.0, desc="All models processed!")
+        return results
+
     except Exception as e:
         print(f"An error occurred: {traceback.format_exc()}")
-        raise gr.Error(f"An error occurred during processing: {e}")
-
-def clear_outputs():
-    return None, None
+        raise gr.Error(f"An error occurred during batch processing: {e}")
 
 # --- Gradio Interface Definition ---
 title = """<h1 align="center">Image Upscaler</h1>
@@ -213,7 +210,7 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
                 label="Model (alphabetically sorted)",
             )
             run_button = gr.Button("Upscale", variant="primary")
-            run_all_button = gr.Button("Upscale with All Models", variant="secondary")
+            batch_button = gr.Button("Upscale with All Models (takes time!)", variant="secondary")
             
         with gr.Column(scale=2):
             result_slider = ImageSlider(
@@ -230,9 +227,19 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
             )
             
             download_output = gr.File(label="Download Full-Quality Upscaled Image (Lossless PNG)")
-            
-            # --- All models output ---
-            all_models_output = gr.HTML(label="All Models Results", visible=True)
+    
+    # --- Gallery for batch results ---
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Batch Processing Results")
+            batch_gallery = gr.Gallery(
+                label="All Models Results",
+                show_label=True,
+                columns=4,
+                rows=None,
+                height="auto",
+                object_fit="contain"
+            )
 
     # --- Event Handling ---
     run_button.click(
@@ -246,10 +253,15 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
         outputs=[result_slider, download_output],
     )
     
-    run_all_button.click(
+    batch_button.click(
+        fn=lambda: None,
+        inputs=None,
+        outputs=batch_gallery,
+        queue=False
+    ).then(
         fn=upscale_with_all_models,
         inputs=[input_image],
-        outputs=[all_models_output],
+        outputs=batch_gallery,
     )
 
 # --- Pre-load the default model for a faster first-time user experience ---