Update app.py

app.py

@@ -291,12 +291,46 @@ def get_random_seed(randomize_seed, seed):
         logger.error(f"Error in get_random_seed: {str(e)}")
         raise
 
+
+def download_image(url: str, save_path: str) -> str:
+    """Download an image from a URL and save it locally."""
+    try:
+        logger.info(f"Downloading image from {url}")
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        with open(save_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        logger.info(f"Saved image to {save_path}")
+        return save_path
+    except Exception as e:
+        logger.error(f"Failed to download image from {url}: {str(e)}")
+        raise
+
 @spaces.GPU()
 @torch.no_grad()
-def run_segmentation(image: str):
+def run_segmentation(image):
     try:
         logger.info("Running segmentation")
-        image = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
+        # Handle FileData dict or URL
+        if isinstance(image, dict):
+            image_path = image.get("path") or image.get("url")
+            if not image_path:
+                logger.error("Invalid image input: no path or URL provided")
+                raise ValueError("Invalid image input: no path or URL provided")
+            if image_path.startswith("http"):
+                temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
+                image_path = download_image(image_path, temp_image_path)
+        elif isinstance(image, str) and image.startswith("http"):
+            temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
+            image_path = download_image(image, temp_image_path)
+        else:
+            image_path = image
+            if not isinstance(image, (str, bytes)) or (isinstance(image, str) and not os.path.exists(image)):
+                logger.error(f"Invalid image type or path: {type(image)}")
+                raise ValueError(f"Expected str (path/URL), bytes, or FileData dict, got {type(image)}")
+
+        image = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
         logger.info("Segmentation complete")
         return image
     except Exception as e: