Update app.py

app.py

@@ -171,34 +171,35 @@ def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Im
     return mask, bbox, time_log
 
 def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str | None = None) -> tuple[tuple[Image.Image, Image.Image, Image.Image], gr.DownloadButton]:
-    if img.width > 2048 or img.height > 2048:
-        orig_res = max(img.width, img.height)
-        img.thumbnail((2048, 2048))
-        if isinstance(prompt, tuple):
-            x0, y0, x1, y1 = (int(x * 2048 / orig_res) for x in prompt)
-            prompt = (x0, y0, x1, y1)
+    try:
+        if img.width > 2048 or img.height > 2048:
+            orig_res = max(img.width, img.height)
+            img.thumbnail((2048, 2048))
+            if isinstance(prompt, tuple):
+                x0, y0, x1, y1 = (int(x * 2048 / orig_res) for x in prompt)
+                prompt = (x0, y0, x1, y1)
 
-    mask, bbox, time_log = _gpu_process(img, prompt)
-    masked_alpha = apply_mask(img, mask, defringe=True)
+        mask, bbox, time_log = _gpu_process(img, prompt)
+        masked_alpha = apply_mask(img, mask, defringe=True)
 
-    if bg_prompt:
-        try:
+        if bg_prompt:
             background = generate_background(bg_prompt, img.width, img.height)
             combined = combine_with_background(masked_alpha, background)
-        except Exception as e:
-            raise gr.Error(f"Background processing failed: {str(e)}")
-    else:
-        combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
+        else:
+            combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
 
-    thresholded = mask.point(lambda p: 255 if p > 10 else 0)
-    bbox = thresholded.getbbox()
-    to_dl = masked_alpha.crop(bbox)
+        thresholded = mask.point(lambda p: 255 if p > 10 else 0)
+        bbox = thresholded.getbbox()
+        to_dl = masked_alpha.crop(bbox)
 
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-    to_dl.save(temp, format="PNG")
-    temp.close()
+        temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        to_dl.save(temp, format="PNG")
+        temp.close()
 
-    return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)
+        return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)
+        
+    except Exception as e:
+        raise gr.Error(f"Processing failed: {str(e)}")
 
 def process_bbox(img: Image.Image, box_input: str) -> tuple[list[Image.Image], str]:
     try: