Commit a971160
Changed files (1)
src
ai
images
src/ai/images/openai_img.py
@@ -90,18 +90,20 @@ async def openai_image_generation(
func = openai.images.generate
async for chunk in await func(**params):
resp = trim_none(chunk.model_dump())
- if b64_json := resp.pop("b64_json", None):
+ b64_json = resp.pop("b64_json", None)
+ logger.trace(resp)
+ rtype = resp.pop("type", None)
+ is_completed = rtype in {"image_generation.completed", "image_edit.completed"}
+ if b64_json:
image_bytes = base64.b64decode(b64_json)
save_path = Path(DOWNLOAD_DIR) / f"{rand_string(10)}.{params.get('output_format', 'jpeg')}"
async with await anyio.open_file(save_path, "wb") as f:
await f.write(image_bytes)
- rtype = resp.pop("type", None)
- caption = f"{EMOJI_IMG_BOT}**{model_name}**\n" if rtype == "image_generation.completed" else f"⌛️**{model_name}** 中间结果...\n"
+ caption = f"{EMOJI_IMG_BOT}**{model_name}**\n" if is_completed else f"⌛️**{model_name}** 中间结果...\n"
status_msg = await status_msg.edit_media(InputMediaPhoto(str(save_path), caption=caption + prettify(resp)))
save_path.unlink(missing_ok=True)
- elif resp.pop("type", None) == "image_generation.completed":
+ elif is_completed:
await status_msg.edit_caption(f"{EMOJI_IMG_BOT}**{model_name}**\n{prettify(resp)}")
- logger.trace(resp)
if status_msg.photo:
return {"success": True}
except Exception as e: