509 lines
17 KiB
Python
509 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
import httpx
|
|
from PIL import Image
|
|
from typing_extensions import override
|
|
|
|
from sn_image_base.configs import global_configs, is_valid_base_url
|
|
from sn_image_base.exceptions import InvalidBaseUrlError, MissingApiKeyError
|
|
from sn_image_base.generation.core import ensure_output_path
|
|
from sn_image_base.generation.core.client_base import (
|
|
DEFAULT_HTTP_REQUEST_TIMEOUT,
|
|
DEFAULT_MAX_CONNECTIONS,
|
|
T2IBaseClient,
|
|
)
|
|
from sn_image_base.utils.error_utils import U1HttpErrorBase
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
DEFAULT_RESOLUTION: Literal["1K", "2K", "4K"] = "2K"
|
|
DEFAULT_ASPECT_RATIO = "16:9"
|
|
DEFAULT_POLL_INTERVAL = 5.0
|
|
OUTPUT_DIR = Path("/tmp/openclaw-sn-image")
|
|
|
|
|
|
IMAGE_GEN_ENDPOINT = "/images/generations"
|
|
|
|
|
|
class SensenovaText2ImageClient(T2IBaseClient):
|
|
"""Async client for Sensenova text-to-image API."""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str,
|
|
base_url: str | None = None,
|
|
*,
|
|
model: str | None = None,
|
|
max_connections: int = DEFAULT_MAX_CONNECTIONS,
|
|
timeout: float = DEFAULT_HTTP_REQUEST_TIMEOUT,
|
|
ssl_verify: bool = True,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Initialize the SensenovaText2ImageClient.
|
|
|
|
Args:
|
|
api_key (str):
|
|
API key for authentication.
|
|
base_url (str | None, optional):
|
|
API base URL. If None, reads from SN_IMAGE_GEN_BASE_URL env var.
|
|
model (str | None, optional):
|
|
Model name. If None, reads from SN_IMAGE_GEN_MODEL env var.
|
|
max_connections (int, optional):
|
|
Maximum number of connections. Defaults to 100.
|
|
timeout (float, optional):
|
|
Total timeout in seconds for HTTP requests.
|
|
Defaults to DEFAULT_HTTP_REQUEST_TIMEOUT.
|
|
ssl_verify (bool, optional):
|
|
If True, enable TLS verification. Defaults to True.
|
|
"""
|
|
api_key = api_key or global_configs.SN_IMAGE_GEN_API_KEY
|
|
if not api_key:
|
|
raise MissingApiKeyError(
|
|
"API key is missing: {}".format(
|
|
global_configs.get_env_var_help("SN_IMAGE_GEN_API_KEY")
|
|
)
|
|
)
|
|
base_url = base_url or global_configs.SN_IMAGE_GEN_BASE_URL
|
|
if not base_url:
|
|
raise InvalidBaseUrlError(
|
|
"Base URL is missing: {}".format(
|
|
global_configs.get_env_var_help("SN_IMAGE_GEN_BASE_URL")
|
|
)
|
|
)
|
|
if not is_valid_base_url(base_url):
|
|
raise InvalidBaseUrlError(
|
|
f"Base URL is not a valid base URL: {base_url}. "
|
|
f"Try setting environment variable(s): {global_configs.get_env_var_help('SN_IMAGE_GEN_BASE_URL')}"
|
|
)
|
|
super().__init__(
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
model=model,
|
|
max_connections=max_connections,
|
|
timeout=timeout,
|
|
ssl_verify=ssl_verify,
|
|
**kwargs,
|
|
)
|
|
|
|
@override
|
|
async def generate(
|
|
self,
|
|
prompt: str,
|
|
negative_prompt: str = "",
|
|
*,
|
|
model: str | None = None,
|
|
image_size: Literal["1K", "2K", "4K"] = DEFAULT_RESOLUTION,
|
|
aspect_ratio: str = DEFAULT_ASPECT_RATIO,
|
|
output_path: Path | None = None,
|
|
**kwargs: Any,
|
|
) -> dict:
|
|
"""Generate an image from text prompt.
|
|
|
|
Args:
|
|
prompt (str):
|
|
Text prompt for image generation.
|
|
negative_prompt (str, optional):
|
|
Negative prompt. Defaults to "".
|
|
model (str | None, optional):
|
|
Model name override. Defaults to None.
|
|
image_size (str, optional):
|
|
Image size preset ("1K", "2K", "4K"). Defaults to DEFAULT_RESOLUTION.
|
|
aspect_ratio (str, optional):
|
|
Aspect ratio (e.g. "16:9", "1:1"). Defaults to DEFAULT_ASPECT_RATIO.
|
|
output_path (Path | None, optional):
|
|
Output path for the generated image. Defaults to None.
|
|
**kwargs:
|
|
Additional arguments reserved for backend compatibility.
|
|
|
|
Returns:
|
|
dict:
|
|
Dictionary with keys: status, output (path), message.
|
|
"""
|
|
model = model or self.model or global_configs.SN_IMAGE_GEN_MODEL
|
|
# Normalize image_size to uppercase for NanoBanana API
|
|
image_size = image_size.upper() # type: ignore[assignment]
|
|
output_format = "png"
|
|
size = self._resolve_size(image_size, aspect_ratio)
|
|
payload = self.build_payload(
|
|
prompt=prompt,
|
|
model=model,
|
|
size=size,
|
|
aspect_ratio=aspect_ratio,
|
|
output_format=output_format,
|
|
)
|
|
headers = self.headers
|
|
api_url = self.get_api_url(model)
|
|
if output_path is None:
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
|
output_path = OUTPUT_DIR / f"t2i_{timestamp}.png"
|
|
output_path = ensure_output_path(output_path)
|
|
|
|
client = await self._get_client()
|
|
|
|
try:
|
|
create_response = await client.post(
|
|
api_url,
|
|
json=payload,
|
|
headers=headers,
|
|
)
|
|
data = self.parse_response(create_response)
|
|
except U1HttpErrorBase as exc:
|
|
details = exc.detail or ""
|
|
field_name = None
|
|
if exc.code == 404:
|
|
field_name = "SN_IMAGE_GEN_BASE_URL"
|
|
elif exc.code == 401:
|
|
field_name = "SN_IMAGE_GEN_API_KEY"
|
|
# elif exc.code == 400:
|
|
# warnings.warn(f"Bad request: {exc.message}; body: {payload}", stacklevel=2)
|
|
if field_name is not None:
|
|
field_hint = global_configs.get_annotated_field(field_name)
|
|
if field_hint is not None:
|
|
env_names = list(field_hint.env_names) if field_hint.env_names else []
|
|
if env_names:
|
|
if len(env_names) == 1:
|
|
details += (
|
|
f"\nIs the environment variable `{env_names[0]}` set correctly?"
|
|
)
|
|
else:
|
|
env_names_str = ", ".join([f"`{n}`" for n in env_names])
|
|
details += f"\nIs any of the following environment variable(s) set correctly: {env_names_str}?"
|
|
return {
|
|
"status": "failed",
|
|
"error": f"HTTP {exc.code}: {exc.message}",
|
|
"message": details,
|
|
}
|
|
try:
|
|
images_urls: list[str] = data["images_urls"]
|
|
if not images_urls:
|
|
return {
|
|
"status": "failed",
|
|
"error": "No image generated from the model",
|
|
}
|
|
url = images_urls[-1]
|
|
suffix = f".{output_format}"
|
|
save_path = output_path.with_suffix(suffix)
|
|
saved_path = await download_image(url, save_path)
|
|
return {
|
|
"status": "ok",
|
|
"output": str(saved_path),
|
|
"message": "Image generated successfully",
|
|
}
|
|
except httpx.HTTPStatusError as exc:
|
|
return {
|
|
"status": "failed",
|
|
"error": f"HTTP {exc.response.status_code}",
|
|
"message": f"http error: {exc.response.status_code} {exc.response.text}",
|
|
}
|
|
except (httpx.HTTPError, OSError, ValueError) as exc:
|
|
return {
|
|
"status": "failed",
|
|
"error": type(exc).__name__,
|
|
"message": f"request error: {exc}",
|
|
}
|
|
|
|
@property
|
|
@override
|
|
def api_key(self) -> str:
|
|
api_key = self._api_key or global_configs.SN_IMAGE_GEN_API_KEY
|
|
if not api_key:
|
|
raise ValueError(
|
|
"API key is missing: {}".format(
|
|
global_configs.get_env_var_help("SN_IMAGE_GEN_API_KEY")
|
|
)
|
|
)
|
|
return api_key
|
|
|
|
@property
|
|
@override
|
|
def base_url(self) -> str:
|
|
base_url = self._base_url or global_configs.SN_IMAGE_GEN_BASE_URL
|
|
if not base_url:
|
|
raise ValueError(
|
|
"Base URL is missing: {}".format(
|
|
global_configs.get_env_var_help("SN_IMAGE_GEN_BASE_URL")
|
|
)
|
|
)
|
|
if not is_valid_base_url(base_url):
|
|
raise ValueError(
|
|
f"Base URL is not a valid base URL: {base_url}. "
|
|
f"Try setting environment variable(s): {global_configs.get_env_var_help('SN_IMAGE_GEN_BASE_URL')}"
|
|
)
|
|
return base_url
|
|
|
|
@override
|
|
def get_api_url(self, _model: str | None = None) -> str:
|
|
base_url = self.base_url.rstrip("/")
|
|
path = IMAGE_GEN_ENDPOINT.lstrip("/")
|
|
api_url = f"{base_url}/{path}"
|
|
return api_url
|
|
|
|
@override
|
|
def build_payload(
|
|
self,
|
|
prompt: str,
|
|
model: str,
|
|
*,
|
|
size: str | None = None,
|
|
modalities: Sequence[str] = ("text", "image"),
|
|
output_format: Literal["png"] = "png",
|
|
response_format: Literal["url"] = "url",
|
|
**kwargs: Any,
|
|
) -> dict[str, Any]:
|
|
"""Build the payload for the SenseNova image-generation endpoint.
|
|
|
|
Args:
|
|
prompt (str): The prompt to generate an image for.
|
|
model (str): The model to use for generation.
|
|
size (str | None): Pixel size string (for example, "1920x1920").
|
|
modalities (Sequence[str]): Reserved for compatibility; currently not sent.
|
|
output_format (Literal["png"]): The output format of the image. Defaults to "png".
|
|
response_format (Literal["url"]): The response format of the image. Defaults to "url".
|
|
**kwargs (Any, optional): Additional parameters to pass to the API.
|
|
|
|
Example:
|
|
{
|
|
"model": "sensenova-u1-fast",
|
|
"prompt": "A cat wearing a hat",
|
|
"size": "1024x1024",
|
|
"response_format": "url",
|
|
"output_format": "png",
|
|
}
|
|
"""
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
# "modalities": modalities,
|
|
"size": size,
|
|
# "n": 1,
|
|
"response_format": response_format,
|
|
"output_format": output_format,
|
|
**kwargs,
|
|
}
|
|
return payload
|
|
|
|
@property
|
|
@override
|
|
def headers(self) -> dict[str, str]:
|
|
if not self.api_key:
|
|
raise MissingApiKeyError(
|
|
"API key is missing: {}".format(
|
|
global_configs.get_env_var_help("SN_IMAGE_GEN_API_KEY")
|
|
)
|
|
)
|
|
return {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
@classmethod
|
|
def _resolve_size(
|
|
cls,
|
|
resolution: Literal["1K", "2K"] | str | None = None,
|
|
aspect_ratio: ASPECT_RATIO_LITERALS | str | None = None,
|
|
) -> str | None:
|
|
"""Convert (resolution, aspect_ratio) to a pixel size string.
|
|
|
|
If aspect_ratio is None, returns the resolution as-is (e.g. "1K").
|
|
"""
|
|
if not resolution and not aspect_ratio:
|
|
return None
|
|
resolution = resolution or "2K"
|
|
aspect_ratio = aspect_ratio or "1:1"
|
|
if resolution == "1K":
|
|
buckets = BUCKETS_1K
|
|
elif resolution == "2K":
|
|
buckets = BUCKETS_2K
|
|
else:
|
|
raise ValueError(f"Unsupported resolution: {resolution!r}. Must be '1K' or '2K'.")
|
|
try:
|
|
ws, _, hs = aspect_ratio.strip().partition(":")
|
|
width = int(ws)
|
|
height = int(hs)
|
|
ratio = width / height
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid aspect ratio: {aspect_ratio!r}") from e
|
|
if ratio > 16 / 9:
|
|
raise ValueError(f"Aspect ratio {aspect_ratio!r} is too wide. Maximum is 16:9")
|
|
if ratio < 9 / 21:
|
|
raise ValueError(f"Aspect ratio {aspect_ratio!r} is too high. Maximum is 9:21")
|
|
w, h = _find_nearest_aspect_ratio(ratio, buckets)
|
|
return f"{w}x{h}"
|
|
|
|
@override
|
|
def parse_response(self, response: httpx.Response) -> dict:
|
|
"""Parse the response from the SenseNova image-generation endpoint.
|
|
|
|
Example response data:
|
|
|
|
```json
|
|
{
|
|
"data": [{
|
|
"url": "https://cdn.sensenova.dev/gen/..."
|
|
}]
|
|
}
|
|
```
|
|
|
|
Args:
|
|
response: The HTTP response from the SenseNova image-generation endpoint.
|
|
|
|
Returns:
|
|
dict: Parsed data with key ``images_urls``.
|
|
"""
|
|
raw_data = super().parse_response(response)
|
|
|
|
images_urls: list[str] = []
|
|
for item in raw_data.get("data", []):
|
|
url = item.get("url")
|
|
if url:
|
|
images_urls.append(url)
|
|
return {"images_urls": images_urls}
|
|
|
|
|
|
async def download_image(
|
|
url: str,
|
|
save_path: Path,
|
|
timeout: float = DEFAULT_HTTP_REQUEST_TIMEOUT,
|
|
) -> Path:
|
|
"""Download an image from a URL.
|
|
|
|
Args:
|
|
url: The URL of the image to download.
|
|
timeout: The timeout for the request.
|
|
|
|
Returns:
|
|
Path: The path to the downloaded image file.
|
|
"""
|
|
save_path.parent.mkdir(parents=True, exist_ok=True)
|
|
temp_path: Path | None = None
|
|
bytes_written = 0
|
|
expected_length: int | None = None
|
|
try:
|
|
with tempfile.NamedTemporaryFile(
|
|
dir=save_path.parent,
|
|
prefix=f".{save_path.name}.",
|
|
suffix=".tmp",
|
|
delete=False,
|
|
) as temp_file:
|
|
temp_path = Path(temp_file.name)
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
async with client.stream("GET", url) as response:
|
|
response.raise_for_status()
|
|
content_length = response.headers.get("content-length")
|
|
if content_length is not None:
|
|
expected_length = int(content_length)
|
|
async for chunk in response.aiter_bytes():
|
|
bytes_written += len(chunk)
|
|
temp_file.write(chunk)
|
|
temp_file.flush()
|
|
os.fsync(temp_file.fileno())
|
|
|
|
if expected_length is not None and bytes_written != expected_length:
|
|
raise OSError(
|
|
f"Downloaded image is incomplete: got {bytes_written} bytes, "
|
|
f"expected {expected_length} bytes"
|
|
)
|
|
|
|
assert temp_path is not None
|
|
_validate_image_file(temp_path)
|
|
temp_path.replace(save_path)
|
|
return save_path
|
|
except Exception:
|
|
if temp_path is not None:
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def _validate_image_file(image_path: Path) -> None:
|
|
"""Verify that the downloaded image can be decoded completely."""
|
|
with Image.open(image_path) as image:
|
|
image.verify()
|
|
with Image.open(image_path) as image:
|
|
image.load()
|
|
|
|
|
|
def mime_type_to_suffix(mime_type: str) -> str:
|
|
"""Convert MIME type to file suffix.
|
|
|
|
Args:
|
|
mime_type: MIME type.
|
|
|
|
Returns:
|
|
str: File suffix.
|
|
"""
|
|
if mime_type == "image/jpeg":
|
|
return ".jpg"
|
|
elif mime_type == "image/png":
|
|
return ".png"
|
|
elif mime_type == "image/webp":
|
|
return ".webp"
|
|
else:
|
|
return ".png"
|
|
|
|
|
|
ASPECT_RATIO_LITERALS = Literal[
|
|
"2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "1:1", "16:9", "9:16", "9:21"
|
|
]
|
|
BUCKETS_1K: dict[ASPECT_RATIO_LITERALS, tuple[int, int]] = {
|
|
"2:3": (1088, 1632),
|
|
"3:2": (1632, 1088),
|
|
"3:4": (1152, 1536),
|
|
"4:3": (1536, 1152),
|
|
"4:5": (1184, 1472),
|
|
"5:4": (1472, 1184),
|
|
"1:1": (1344, 1344),
|
|
"16:9": (1792, 992),
|
|
"9:16": (992, 1792),
|
|
"9:21": (864, 2048),
|
|
}
|
|
BUCKETS_2K: dict[ASPECT_RATIO_LITERALS, tuple[int, int]] = {
|
|
"2:3": (1664, 2496),
|
|
"3:2": (2496, 1664),
|
|
"3:4": (1760, 2368),
|
|
"4:3": (2368, 1760),
|
|
"4:5": (1824, 2272),
|
|
"5:4": (2272, 1824),
|
|
"1:1": (2048, 2048),
|
|
"16:9": (2752, 1536),
|
|
"9:16": (1536, 2752),
|
|
"9:21": (1344, 3136),
|
|
}
|
|
|
|
|
|
def _find_nearest_aspect_ratio(
|
|
ratio: float,
|
|
buckets: dict[ASPECT_RATIO_LITERALS, tuple[int, int]],
|
|
) -> tuple[int, int]:
|
|
wh_pairs = sorted(
|
|
buckets.values(),
|
|
key=lambda wh: abs(wh[0] / wh[1] - ratio),
|
|
)
|
|
return wh_pairs[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
|
|
async def main_async():
|
|
client = SensenovaText2ImageClient(
|
|
api_key=global_configs.SN_IMAGE_GEN_API_KEY,
|
|
base_url=global_configs.SN_IMAGE_GEN_BASE_URL,
|
|
)
|
|
|
|
result = await client.generate(
|
|
prompt="A cat wearing a hat",
|
|
image_size="1K",
|
|
aspect_ratio="16:9",
|
|
)
|
|
print(result)
|
|
|
|
asyncio.run(main_async())
|