update code

app.py

@@ -9,6 +9,7 @@ import requests
 from io import BytesIO
 from kan_linear import KANLinear
 import logging
+import os
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -46,6 +47,13 @@ class CustomImageLoadingError(Exception):
 def load_image_from_url(url):
     try:
         logging.info(f"Loading image from URL: {url}")
+
+        # Check the file extension
+        valid_extensions = ['jpg', 'jpeg', 'png', 'webp']
+        file_extension = os.path.splitext(url)[1][1:].lower()
+        if file_extension not in valid_extensions:
+            raise CustomImageLoadingError(f"URL does not point to an image with a valid extension: {file_extension}")
+
         response = requests.get(url)
         response.raise_for_status()  # Check if the request was successful