Update app.py

app.py

@@ -532,20 +532,18 @@ def process_input(input_image, upscale_factor, **kwargs):
 
     was_resized = False
 
-    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-        warnings.warn(
-            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels."
-        )
-        gr.Info(
-            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing input to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels budget."
-        )
-        input_image = input_image.resize(
-            (
-                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
-                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
-            )
-        )
+    max_size = int(np.sqrt(MAX_PIXEL_BUDGET / (upscale_factor ** 2)))
+    if w > max_size or h > max_size:
+        if w > h:
+            w_new = max_size
+            h_new = int(w_new / aspect_ratio)
+        else:
+            h_new = max_size
+            w_new = int(h_new * aspect_ratio)
+        
+        input_image = input_image.resize((w_new, h_new), Image.LANCZOS)
         was_resized = True
+        gr.Info(f"Input image resized to {w_new}x{h_new} to fit within pixel budget after upscaling.")
 
     # resize to multiple of 8
     w, h = input_image.size