fix(generate_image): improve input image handling and validate output filename extensions

This commit is contained in:
nblog
2026-02-09 17:24:24 +08:00
parent 5efb7329a3
commit b8bbc75db2

View File

@@ -71,7 +71,7 @@ def build_message_content(prompt: str, input_images):
content.append({"type": "image_url", "image_url": {"url": data_url}}) content.append({"type": "image_url", "image_url": {"url": data_url}})
return content return content
def parse_data_url(data_url: str): def parse_data_url(data_url: str):
if not data_url.startswith("data:") or ";base64," not in data_url: if not data_url.startswith("data:") or ";base64," not in data_url:
raise ValueError("Image URL is not a base64 data URL.") raise ValueError("Image URL is not a base64 data URL.")
header, encoded = data_url.split(",", 1) header, encoded = data_url.split(",", 1)
@@ -86,11 +86,18 @@ def build_message_content(prompt: str, input_images):
def resolve_output_paths(filename: str, image_count: int, mime: str): def resolve_output_paths(filename: str, image_count: int, mime: str):
output_path = Path(filename) output_path = Path(filename)
suffix = output_path.suffix suffix = output_path.suffix
if not suffix:
suffix = MIME_TO_EXT.get(mime, ".png") # Validate/correct suffix matches MIME type
expected_suffix = MIME_TO_EXT.get(mime, ".png")
if suffix and suffix.lower() != expected_suffix.lower():
print(f"Warning: filename extension '{suffix}' doesn't match returned MIME type '{mime}'. Using '{expected_suffix}' instead.")
suffix = expected_suffix
elif not suffix:
suffix = expected_suffix
output_path = output_path.with_suffix(suffix) output_path = output_path.with_suffix(suffix)
# Create parent directory if it doesn't exist (for absolute paths) # Create parent directory if it doesn't exist (for paths with parent directories, absolute or relative)
if output_path.parent and str(output_path.parent) != '.': if output_path.parent and str(output_path.parent) != '.':
output_path.parent.mkdir(parents=True, exist_ok=True) output_path.parent.mkdir(parents=True, exist_ok=True)