Update fetch_url_util.py

fetch_url_util.py

@@ -1,54 +1,73 @@
 import io
 import urllib.parse
+import pillow_jxl
 from PIL import Image
 from curl_cffi import requests
 
+# Max file size: 150MB
+MAX_SIZE = 150 * 1024 * 1024 
+MAX_SIZE_MB = MAX_SIZE / (1024 * 1024)
+
 def fetch_image_from_url(url):
-    """
-    Robustly fetches an image from a URL using TLS impersonation and 
-    dynamic Referer headers. Automatically handles missing http/https prefixes.
-    """
     if not url:
         return None
     
-    # 1. Handle missing protocol (e.g., pixiv.net -> https://pixiv.net)
     url = url.strip()
     if not url.startswith(("http://", "https://")):
         url = "https://" + url
 
-    # 2. Extract base domain for the Referer header
-    # This is crucial for Pixiv and similar sites to prevent 403 Forbidden
     parsed_url = urllib.parse.urlparse(url)
     base_domain = f"{parsed_url.scheme}://{parsed_url.netloc}/"
 
     headers = {
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
-        "Accept": "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.7559.59 Safari/537.36",
+        "Accept": "image/avif,image/webp,image/jxl,image/apng,image/png,image/jpeg,image/gif,image/svg+xml,image/*,*/*;q=0.8",
         "Accept-Language": "en-US,en;q=0.9",
         "Referer": base_domain,
         "Cache-Control": "no-cache",
     }
 
     try:
-        # Changed 'follow_redirects' to 'allow_redirects' for curl_cffi compatibility
-        response = requests.get(
+        # We use stream=True to check file size BEFORE downloading everything
+        with requests.get(
             url, 
             headers=headers, 
-            impersonate="chrome120", 
+            impersonate="chrome120", # Handshake logic (latest stable in curl_cffi)
             timeout=15,
-            allow_redirects=True 
-        )
-        
-        response.raise_for_status()
-        
-        image_bytes = io.BytesIO(response.content)
-        img = Image.open(image_bytes)
-        img.load()
-        
-        return img.convert("RGB")
+            allow_redirects=True,
+            stream=True 
+        ) as response:
+            
+            response.raise_for_status()
+
+            # 1. Check Content-Length header if available
+            cl = response.headers.get("Content-Length")
+            if cl and int(cl) > MAX_SIZE:
+                raise RuntimeError(f"File too large (limit: {MAX_SIZE_MB:.0f}MB)")
+
+            # 2. Download in chunks to verify actual size (for servers with no header)
+            downloaded_data = io.BytesIO()
+            size = 0
+            for chunk in response.iter_content(chunk_size=8192):
+                size += len(chunk)
+                if size > MAX_SIZE:
+                    raise RuntimeError(f"File exceeds {MAX_SIZE_MB:.0f}MB limit.")
+                downloaded_data.write(chunk)
+
+            # Open and verify
+            downloaded_data.seek(0)
+            img = Image.open(downloaded_data)
+            img.load()
+            
+            return img.convert("RGB")
     
     except Exception as e:
-        print(f"Error fetching {url}: {str(e)}")
-        import gradio as gr
-        # This will show a red toast notification in the UI
-        raise gr.Error(f"Failed to fetch image: {str(e)}")
+        error_msg = f"Fetch Error: {str(e)}"
+        print(error_msg)
+        
+        # Try to raise a pretty Gradio error, otherwise raise standard RuntimeError
+        try:
+            import gradio as gr
+            raise gr.Error(error_msg)
+        except ImportError:
+            raise RuntimeError(error_msg)