Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
 import gradio as gr
 import spaces
+import os
 
 def resize_image(image, max_size=2048):
     width, height = image.size
@@ -44,11 +45,11 @@ def upscale_chunk(chunk, model, processor, device):
     return Image.fromarray(output_image)
 
 @spaces.GPU
-def main(image, model_choice, save_as_jpg=True, use_tiling=True, auto_cpu=True):
+def main(image, model_choice, save_as_jpg=True, use_tiling=True):
     # Resize the input image
     image = resize_image(image)
     
-    device = torch.device("cuda" if torch.cuda.is_available() and not auto_cpu else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     model_paths = {
         "Pixel Perfect": "caidas/swin2SR-classical-sr-x4-64",
@@ -77,16 +78,22 @@ def main(image, model_choice, save_as_jpg=True, use_tiling=True, auto_cpu=True):
         # Process the entire image at once
         upscaled_image = upscale_chunk(image, model, processor, device)
     
+    # Generate output filename
+    original_filename = os.path.splitext(image.filename)[0] if image.filename else "image"
+    output_filename = f"{original_filename}_upscaled"
+    
     if save_as_jpg:
-        upscaled_image.save("upscaled_image.jpg", quality=95)
-        return "upscaled_image.jpg"
+        output_filename += ".jpg"
+        upscaled_image.save(output_filename, quality=95)
     else:
-        upscaled_image.save("upscaled_image.png")
-        return "upscaled_image.png"
+        output_filename += ".png"
+        upscaled_image.save(output_filename)
+    
+    return output_filename
 
-def gradio_interface(image, model_choice, save_as_jpg, use_tiling, auto_cpu):
+def gradio_interface(image, model_choice, save_as_jpg, use_tiling):
     try:
-        result = main(image, model_choice, save_as_jpg, use_tiling, auto_cpu)
+        result = main(image, model_choice, save_as_jpg, use_tiling)
         return result, None
     except Exception as e:
         return None, str(e)
@@ -102,14 +109,13 @@ interface = gr.Interface(
         ),
         gr.Checkbox(value=True, label="Save as JPEG"),
         gr.Checkbox(value=True, label="Use Tiling"),
-        gr.Checkbox(value=True, label="Auto CPU"),
     ],
     outputs=[
         gr.File(label="Download Upscaled Image"),
         gr.Textbox(label="Error Message", visible=True)
     ],
     title="Image Upscaler",
-    description="Upload an image, select a model, and upscale it. Images larger than 2048x2048 will be resized while maintaining aspect ratio. Use tiling for efficient processing of large images. Auto CPU will use CPU if GPU is not available.",
+    description="Upload an image, select a model, and upscale it. Images larger than 2048x2048 will be resized while maintaining aspect ratio. Use tiling for efficient processing of large images.",
 )
 
 interface.launch()