From 0a1426918707899773848f6aef2ae2dddb02854c Mon Sep 17 00:00:00 2001 From: hp0912 <809211365@qq.com> Date: Wed, 6 May 2026 00:27:31 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BC=98=E5=8C=96=E5=8F=91=E9=80=81?= =?UTF-8?q?=E5=9B=BE=E7=89=87=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../image-to-image/scripts/image_to_image.py | 55 +++++++++++++-- skills/text-to-image/scripts/text_to_image.py | 67 +++++++++++++++++-- 2 files changed, 110 insertions(+), 12 deletions(-) diff --git a/skills/image-to-image/scripts/image_to_image.py b/skills/image-to-image/scripts/image_to_image.py index 404b53b..a9f536d 100644 --- a/skills/image-to-image/scripts/image_to_image.py +++ b/skills/image-to-image/scripts/image_to_image.py @@ -281,12 +281,55 @@ def _rewrite_openai_image_url(url: str) -> str: return url -def _openai_images_from_response(response) -> list[str]: +def _extension_from_output_format(output_format: str) -> str: + if output_format == "jpeg": + return ".jpg" + if output_format == "webp": + return ".webp" + return ".png" + + +def _openai_response_value(item, key: str): + if isinstance(item, dict): + return item.get(key) + return getattr(item, key, None) + + +def _write_openai_b64_image(b64_json: str, output_format: str) -> str: + encoded = b64_json.strip() + suffix = _extension_from_output_format(output_format) + if encoded.startswith("data:"): + header, encoded = encoded.split(",", 1) + mime_type = header[5:].split(";", 1)[0].strip().lower() + if mime_type: + suffix = _extension_from_mime(mime_type) + + encoded = "".join(encoded.split()) + padding = len(encoded) % 4 + if padding: + encoded = f"{encoded}{'=' * (4 - padding)}" + + image_bytes = base64.b64decode(encoded) + with tempfile.NamedTemporaryFile(prefix="wechat-openai-image-", suffix=suffix, delete=False) as temp_file: + temp_file.write(image_bytes) + return temp_file.name + + +def _openai_images_from_response(response, output_format: str) -> list[str]: outputs: list[str] = [] - for item in getattr(response, "data", []) or []: - url = getattr(item, "url", None) - if url: - outputs.append(_rewrite_openai_image_url(str(url))) + try: + for item in getattr(response, "data", []) or []: + b64_json = _openai_response_value(item, "b64_json") + if b64_json: + outputs.append(_write_openai_b64_image(str(b64_json), output_format)) + continue + + url = _openai_response_value(item, "url") + if url: + outputs.append(_rewrite_openai_image_url(str(url))) + except Exception: + _cleanup_openai_temp_files(outputs) + raise return outputs @@ -541,7 +584,7 @@ def call_openai(config: dict, prompt: str, model: str, images: list[str], input_file.close() _debug_response("openai images.edit response", response) - return _openai_images_from_response(response) + return _openai_images_from_response(response, output_format) # --------------------------------------------------------------------------- diff --git a/skills/text-to-image/scripts/text_to_image.py b/skills/text-to-image/scripts/text_to_image.py index e5cd31a..2691b0e 100644 --- a/skills/text-to-image/scripts/text_to_image.py +++ b/skills/text-to-image/scripts/text_to_image.py @@ -3,11 +3,14 @@ from __future__ import annotations import argparse +import base64 import json +import mimetypes import os import re import subprocess import sys +import tempfile import time import traceback import urllib.parse @@ -280,12 +283,64 @@ def _rewrite_openai_image_url(url: str) -> str: return url -def _openai_images_from_response(response) -> list[str]: +def _extension_from_mime(mime_type: str) -> str: + if mime_type == "image/jpeg": + return ".jpg" + guessed = mimetypes.guess_extension(mime_type) + if guessed in {".png", ".jpg", ".jpeg", ".webp"}: + return guessed + return ".png" + + +def _extension_from_output_format(output_format: str) -> str: + if output_format == "jpeg": + return ".jpg" + if output_format == "webp": + return ".webp" + return ".png" + + +def _openai_response_value(item, key: str): + if isinstance(item, dict): + return item.get(key) + return getattr(item, key, None) + + +def _write_openai_b64_image(b64_json: str, output_format: str) -> str: + encoded = b64_json.strip() + suffix = _extension_from_output_format(output_format) + if encoded.startswith("data:"): + header, encoded = encoded.split(",", 1) + mime_type = header[5:].split(";", 1)[0].strip().lower() + if mime_type: + suffix = _extension_from_mime(mime_type) + + encoded = "".join(encoded.split()) + padding = len(encoded) % 4 + if padding: + encoded = f"{encoded}{'=' * (4 - padding)}" + + image_bytes = base64.b64decode(encoded) + with tempfile.NamedTemporaryFile(prefix="wechat-openai-image-", suffix=suffix, delete=False) as temp_file: + temp_file.write(image_bytes) + return temp_file.name + + +def _openai_images_from_response(response, output_format: str) -> list[str]: outputs: list[str] = [] - for item in getattr(response, "data", []) or []: - url = getattr(item, "url", None) - if url: - outputs.append(_rewrite_openai_image_url(str(url))) + try: + for item in getattr(response, "data", []) or []: + b64_json = _openai_response_value(item, "b64_json") + if b64_json: + outputs.append(_write_openai_b64_image(str(b64_json), output_format)) + continue + + url = _openai_response_value(item, "url") + if url: + outputs.append(_rewrite_openai_image_url(str(url))) + except Exception: + _cleanup_openai_temp_files(outputs) + raise return outputs @@ -498,7 +553,7 @@ def call_openai(config: dict, prompt: str, model: str, response = client.images.generate(**kwargs) _debug_response("openai images.generate response", response) - return _openai_images_from_response(response) + return _openai_images_from_response(response, output_format) # ---------------------------------------------------------------------------