Update app.py

app.py

@@ -151,7 +151,7 @@ def upscale_image(image, model_selection: str, progress=gr.Progress(track_tqdm=T
 
 @spaces.GPU
 def upscale_with_all_models(image, progress=gr.Progress(track_tqdm=True)):
-    """Upscale image with all available models and return results as a list."""
+    """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.")
     
@@ -165,10 +165,26 @@ def upscale_with_all_models(image, progress=gr.Progress(track_tqdm=True)):
             print(f"Processing with model: {model_name}")
             upscaler = get_upscaler(model_name)
             upscaled = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)
-            all_results.append((upscaled, model_name))
+            
+            # 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))
         
         progress(1.0, desc="All upscaling complete!")
-        return all_results
+        
+        # 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
     
     except Exception as e:
         print(f"An error occurred: {traceback.format_exc()}")
@@ -215,15 +231,8 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
             
             download_output = gr.File(label="Download Full-Quality Upscaled Image (Lossless PNG)")
             
-            # --- Gallery for all models ---
-            all_models_gallery = gr.Gallery(
-                label="All Models Results",
-                show_label=True,
-                columns=2,
-                rows=1,
-                object_fit="scale-down",
-                show_download_button=True
-            )
+            # --- All models output ---
+            all_models_output = gr.HTML(label="All Models Results", visible=True)
 
     # --- Event Handling ---
     run_button.click(
@@ -240,11 +249,7 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
     run_all_button.click(
         fn=upscale_with_all_models,
         inputs=[input_image],
-        outputs=[all_models_gallery],
-    ).then(
-        lambda x: x,
-        inputs=[all_models_gallery],
-        outputs=[all_models_gallery]
+        outputs=[all_models_output],
     )
 
 # --- Pre-load the default model for a faster first-time user experience ---