返回提交历史
Modified
g4f/image/copy_images.py
+83
-43
XFEstudio/gpt4free
feat(g4f/image/copy_images.py): improve image handling with Unicode support and safer encoding
d93853af
代码差异
1 个文件
+83
-43
@@ -6,7 +6,7 @@ import uuid
6
6
import asyncio
7
7
import hashlib
8
8
import re
9
from urllib.parse import quote_plus
9
from urllib.parse import quote, unquote
10
10
from aiohttp import ClientSession, ClientError
11
11
12
12
from ..typing import Optional, Cookies
@@ -15,26 +15,24 @@ from ..Provider.template import BackendApi
15
15
from . import is_accepted_format, extract_data_uri
16
16
from .. import debug
17
17
18
# Define the directory for generated images
18
# Directory for storing generated images
19
19
images_dir = "./generated_images"
20
20
21
21
def get_image_extension(image: str) -> str:
22
match = re.search(r"\.(?:jpe?g|png|webp)", image)
23
if match:
24
return match.group(0)
25
return ".jpg"
22
"""Extract image extension from URL or filename, default to .jpg"""
23
match = re.search(r"\.(jpe?g|png|webp)$", image, re.IGNORECASE)
24
return f".{match.group(1).lower()}" if match else ".jpg"
26
25
27
# Function to ensure the images directory exists
28
26
def ensure_images_dir():
27
"""Create images directory if it doesn't exist"""
29
28
os.makedirs(images_dir, exist_ok=True)
30
29
31
30
def get_source_url(image: str, default: str = None) -> str:
32
source_url = image.split("url=", 1)
33
if len(source_url) > 1:
34
source_url = source_url[1]
35
source_url = source_url.replace("%2F", "/").replace("%3A", ":").replace("%3F", "?").replace("%3D", "=")
36
if source_url.startswith("https://"):
37
return source_url
31
"""Extract original URL from image parameter if present"""
32
if "url=" in image:
33
decoded_url = unquote(image.split("url=", 1)[1])
34
if decoded_url.startswith(("http://", "https://")):
35
return decoded_url
38
36
return default
39
37
40
38
async def copy_images(
@@ -47,45 +45,87 @@ async def copy_images(
47
45
target: str = None,
48
46
ssl: bool = None
49
47
) -> list[str]:
48
"""
49
Download and store images locally with Unicode-safe filenames
50
Returns list of relative image URLs
51
"""
50
52
if add_url:
51
53
add_url = not cookies
52
54
ensure_images_dir()
55
53
56
async with ClientSession(
54
57
connector=get_connector(proxy=proxy),
55
58
cookies=cookies,
56
59
headers=headers,
57
60
) as session:
58
async def copy_image(image: str, target: str = None, headers: dict = headers, ssl: bool = ssl) -> str:
59
if target is None or len(images) > 1:
60
hash = hashlib.sha256(image.encode()).hexdigest()
61
target = f"{quote_plus('+'.join(alt.split()[:10]), '')[:100]}_{hash[:16]}" if alt else str(uuid.uuid4())
62
target = f"{int(time.time())}_{target}{get_image_extension(image)}"
63
target = os.path.join(images_dir, target)
64
if image.startswith("data:"):
65
with open(target, "wb") as f:
66
f.write(extract_data_uri(image))
67
else:
68
try:
69
if BackendApi.working and image.startswith(BackendApi.url) and headers is None:
70
headers = BackendApi.headers
71
ssl = BackendApi.ssl
72
async with session.get(image, ssl=ssl, headers=headers) as response:
61
async def copy_image(image: str, target: str = None) -> str:
62
"""Process individual image and return its local URL"""
63
target_path = None
64
try:
65
# Generate filename components
66
file_hash = hashlib.sha256(image.encode()).hexdigest()[:16]
67
timestamp = int(time.time())
68
69
# Sanitize alt text for filename (Unicode-safe)
70
if alt:
71
# Keep letters, numbers, basic punctuation and all Unicode chars
72
clean_alt = re.sub(
73
r'[^\w\s.-]', # Allow all Unicode word chars
74
'_',
75
unquote(alt).strip(),
76
flags=re.UNICODE
77
)
78
clean_alt = re.sub(r'[\s_]+', '_', clean_alt)[:100]
79
else:
80
clean_alt = "image"
81
82
# Build safe filename with full Unicode support
83
extension = get_image_extension(image)
84
filename = (
85
f"{timestamp}_"
86
f"{clean_alt}_"
87
f"{file_hash}"
88
f"{extension}"
89
)
90
target_path = os.path.join(images_dir, filename)
91
92
# Handle different image types
93
if image.startswith("data:"):
94
with open(target_path, "wb") as f:
95
f.write(extract_data_uri(image))
96
else:
97
# Apply BackendApi settings if needed
98
if BackendApi.working and image.startswith(BackendApi.url):
99
request_headers = BackendApi.headers if headers is None else headers
100
request_ssl = BackendApi.ssl
101
else:
102
request_headers = headers
103
request_ssl = ssl
104
105
async with session.get(image, ssl=request_ssl, headers=request_headers) as response:
73
106
response.raise_for_status()
74
with open(target, "wb") as f:
107
with open(target_path, "wb") as f:
75
108
async for chunk in response.content.iter_chunked(4096):
76
109
f.write(chunk)
77
except ClientError as e:
78
debug.log(f"copy_images failed: {e.__class__.__name__}: {e}")
79
if os.path.exists(target):
80
os.unlink(target)
81
return get_source_url(image, image)
82
if "." not in target:
83
with open(target, "rb") as f:
84
extension = is_accepted_format(f.read(12)).split("/")[-1]
85
extension = "jpg" if extension == "jpeg" else extension
86
new_target = f"{target}.{extension}"
87
os.rename(target, new_target)
88
target = new_target
89
return f"/images/{os.path.basename(target)}{'?url=' + image if add_url and not image.startswith('data:') else ''}"
90
110
91
return await asyncio.gather(*[copy_image(image, target) for image in images])
111
# Verify file format
112
if not os.path.splitext(target_path)[1]:
113
with open(target_path, "rb") as f:
114
file_header = f.read(12)
115
detected_type = is_accepted_format(file_header)
116
if detected_type:
117
new_ext = f".{detected_type.split('/')[-1]}"
118
os.rename(target_path, f"{target_path}{new_ext}")
119
target_path = f"{target_path}{new_ext}"
120
121
# Build URL with safe encoding
122
url_filename = quote(os.path.basename(target_path))
123
return f"/images/{url_filename}{'?url=' + quote(image) if add_url and not image.startswith('data:') else ''}"
124
125
except (ClientError, IOError, OSError) as e:
126
debug.log(f"Image processing failed: {e.__class__.__name__}: {e}")
127
if target_path and os.path.exists(target_path):
128
os.unlink(target_path)
129
return get_source_url(image, image)
130
131
return await asyncio.gather(*[copy_image(img, target) for img in images])