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

@@ -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)