fetch_url_util.py

import io
import urllib.parse
from PIL import Image
from curl_cffi import requests

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",
        "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(
            url, 
            headers=headers, 
            impersonate="chrome120", 
            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")
    
    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)}")