File size: 3,598 Bytes
50f35c5 18120fc 50f35c5 c153ed5 18120fc f246636 c2421e6 18120fc 50f35c5 083a50d 50f35c5 083a50d 50f35c5 18120fc 50f35c5 c2421e6 50f35c5 c153ed5 50f35c5 c153ed5 50f35c5 c153ed5 50f35c5 c2421e6 c153ed5 c2421e6 c153ed5 c2421e6 c153ed5 18120fc c153ed5 c2421e6 c153ed5 18120fc c153ed5 50f35c5 c2421e6 c153ed5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import io
import os
import hashlib
import urllib.parse
from PIL import Image
from curl_cffi import requests
try:
import pillow_jxl
except ImportError:
pass
CACHE_DIR = "/tmp/image_fetch_cache"
MAX_CACHE_COUNT = 100
MAX_SIZE = 150 * 1024 * 1024
MAX_SIZE_MB = MAX_SIZE / (1024 * 1024)
# Ensure cache directory exists
os.makedirs(CACHE_DIR, exist_ok=True)
def _manage_cache_limit():
"""Removes the oldest files if cache exceeds MAX_CACHE_COUNT."""
files = [os.path.join(CACHE_DIR, f) for f in os.listdir(CACHE_DIR)]
if len(files) > MAX_CACHE_COUNT:
# Sort by access time (oldest first)
files.sort(key=os.path.getatime)
for i in range(len(files) - MAX_CACHE_COUNT):
try:
os.remove(files[i])
except OSError:
pass
def fetch_image_from_url(url):
if not url:
return None
url = url.strip()
if not url.startswith(("http://", "https://")):
url = "https://" + url
# Generate a unique filename based on the URL
url_hash = hashlib.sha256(url.encode()).hexdigest()
cache_path = os.path.join(CACHE_DIR, f"{url_hash}.png")
# 1. Check if image is in cache
if os.path.exists(cache_path):
try:
img = Image.open(cache_path)
img.load()
# Update access time so it's not cleared by cache manager
os.utime(cache_path, None)
return img.convert("RGB")
except Exception:
# If cached file is corrupt, remove it and proceed to fetch
os.remove(cache_path)
# 2. If not in cache, prepare to fetch
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/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",
}
response = None
try:
response = requests.get(
url,
headers=headers,
impersonate="chrome120",
timeout=15,
allow_redirects=True,
stream=True
)
response.raise_for_status()
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)")
downloaded_data = io.BytesIO()
size = 0
for chunk in response.iter_content(chunk_size=65536):
size += len(chunk)
if size > MAX_SIZE:
raise RuntimeError(f"File exceeds {MAX_SIZE_MB:.0f}MB limit.")
downloaded_data.write(chunk)
downloaded_data.seek(0)
img = Image.open(downloaded_data)
img.load()
# 3. Save to cache
try:
img.save(cache_path, "PNG")
_manage_cache_limit()
except Exception as cache_err:
print(f"Cache write failed: {cache_err}")
return img.convert("RGB")
except Exception as e:
error_msg = f"Fetch Error: {str(e)}"
print(error_msg)
try:
import gradio as gr
raise gr.Error(error_msg)
except ImportError:
raise RuntimeError(error_msg)
finally:
if response is not None:
response.close() |