Zhengshou Lai
334ff647fc
- Add bin/api/ package with generate_image() and synthesize_tts() - Refactor bin/commands/image.py to call bin.api.image.generate_image - Refactor bin/commands/voice.py to call bin.api.voice.synthesize_tts - API layer is free of Click CLI dependencies, usable by external scripts
58 lines
1.4 KiB
Python
58 lines
1.4 KiB
Python
"""Image generation API client (Doubao/Seedream)."""
|
|
|
|
import base64
|
|
import os
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
_RATIO_MAP = {
|
|
"1:1": "1024x1024",
|
|
"4:3": "1024x768",
|
|
"3:4": "768x1024",
|
|
"16:9": "1280x720",
|
|
"9:16": "720x1280",
|
|
}
|
|
|
|
|
|
def generate_image(
|
|
api_key: str,
|
|
prompt: str,
|
|
size: str | None = None,
|
|
ratio: str | None = None,
|
|
output: str | None = None,
|
|
watermark: bool = True,
|
|
b64: bool = False,
|
|
) -> str:
|
|
"""Generate image using Doubao/Seedream API.
|
|
|
|
Returns absolute file path if saved, or URL string if output is None.
|
|
"""
|
|
from openai import OpenAI
|
|
|
|
resolved_size = size or (ratio and _RATIO_MAP.get(ratio)) or "2K"
|
|
|
|
client = OpenAI(
|
|
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
|
api_key=api_key,
|
|
)
|
|
|
|
resp = client.images.generate(
|
|
model="doubao-seedream-5-0-260128",
|
|
prompt=prompt,
|
|
size=resolved_size,
|
|
response_format="b64_json" if b64 else "url",
|
|
extra_body={"watermark": watermark},
|
|
)
|
|
|
|
if b64:
|
|
image_bytes = base64.b64decode(resp.data[0].b64_json)
|
|
save_path = output or "generated_image.png"
|
|
Path(save_path).write_bytes(image_bytes)
|
|
return os.path.abspath(save_path)
|
|
else:
|
|
url = resp.data[0].url
|
|
if output:
|
|
urllib.request.urlretrieve(url, output)
|
|
return os.path.abspath(output)
|
|
return url
|