Smaug branding; public-address check for links; upload size cap; queue cap; 10 s GPU reservation

README.md

@@ -1,5 +1,5 @@
 ---
-title: AI Image Detector
+title: Smaug AI Image Detector
 emoji: 🔍
 colorFrom: blue
 colorTo: indigo
@@ -11,10 +11,11 @@ pinned: false
 license: apache-2.0
 models:
 - abacusai/Smaug-AI-Image-Detector
-short_description: Is this image AI-generated? Upload one or paste a link.
+short_description: Is this image AI-generated? Free, open-source detector
 ---
 
-Upload an image or paste a link to check whether it was made by an AI image generator, using
+Smaug AI Image Detector by Abacus.AI: upload an image or paste a link to check whether it was made by an AI image
+generator. Free and open source, running
 [abacusai/Smaug-AI-Image-Detector](https://huggingface.co/abacusai/Smaug-AI-Image-Detector).
 
 Example images: a public-domain Flickr photo (CC0, via Megalith-10m) and an image generated with Z-Image-Turbo (Apache 2.0).

app.py

@@ -1,7 +1,9 @@
-"""Demo for abacusai/Smaug-AI-Image-Detector: upload an image or paste a link, get the probability it is AI-generated."""
+"""Smaug AI Image Detector demo: upload an image or paste a link, get the probability it is AI-generated."""
 import io
-import os
+import ipaddress
+import socket
 import sys
+from urllib.parse import urljoin, urlparse
 
 import gradio as gr
 import requests
@@ -12,26 +14,46 @@ from PIL import Image
 
 MODEL = 'abacusai/Smaug-AI-Image-Detector'
 MAX_BYTES = 20 * 1024 * 1024
+UA = 'AbacusAI-SmaugAIImageDetector/1.0 (https://huggingface.co/spaces/abacusai/Smaug-AI-Image-Detector)'
 
-path = snapshot_download(MODEL, token=os.environ.get('HF_TOKEN'))  # token only needed while the model repo is private
+path = snapshot_download(MODEL)
 sys.path.insert(0, path)
 from smaug_image_detect import SmaugImageDetect  # noqa: E402 - ships with the model repo
 
 det = SmaugImageDetect.from_pretrained(path, device="cuda")  # ZeroGPU: place on cuda at startup, real GPU inside @spaces.GPU
 
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=10)  # scoring takes ~1-2 s; a short reservation leaves more of each visitor's GPU quota
 def score(img):
     return det(img)
 
 
-def fetch(url):
-    if not url.startswith(('http://', 'https://')):
+def check_public(url):
+    """Only fetch from public internet addresses: no localhost, private networks or cloud metadata endpoints."""
+    u = urlparse(url)
+    if u.scheme not in ('http', 'https') or not u.hostname:
         raise gr.Error('The link must start with http:// or https://')
     try:
-        r = requests.get(url, timeout=15, stream=True, headers={'User-Agent': 'AbacusAI-AIImageDetectorDemo/1.0 (https://huggingface.co/spaces/abacusai/Smaug-AI-Image-Detector-demo)'})
-        r.raise_for_status()
-        data = r.raw.read(MAX_BYTES + 1, decode_content=True)
+        addrs = {a[4][0] for a in socket.getaddrinfo(u.hostname, u.port or (443 if u.scheme == 'https' else 80))}
+    except socket.gaierror:
+        raise gr.Error('Could not find that website.')
+    if not all(ipaddress.ip_address(a.split('%')[0]).is_global for a in addrs):
+        raise gr.Error('That link does not point to a public website.')
+
+
+def fetch(url):
+    try:
+        for _ in range(4):  # the link plus up to 3 redirects, each checked before it is followed
+            check_public(url)
+            r = requests.get(url, timeout=15, stream=True, allow_redirects=False, headers={'User-Agent': UA})
+            if r.is_redirect:
+                url = urljoin(url, r.headers['Location'])
+                continue
+            r.raise_for_status()
+            data = r.raw.read(MAX_BYTES + 1, decode_content=True)
+            break
+        else:
+            raise gr.Error('That link redirects too many times.')
     except requests.RequestException as e:
         raise gr.Error(f'Could not download the image: {e}')
     if len(data) > MAX_BYTES:
@@ -54,9 +76,9 @@ def predict(image, url):
     return {'AI-generated': p, 'Real': 1 - p}, f'### {verdict}\nProbability AI-generated: **{p:.1%}** (threshold 50%)'
 
 
-with gr.Blocks(title='AI Image Detector') as demo:
-    gr.Markdown('# AI Image Detector\nIs this image AI-generated? Upload one or paste a link. '
-                f'Model: [{MODEL}](https://huggingface.co/{MODEL}) by [Abacus.AI](https://abacus.ai).')
+with gr.Blocks(title='Smaug AI Image Detector: is this image AI-generated?') as demo:
+    gr.Markdown('# Smaug AI Image Detector\nIs this image AI-generated? Upload one or paste a link. Free and open source: '
+                f'[{MODEL}](https://huggingface.co/{MODEL}) by [Abacus.AI](https://abacus.ai).')
     with gr.Row():
         with gr.Column():
             image = gr.Image(type='pil', label='Image', sources=['upload', 'clipboard'])
@@ -70,4 +92,4 @@ with gr.Blocks(title='AI Image Detector') as demo:
                 'with a small AI-edited region as real. Do not use it as the only basis for decisions about people or their work.')
     btn.click(predict, [image, url], [label, text])
 
-demo.launch()
+demo.queue(max_size=40).launch(max_file_size='20mb')