93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
"""Async image generation and editing clients."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
import openai
|
|
|
|
if TYPE_CHECKING:
|
|
from io import BufferedReader, BytesIO
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_image_gen_client: openai.AsyncOpenAI | None = None
|
|
_image_edit_client: openai.AsyncOpenAI | None = None
|
|
|
|
|
|
def get_image_gen_client() -> openai.AsyncOpenAI:
|
|
"""Return the shared async image-generation client, building it once."""
|
|
global _image_gen_client
|
|
if _image_gen_client is None:
|
|
from vibe_bot.config import IMAGE_GEN_ENDPOINT, IMAGE_GEN_ENDPOINT_KEY
|
|
|
|
_image_gen_client = openai.AsyncOpenAI(
|
|
base_url=IMAGE_GEN_ENDPOINT,
|
|
api_key=IMAGE_GEN_ENDPOINT_KEY,
|
|
max_retries=0,
|
|
)
|
|
return _image_gen_client
|
|
|
|
|
|
def get_image_edit_client() -> openai.AsyncOpenAI:
|
|
"""Return the shared async image-edit client, building it once."""
|
|
global _image_edit_client
|
|
if _image_edit_client is None:
|
|
from vibe_bot.config import IMAGE_EDIT_ENDPOINT, IMAGE_EDIT_ENDPOINT_KEY
|
|
|
|
_image_edit_client = openai.AsyncOpenAI(
|
|
base_url=IMAGE_EDIT_ENDPOINT, api_key=IMAGE_EDIT_ENDPOINT_KEY
|
|
)
|
|
return _image_edit_client
|
|
|
|
|
|
async def image_generation(
|
|
prompt: str,
|
|
*,
|
|
model: str = "gen",
|
|
n: int = 1,
|
|
size: str = "1024x1024",
|
|
) -> str:
|
|
"""Generate an image; return base64 data ("" on failure)."""
|
|
client = get_image_gen_client()
|
|
try:
|
|
response = await client.images.generate(
|
|
prompt=prompt,
|
|
n=n,
|
|
size=size,
|
|
model=model,
|
|
timeout=300.0,
|
|
)
|
|
except openai.OpenAIError as e:
|
|
logger.warning("Image generation failed: %s", e)
|
|
return ""
|
|
if response.data:
|
|
return response.data[0].b64_json or ""
|
|
return ""
|
|
|
|
|
|
async def image_edit(
|
|
image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
|
|
prompt: str,
|
|
*,
|
|
model: str = "edit",
|
|
n: int = 1,
|
|
) -> str:
|
|
"""Edit an image; return base64 data ("" on failure)."""
|
|
client = get_image_edit_client()
|
|
try:
|
|
response = await client.images.edit(
|
|
image=image,
|
|
prompt=prompt,
|
|
n=n,
|
|
size="768x768",
|
|
model=model,
|
|
)
|
|
except openai.OpenAIError as e:
|
|
logger.warning("Image edit failed: %s", e)
|
|
return ""
|
|
if response.data:
|
|
return response.data[0].b64_json or ""
|
|
return ""
|