Update app.py

app.py

@@ -149,6 +149,31 @@ 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}")
 
+@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."""
+    if image is None:
+        raise gr.Error("No image uploaded. Please upload an image to upscale.")
+    
+    try:
+        original = load_image(image)
+        all_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)
+            all_results.append((upscaled, model_name))
+        
+        progress(1.0, desc="All upscaling complete!")
+        return all_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
 
@@ -172,6 +197,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")
             
         with gr.Column(scale=2):
             result_slider = ImageSlider(
@@ -188,6 +214,17 @@ 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,
+                elem_id="gallery",
+                columns=2,
+                rows=2,
+                object_fit="scale-down",
+                show_download_button=True
+            )
 
     # --- Event Handling ---
     run_button.click(
@@ -200,6 +237,12 @@ with gr.Blocks(delete_cache=(3600, 3600)) as demo:
         inputs=[input_image, model_selection],
         outputs=[result_slider, download_output],
     )
+    
+    run_all_button.click(
+        fn=upscale_with_all_models,
+        inputs=[input_image],
+        outputs=[all_models_gallery],
+    )
 
 # --- Pre-load the default model for a faster first-time user experience ---
 try: