Release standalone PPT template workflow 0.5.0
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Internal backend implementations for image_gen.py."""
|
||||
@@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Black Forest Labs FLUX image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
BFL_API_KEY (required)
|
||||
BFL_BASE_URL (optional)
|
||||
BFL_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend bfl")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
poll_json,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
VALID_ASPECT_RATIOS = [
|
||||
"1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
]
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.bfl.ai"
|
||||
DEFAULT_MODEL = "flux-pro-1.1-ultra"
|
||||
|
||||
MODEL_ENDPOINTS = {
|
||||
"flux-pro-1.1": "/v1/flux-pro-1.1",
|
||||
"flux-pro-1.1-ultra": "/v1/flux-pro-1.1-ultra",
|
||||
"flux-dev": "/v1/flux-dev",
|
||||
}
|
||||
|
||||
ASPECT_RATIO_TO_DIMENSIONS = {
|
||||
"1:1": (1024, 1024),
|
||||
"2:3": (1024, 1536),
|
||||
"3:2": (1536, 1024),
|
||||
"3:4": (1024, 1365),
|
||||
"4:3": (1365, 1024),
|
||||
"4:5": (1024, 1280),
|
||||
"5:4": (1280, 1024),
|
||||
"9:16": (1024, 1820),
|
||||
"16:9": (1820, 1024),
|
||||
"21:9": (2048, 878),
|
||||
}
|
||||
|
||||
|
||||
def _submit_request(url: str, headers: dict, payload: dict) -> dict:
|
||||
"""Submit a BFL generation request and return the JSON response."""
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=180)
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "BFL generation request")
|
||||
return response.json()
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_BASE_URL) -> str:
|
||||
"""Generate one image with the Black Forest Labs backend."""
|
||||
del image_size # BFL quality is primarily controlled by model choice.
|
||||
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for BFL backend. "
|
||||
f"Supported: {VALID_ASPECT_RATIOS}"
|
||||
)
|
||||
|
||||
normalized_model = model.strip().lower()
|
||||
endpoint = MODEL_ENDPOINTS.get(normalized_model)
|
||||
if not endpoint:
|
||||
supported = sorted(MODEL_ENDPOINTS)
|
||||
raise ValueError(f"Unsupported BFL model '{model}'. Supported: {supported}")
|
||||
|
||||
headers = {
|
||||
"x-key": api_key,
|
||||
"accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"prompt_upsampling": False,
|
||||
"output_format": "png",
|
||||
}
|
||||
|
||||
if normalized_model.endswith("-ultra"):
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
payload["raw"] = False
|
||||
else:
|
||||
width, height = ASPECT_RATIO_TO_DIMENSIONS[aspect_ratio]
|
||||
payload["width"] = width
|
||||
payload["height"] = height
|
||||
|
||||
url = base_url.rstrip("/") + endpoint
|
||||
|
||||
print("[Black Forest Labs]")
|
||||
print(f" Model: {normalized_model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print()
|
||||
print(" [..] Submitting request...", end="", flush=True)
|
||||
start = time.time()
|
||||
request_payload = _submit_request(url, headers, payload)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Request accepted ({elapsed:.1f}s)")
|
||||
|
||||
polling_url = request_payload.get("polling_url")
|
||||
if not polling_url:
|
||||
raise RuntimeError(f"BFL response missing polling_url: {request_payload}")
|
||||
|
||||
print(" [..] Polling result...")
|
||||
result_payload = poll_json(
|
||||
polling_url,
|
||||
{"x-key": api_key, "accept": "application/json"},
|
||||
status_label="status",
|
||||
ready_values=["Ready"],
|
||||
failed_values=["Error", "Failed", "Request Moderated", "Content Moderated"],
|
||||
)
|
||||
|
||||
image_url = ((result_payload.get("result") or {}).get("sample"))
|
||||
if not image_url:
|
||||
raise RuntimeError(f"BFL result missing sample URL: {result_payload}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the BFL backend."""
|
||||
api_key = require_api_key(
|
||||
"BFL_API_KEY",
|
||||
message="No API key found. Set BFL_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("BFL_BASE_URL") or DEFAULT_BASE_URL
|
||||
resolved_model = model or os.environ.get("BFL_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,377 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Shared helpers for image generation backends.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("This is an internal helper module used by image_gen.py backends.")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage, ImageOps as PILImageOps
|
||||
HAS_PIL = True
|
||||
except ImportError:
|
||||
HAS_PIL = False
|
||||
|
||||
|
||||
MAX_RETRIES = 3
|
||||
RETRY_BASE_DELAY = 10
|
||||
RETRY_BACKOFF = 2
|
||||
|
||||
|
||||
def resolve_output_path(prompt: str, output_dir: str = None,
|
||||
filename: str = None, ext: str = ".png") -> str:
|
||||
"""Compute the final output file path based on parameters."""
|
||||
if filename:
|
||||
file_name = os.path.splitext(filename)[0]
|
||||
else:
|
||||
safe = "".join(c for c in prompt if c.isalnum() or c in (" ", "_")).rstrip()
|
||||
safe = safe.replace(" ", "_").lower()[:30]
|
||||
file_name = safe or "generated_image"
|
||||
|
||||
full_name = f"{file_name}{ext}"
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
return os.path.join(output_dir, full_name)
|
||||
return full_name
|
||||
|
||||
|
||||
CONTENT_TYPE_TO_EXT = {
|
||||
"image/png": ".png",
|
||||
"image/jpeg": ".jpg",
|
||||
"image/jpg": ".jpg",
|
||||
"image/webp": ".webp",
|
||||
"image/gif": ".gif",
|
||||
"image/bmp": ".bmp",
|
||||
"image/tiff": ".tiff",
|
||||
}
|
||||
|
||||
EXT_TO_PIL_FORMAT = {
|
||||
".png": "PNG",
|
||||
".jpg": "JPEG",
|
||||
".jpeg": "JPEG",
|
||||
".webp": "WEBP",
|
||||
".gif": "GIF",
|
||||
".bmp": "BMP",
|
||||
".tiff": "TIFF",
|
||||
".tif": "TIFF",
|
||||
}
|
||||
|
||||
|
||||
def detect_image_extension(image_bytes: bytes, content_type: str = None) -> str | None:
|
||||
"""Best-effort detection of the real image format."""
|
||||
if image_bytes.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
return ".png"
|
||||
if image_bytes.startswith(b"\xff\xd8\xff"):
|
||||
return ".jpg"
|
||||
if image_bytes.startswith(b"GIF87a") or image_bytes.startswith(b"GIF89a"):
|
||||
return ".gif"
|
||||
if image_bytes.startswith(b"RIFF") and image_bytes[8:12] == b"WEBP":
|
||||
return ".webp"
|
||||
if image_bytes.startswith(b"BM"):
|
||||
return ".bmp"
|
||||
if image_bytes.startswith((b"II*\x00", b"MM\x00*")):
|
||||
return ".tiff"
|
||||
if content_type:
|
||||
clean_type = content_type.split(";", 1)[0].strip().lower()
|
||||
if clean_type in CONTENT_TYPE_TO_EXT:
|
||||
return CONTENT_TYPE_TO_EXT[clean_type]
|
||||
return None
|
||||
|
||||
|
||||
DATA_URI_HEADER = re.compile(
|
||||
r"data:(?P<mime>image/[A-Za-z0-9.+-]+)(?P<params>;[^,]*)?,",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def decode_data_uri(value: str) -> tuple[bytes, str | None]:
|
||||
"""
|
||||
Decode a base64 image data URI into raw bytes plus its declared content type.
|
||||
|
||||
The declared type is returned so callers can hand it to `save_image_bytes` instead of
|
||||
assuming the payload matches the output extension.
|
||||
"""
|
||||
header = DATA_URI_HEADER.match(value.strip())
|
||||
if not header:
|
||||
raise ValueError("Expected a base64 image data URI (data:image/...;base64,...).")
|
||||
|
||||
params = (header.group("params") or "").lower()
|
||||
if "base64" not in params:
|
||||
raise ValueError("Only base64-encoded image data URIs are supported.")
|
||||
|
||||
payload = "".join(value.strip()[header.end():].split())
|
||||
payload += "=" * (-len(payload) % 4)
|
||||
return base64.urlsafe_b64decode(payload), header.group("mime").lower()
|
||||
|
||||
|
||||
def find_data_uri(content) -> str | None:
|
||||
"""
|
||||
Return the first base64 image data URI inside a chat completion `content` value.
|
||||
|
||||
OpenAI-compatible gateways differ here: some return a dedicated image field, others
|
||||
inline the image in the message text (often as ``)
|
||||
or in a content-part list.
|
||||
"""
|
||||
if isinstance(content, str):
|
||||
header = DATA_URI_HEADER.search(content)
|
||||
if not header:
|
||||
return None
|
||||
payload = re.match(r"[A-Za-z0-9+/=_-]*", content[header.end():]).group(0)
|
||||
return content[header.start():header.end()] + payload
|
||||
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict):
|
||||
nested = part.get("image_url")
|
||||
if isinstance(nested, dict):
|
||||
nested = nested.get("url")
|
||||
found = find_data_uri(nested if nested else part.get("text"))
|
||||
else:
|
||||
found = find_data_uri(part)
|
||||
if found:
|
||||
return found
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_extension(ext: str) -> str:
|
||||
"""Normalize equivalent image extensions to a canonical form."""
|
||||
ext = ext.lower()
|
||||
if ext == ".jpeg":
|
||||
return ".jpg"
|
||||
if ext == ".tif":
|
||||
return ".tiff"
|
||||
return ext
|
||||
|
||||
|
||||
def save_image_bytes(image_bytes: bytes, path: str, content_type: str = None) -> str:
|
||||
"""
|
||||
Save image bytes to disk while keeping the file extension and the real bytes aligned.
|
||||
|
||||
If the target extension differs from the actual bytes, transcode through Pillow when
|
||||
available. Otherwise fail loudly instead of writing a misleading file.
|
||||
"""
|
||||
target_ext = _normalize_extension(os.path.splitext(path)[1])
|
||||
actual_ext = _normalize_extension(detect_image_extension(image_bytes, content_type) or "")
|
||||
|
||||
if not target_ext:
|
||||
raise ValueError(f"Output path must include an image extension: {path}")
|
||||
|
||||
if actual_ext and target_ext == actual_ext:
|
||||
with open(path, "wb") as f:
|
||||
f.write(image_bytes)
|
||||
print(f" File saved to: {path}")
|
||||
report_resolution(path)
|
||||
return path
|
||||
|
||||
if not HAS_PIL:
|
||||
actual_label = actual_ext or "unknown"
|
||||
raise RuntimeError(
|
||||
f"Image format mismatch for {path}: target extension is {target_ext}, "
|
||||
f"but the actual image bytes are {actual_label}. "
|
||||
"Install Pillow to enable automatic format conversion."
|
||||
)
|
||||
|
||||
target_format = EXT_TO_PIL_FORMAT.get(target_ext)
|
||||
if not target_format:
|
||||
raise ValueError(f"Unsupported output image extension: {target_ext}")
|
||||
|
||||
with PILImage.open(io.BytesIO(image_bytes)) as source:
|
||||
image = PILImageOps.exif_transpose(source)
|
||||
try:
|
||||
if target_format == "JPEG":
|
||||
has_alpha = (
|
||||
image.mode in ("RGBA", "LA")
|
||||
or "transparency" in getattr(image, "info", {})
|
||||
)
|
||||
if has_alpha:
|
||||
rgba = image.convert("RGBA")
|
||||
alpha = rgba.getchannel("A")
|
||||
rgb = rgba.convert("RGB")
|
||||
converted = PILImage.new("RGB", image.size, (255, 255, 255))
|
||||
converted.paste(rgb, mask=alpha)
|
||||
rgb.close()
|
||||
alpha.close()
|
||||
rgba.close()
|
||||
if image is not source:
|
||||
image.close()
|
||||
image = converted
|
||||
elif image.mode != "RGB":
|
||||
converted = image.convert("RGB")
|
||||
if image is not source:
|
||||
image.close()
|
||||
image = converted
|
||||
image.save(path, format=target_format)
|
||||
finally:
|
||||
if image is not source:
|
||||
image.close()
|
||||
|
||||
if actual_ext and actual_ext != target_ext:
|
||||
print(f" Converted: {actual_ext} -> {target_ext}")
|
||||
print(f" File saved to: {path}")
|
||||
report_resolution(path)
|
||||
return path
|
||||
|
||||
|
||||
def validate_image_file(path: str) -> str:
|
||||
"""Require an existing regular file that Pillow can read as an image."""
|
||||
image_path = Path(path)
|
||||
if not image_path.exists():
|
||||
raise RuntimeError(f"Image output path does not exist: {path}")
|
||||
if not image_path.is_file():
|
||||
raise RuntimeError(f"Image output path is not a file: {path}")
|
||||
if not HAS_PIL:
|
||||
raise RuntimeError(
|
||||
"Pillow is required to verify generated images. "
|
||||
"Install it with: pip install Pillow"
|
||||
)
|
||||
|
||||
try:
|
||||
with PILImage.open(image_path) as image:
|
||||
image.verify()
|
||||
except (OSError, ValueError, SyntaxError) as exc:
|
||||
raise RuntimeError(f"Image output is not readable: {path}: {exc}") from exc
|
||||
return str(image_path)
|
||||
|
||||
|
||||
def report_resolution(path: str) -> None:
|
||||
"""Try to report image resolution using PIL."""
|
||||
if HAS_PIL:
|
||||
try:
|
||||
img = PILImage.open(path)
|
||||
print(f" Resolution: {img.size[0]}x{img.size[1]}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def normalize_image_size(image_size: str) -> str:
|
||||
"""Normalize image size input to standard format."""
|
||||
s = image_size.strip()
|
||||
upper = s.upper()
|
||||
if upper in ("1K", "2K", "4K"):
|
||||
return upper
|
||||
if upper in ("512PX", "512"):
|
||||
return "512px"
|
||||
return s
|
||||
|
||||
|
||||
def is_rate_limit_error(exc: Exception) -> bool:
|
||||
"""Check whether the exception appears to be rate limiting."""
|
||||
err_str = str(exc).lower()
|
||||
status_code = getattr(exc, "status_code", None)
|
||||
error_code = getattr(exc, "code", None)
|
||||
response = getattr(exc, "response", None)
|
||||
error_name = type(exc).__name__.lower()
|
||||
if (
|
||||
status_code == 429
|
||||
or error_code == 429
|
||||
or getattr(response, "status_code", None) == 429
|
||||
or error_name in {"ratelimiterror", "toomanyrequestserror"}
|
||||
):
|
||||
return True
|
||||
return (
|
||||
"429" in err_str
|
||||
or "rate limit" in err_str
|
||||
or "rate-limit" in err_str
|
||||
or "rate_limit" in err_str
|
||||
or "too many requests" in err_str
|
||||
or "quota" in err_str
|
||||
or "resource_exhausted" in err_str
|
||||
or "resource exhausted" in err_str
|
||||
or "throttl" in err_str
|
||||
)
|
||||
|
||||
|
||||
def retry_delay(attempt: int, rate_limited: bool) -> int:
|
||||
"""Return the retry delay for a given attempt."""
|
||||
if rate_limited:
|
||||
return RETRY_BASE_DELAY * (RETRY_BACKOFF ** attempt)
|
||||
return 5
|
||||
|
||||
|
||||
def download_image(url: str, path: str, headers: dict = None, timeout: int = 180) -> str:
|
||||
"""Download an image URL and save it to disk."""
|
||||
response = requests.get(url, headers=headers or {}, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return save_image_bytes(
|
||||
response.content,
|
||||
path,
|
||||
content_type=response.headers.get("Content-Type"),
|
||||
)
|
||||
|
||||
|
||||
def require_api_key(*candidates: str, message: str) -> str:
|
||||
"""Return the first non-empty env var from candidates or raise."""
|
||||
for name in candidates:
|
||||
value = os.environ.get(name)
|
||||
if value:
|
||||
return value
|
||||
raise ValueError(message)
|
||||
|
||||
|
||||
def http_error(response: requests.Response, label: str) -> RuntimeError:
|
||||
"""Convert an HTTP response into a readable RuntimeError."""
|
||||
body = response.text.strip()
|
||||
if len(body) > 500:
|
||||
body = body[:500] + "..."
|
||||
return RuntimeError(f"{label} failed ({response.status_code}): {body}")
|
||||
|
||||
|
||||
def poll_json(
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
*,
|
||||
interval_seconds: float = 2.0,
|
||||
timeout_seconds: int = 300,
|
||||
status_label: str = "status",
|
||||
ready_values: list[str] | None = None,
|
||||
failed_values: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""Poll a JSON endpoint until it reports a ready or failed status."""
|
||||
ready = {value.lower() for value in (ready_values or ["ready", "success", "succeeded"])}
|
||||
failed = {value.lower() for value in (failed_values or ["error", "failed", "fail"])}
|
||||
|
||||
start = time.time()
|
||||
while True:
|
||||
response = requests.get(url, headers=headers, timeout=180)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
raw_status = str(payload.get(status_label, "")).strip()
|
||||
status = raw_status.lower()
|
||||
|
||||
if raw_status:
|
||||
print(f" Status: {raw_status}")
|
||||
|
||||
if status in ready:
|
||||
return payload
|
||||
|
||||
if status in failed:
|
||||
raise RuntimeError(f"Remote generation failed: {payload}")
|
||||
|
||||
if time.time() - start > timeout_seconds:
|
||||
raise RuntimeError(
|
||||
f"Timed out after {timeout_seconds}s while polling {url}"
|
||||
)
|
||||
|
||||
time.sleep(interval_seconds)
|
||||
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
fal.ai image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
FAL_KEY / FAL_API_KEY (required)
|
||||
FAL_BASE_URL (optional)
|
||||
FAL_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend fal")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
VALID_ASPECT_RATIOS = ["1:1", "16:9", "9:16", "3:4", "4:3"]
|
||||
DEFAULT_ENDPOINT = "https://fal.run"
|
||||
DEFAULT_MODEL = "fal-ai/imagen3/fast"
|
||||
|
||||
|
||||
def _resolve_url(base_url: str, model: str) -> str:
|
||||
"""Resolve the full fal endpoint URL for a model."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith(model):
|
||||
return base
|
||||
return f"{base}/{model}"
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the fal.ai backend."""
|
||||
del image_size
|
||||
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for fal backend. "
|
||||
f"Supported: {VALID_ASPECT_RATIOS}"
|
||||
)
|
||||
|
||||
url = _resolve_url(base_url, model)
|
||||
headers = {
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"num_images": 1,
|
||||
}
|
||||
|
||||
print("[fal.ai]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "fal image generation")
|
||||
|
||||
data = response.json()
|
||||
images = data.get("images") or []
|
||||
image_url = images[0].get("url") if images else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"fal response missing image URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the fal.ai backend."""
|
||||
api_key = require_api_key(
|
||||
"FAL_KEY",
|
||||
"FAL_API_KEY",
|
||||
message="No API key found. Set FAL_KEY or FAL_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("FAL_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("FAL_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,235 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Gemini Image Generation Backend
|
||||
|
||||
Generates images via the Google GenAI API (Gemini).
|
||||
Used by image_gen.py as a backend module.
|
||||
|
||||
Configuration keys:
|
||||
GEMINI_API_KEY (required)
|
||||
GEMINI_BASE_URL (optional) Custom API endpoint for proxy services
|
||||
GEMINI_MODEL (optional) Override default model
|
||||
|
||||
Dependencies:
|
||||
pip install google-genai Pillow
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend gemini")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
save_image_bytes,
|
||||
)
|
||||
|
||||
|
||||
# ╔══════════════════════════════���═══════════════════════════════════╗
|
||||
# ║ Constants ║
|
||||
# ╚═══════════��═════════════════════════════��════════════════════════╝
|
||||
|
||||
VALID_ASPECT_RATIOS = [
|
||||
"1:1", "1:4", "1:8",
|
||||
"2:3", "3:2", "3:4", "4:1", "4:3",
|
||||
"4:5", "5:4", "8:1", "9:16", "16:9", "21:9"
|
||||
]
|
||||
|
||||
VALID_IMAGE_SIZES = ["512px", "1K", "2K", "4K"]
|
||||
|
||||
DEFAULT_MODEL = "gemini-3.1-flash-image-preview"
|
||||
|
||||
|
||||
# ╔══════���═══════════════════════════════════════════════��═══════════╗
|
||||
# ║ Image Generation ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = None) -> str:
|
||||
"""
|
||||
Image generation via Gemini API (streaming).
|
||||
|
||||
Returns:
|
||||
Path of the saved image file
|
||||
|
||||
Raises:
|
||||
RuntimeError: When generation fails
|
||||
"""
|
||||
if base_url:
|
||||
client = genai.Client(api_key=api_key, http_options={'base_url': base_url})
|
||||
else:
|
||||
client = genai.Client(api_key=api_key)
|
||||
|
||||
config_kwargs = {
|
||||
"response_modalities": ["IMAGE"],
|
||||
"image_config": types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
),
|
||||
}
|
||||
if "flash" in model.lower():
|
||||
config_kwargs["thinking_config"] = types.ThinkingConfig(
|
||||
thinking_level="MINIMAL",
|
||||
)
|
||||
config = types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
mode_label = "Proxy Mode" if base_url else "Official Mode"
|
||||
print(f"[Gemini - {mode_label}]")
|
||||
if base_url:
|
||||
print(f" Base URL: {base_url}")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Image Size: {image_size}")
|
||||
print()
|
||||
|
||||
start_time = time.time()
|
||||
print(f" [..] Generating...", end="", flush=True)
|
||||
|
||||
heartbeat_stop = threading.Event()
|
||||
|
||||
def _heartbeat():
|
||||
while not heartbeat_stop.is_set():
|
||||
heartbeat_stop.wait(5)
|
||||
if not heartbeat_stop.is_set():
|
||||
elapsed = time.time() - start_time
|
||||
print(f" {elapsed:.0f}s...", end="", flush=True)
|
||||
|
||||
hb_thread = threading.Thread(target=_heartbeat, daemon=True)
|
||||
hb_thread.start()
|
||||
|
||||
last_image_data = None
|
||||
chunk_count = 0
|
||||
total_bytes = 0
|
||||
|
||||
for chunk in client.models.generate_content_stream(
|
||||
model=model,
|
||||
contents=[prompt],
|
||||
config=config,
|
||||
):
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
if chunk.parts is None:
|
||||
continue
|
||||
|
||||
for part in chunk.parts:
|
||||
if part.text is not None:
|
||||
print(f"\n Model says: {part.text}", end="", flush=True)
|
||||
elif part.inline_data is not None:
|
||||
chunk_count += 1
|
||||
data_size = len(part.inline_data.data) if part.inline_data.data else 0
|
||||
total_bytes += data_size
|
||||
size_str = f"{data_size / 1024:.0f}KB" if data_size < 1048576 else f"{data_size / 1048576:.1f}MB"
|
||||
print(f"\n [OK] Chunk #{chunk_count} received ({size_str}, {elapsed:.1f}s)", end="", flush=True)
|
||||
last_image_data = part
|
||||
|
||||
heartbeat_stop.set()
|
||||
hb_thread.join(timeout=1)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n [DONE] Stream complete ({elapsed:.1f}s, {chunk_count} chunk(s), {total_bytes / 1024:.0f}KB total)")
|
||||
|
||||
if last_image_data is not None and last_image_data.inline_data is not None:
|
||||
if chunk_count > 1:
|
||||
print(f" Keeping the final chunk (highest quality).")
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return save_image_bytes(
|
||||
last_image_data.inline_data.data,
|
||||
path,
|
||||
content_type=getattr(last_image_data.inline_data, "mime_type", None),
|
||||
)
|
||||
|
||||
raise RuntimeError("No image was generated. The server may have refused the request.")
|
||||
|
||||
|
||||
# ╔���═══════════════════════════════════���═════════════════════════════╗
|
||||
# ║ Public Entry Point ║
|
||||
# ╚═════════════���═══════════════════════════���════════════════════════╝
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""
|
||||
Gemini image generation with automatic retry.
|
||||
|
||||
Reads credentials from the current process environment or a `.env` file:
|
||||
GEMINI_API_KEY
|
||||
GEMINI_BASE_URL
|
||||
GEMINI_MODEL (optional override)
|
||||
|
||||
Args:
|
||||
prompt: Prompt text
|
||||
aspect_ratio: Aspect ratio (e.g. "16:9", "1:1")
|
||||
image_size: Image size ("512px", "1K", "2K", "4K", case-insensitive)
|
||||
output_dir: Output directory
|
||||
filename: Output filename (without extension)
|
||||
model: Model name (default: gemini-3.1-flash-image-preview)
|
||||
max_retries: Maximum number of retries
|
||||
|
||||
Returns:
|
||||
Path of the saved image file
|
||||
"""
|
||||
api_key = os.environ.get("GEMINI_API_KEY")
|
||||
base_url = os.environ.get("GEMINI_BASE_URL")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"No API key found. Set GEMINI_API_KEY in the current environment or a .env file."
|
||||
)
|
||||
|
||||
if model is None:
|
||||
model = os.environ.get("GEMINI_MODEL") or DEFAULT_MODEL
|
||||
|
||||
image_size = normalize_image_size(image_size)
|
||||
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(f"Invalid aspect ratio '{aspect_ratio}'. Valid: {VALID_ASPECT_RATIOS}")
|
||||
|
||||
if image_size not in VALID_IMAGE_SIZES:
|
||||
raise ValueError(f"Invalid image size '{image_size}'. Valid: {VALID_IMAGE_SIZES}")
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(api_key, prompt,
|
||||
aspect_ratio, image_size, output_dir,
|
||||
filename, model, base_url)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if attempt < max_retries and is_rate_limit_error(e):
|
||||
delay = retry_delay(attempt, rate_limited=True)
|
||||
print(f"\n [WARN] Rate limit hit (attempt {attempt + 1}/{max_retries + 1}). "
|
||||
f"Waiting {delay}s before retry...")
|
||||
time.sleep(delay)
|
||||
elif attempt < max_retries:
|
||||
delay = retry_delay(attempt, rate_limited=False)
|
||||
print(f"\n [WARN] Error (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
else:
|
||||
break
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,165 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Ideogram image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
IDEOGRAM_API_KEY (required)
|
||||
IDEOGRAM_BASE_URL (optional)
|
||||
IDEOGRAM_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend ideogram")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
ASPECT_RATIO_MAP = {
|
||||
"1:1": "1x1",
|
||||
"1:4": "1x4",
|
||||
"2:3": "2x3",
|
||||
"3:2": "3x2",
|
||||
"3:4": "3x4",
|
||||
"4:1": "4x1",
|
||||
"4:3": "4x3",
|
||||
"4:5": "4x5",
|
||||
"5:4": "5x4",
|
||||
"9:16": "9x16",
|
||||
"16:9": "16x9",
|
||||
"21:9": "21x9",
|
||||
}
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.ideogram.ai"
|
||||
DEFAULT_MODEL = "ideogram-v3"
|
||||
MODEL_ALIASES = {"ideogram-v3", "v3"}
|
||||
|
||||
IMAGE_SIZE_TO_SPEED = {
|
||||
"512px": "TURBO",
|
||||
"1K": "DEFAULT",
|
||||
"2K": "QUALITY",
|
||||
"4K": "QUALITY",
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the Ideogram generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/generate"):
|
||||
return base
|
||||
return base + "/v1/ideogram-v3/generate"
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_BASE_URL) -> str:
|
||||
"""Generate one image with the Ideogram backend."""
|
||||
normalized_model = model.strip().lower()
|
||||
if normalized_model not in MODEL_ALIASES:
|
||||
raise ValueError(
|
||||
f"Unsupported Ideogram model '{model}'. Supported: {sorted(MODEL_ALIASES)}"
|
||||
)
|
||||
|
||||
mapped_ratio = ASPECT_RATIO_MAP.get(aspect_ratio)
|
||||
if not mapped_ratio:
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Ideogram backend. "
|
||||
f"Supported: {sorted(ASPECT_RATIO_MAP)}"
|
||||
)
|
||||
|
||||
rendering_speed = IMAGE_SIZE_TO_SPEED.get(image_size, "DEFAULT")
|
||||
url = _resolve_url(base_url)
|
||||
headers = {"Api-Key": api_key}
|
||||
files = {
|
||||
"prompt": (None, prompt),
|
||||
"aspect_ratio": (None, mapped_ratio),
|
||||
"rendering_speed": (None, rendering_speed),
|
||||
}
|
||||
|
||||
print("[Ideogram]")
|
||||
print(f" Model: {normalized_model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio} -> {mapped_ratio}")
|
||||
print(f" Render Speed: {rendering_speed}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, files=files, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "Ideogram generation")
|
||||
|
||||
payload = response.json()
|
||||
data = payload.get("data") or []
|
||||
image_url = data[0].get("url") if data else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"Ideogram response missing image URL: {payload}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Ideogram backend."""
|
||||
api_key = require_api_key(
|
||||
"IDEOGRAM_API_KEY",
|
||||
message="No API key found. Set IDEOGRAM_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("IDEOGRAM_BASE_URL") or DEFAULT_BASE_URL
|
||||
resolved_model = model or os.environ.get("IDEOGRAM_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,219 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MiniMax image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
MINIMAX_API_KEY (required)
|
||||
MINIMAX_BASE_URL (optional)
|
||||
MINIMAX_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend minimax")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
detect_image_extension,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
save_image_bytes,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT = "https://api.minimaxi.com/v1/image_generation"
|
||||
DEFAULT_MODEL = "image-01"
|
||||
|
||||
# International fallback: set MINIMAX_BASE_URL=https://api.minimax.io if needed
|
||||
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": (512, 512),
|
||||
"16:9": (640, 360),
|
||||
"4:3": (576, 432),
|
||||
"3:2": (624, 416),
|
||||
"2:3": (416, 624),
|
||||
"3:4": (432, 576),
|
||||
"9:16": (360, 640),
|
||||
"21:9": (672, 288),
|
||||
},
|
||||
"1K": {
|
||||
"1:1": (1024, 1024),
|
||||
"16:9": (1280, 720),
|
||||
"4:3": (1152, 864),
|
||||
"3:2": (1248, 832),
|
||||
"2:3": (832, 1248),
|
||||
"3:4": (864, 1152),
|
||||
"9:16": (720, 1280),
|
||||
"21:9": (1344, 576),
|
||||
},
|
||||
"2K": {
|
||||
"1:1": (2048, 2048),
|
||||
"16:9": (2048, 1152),
|
||||
"4:3": (2048, 1536),
|
||||
"3:2": (2048, 1368),
|
||||
"2:3": (1368, 2048),
|
||||
"3:4": (1536, 2048),
|
||||
"9:16": (1152, 2048),
|
||||
"21:9": (2048, 880),
|
||||
},
|
||||
"4K": {
|
||||
"1:1": (2048, 2048),
|
||||
"16:9": (2048, 1152),
|
||||
"4:3": (2048, 1536),
|
||||
"3:2": (2048, 1368),
|
||||
"2:3": (1368, 2048),
|
||||
"3:4": (1536, 2048),
|
||||
"9:16": (1152, 2048),
|
||||
"21:9": (2048, 880),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the MiniMax image generation endpoint.
|
||||
|
||||
Accepts three forms of MINIMAX_BASE_URL:
|
||||
- Full endpoint: https://api.minimax.io/v1/image_generation → used as-is
|
||||
- Versioned base: https://api.minimax.io/v1 → appends /image_generation
|
||||
- Root base: https://api.minimax.io → appends /v1/image_generation
|
||||
"""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/image_generation"):
|
||||
return base
|
||||
if base.endswith("/v1"):
|
||||
return base + "/image_generation"
|
||||
return base + "/v1/image_generation"
|
||||
|
||||
|
||||
def _resolve_dimensions(aspect_ratio: str, image_size: str) -> tuple[int, int]:
|
||||
"""Resolve width and height from the unified aspect_ratio/image_size pair."""
|
||||
normalized = normalize_image_size(image_size)
|
||||
dimensions = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not dimensions:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for MiniMax backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return dimensions
|
||||
|
||||
|
||||
def _extract_image_bytes(payload: dict) -> bytes | None:
|
||||
"""Extract image bytes from a MiniMax response payload."""
|
||||
data = payload.get("data") or {}
|
||||
image_base64 = data.get("image_base64") or []
|
||||
if image_base64:
|
||||
return base64.b64decode(image_base64[0])
|
||||
return None
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the MiniMax backend."""
|
||||
width, height = _resolve_dimensions(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"response_format": "base64",
|
||||
"n": 1,
|
||||
}
|
||||
|
||||
print("[MiniMax Image]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {width}x{height} (from image_size={image_size})")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "MiniMax image generation")
|
||||
|
||||
data = response.json()
|
||||
base_resp = data.get("base_resp") or {}
|
||||
if base_resp.get("status_code") not in (None, 0, "0"):
|
||||
raise RuntimeError(f"MiniMax image generation failed: {data}")
|
||||
|
||||
image_bytes = _extract_image_bytes(data)
|
||||
if not image_bytes:
|
||||
raise RuntimeError(f"MiniMax response missing image data: {data}")
|
||||
|
||||
ext = detect_image_extension(image_bytes) or ".jpeg"
|
||||
path = resolve_output_path(prompt, output_dir, filename, ext)
|
||||
return save_image_bytes(image_bytes, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the MiniMax backend."""
|
||||
api_key = require_api_key(
|
||||
"MINIMAX_API_KEY",
|
||||
message="No API key found. Set MINIMAX_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("MINIMAX_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("MINIMAX_MODEL") or DEFAULT_MODEL
|
||||
normalized_size = normalize_image_size(image_size)
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=normalized_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,168 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ModelScope image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
MODELSCOPE_API_KEY (required)
|
||||
MODELSCOPE_MODEL (optional)
|
||||
MODELSCOPE_BASE_URL (optional)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
poll_json,
|
||||
download_image
|
||||
)
|
||||
|
||||
DEFAULT_ENDPOINT = "https://api-inference.modelscope.cn"
|
||||
DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
# Resolution must be 64-aligned.
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": "1024*1024",
|
||||
"3:4": "768*1024",
|
||||
"4:3": "1024*768",
|
||||
"9:16": "576*1024",
|
||||
"16:9": "1024*576"
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1280*1280",
|
||||
"3:4": "960*1280",
|
||||
"4:3": "1280*960",
|
||||
"9:16": "576*1024",
|
||||
"16:9": "1024*576"
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "2048*2048",
|
||||
"3:4": "1536*2048",
|
||||
"4:3": "2048*1536",
|
||||
"9:16": "1152*2048",
|
||||
"16:9": "2048*1152"
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "2048*2048",
|
||||
"3:4": "1920*2560",
|
||||
"4:3": "2560*1920",
|
||||
"9:16": "1728*3072",
|
||||
"16:9": "3072*1728"
|
||||
}
|
||||
}
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the ModelScope generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/v1"):
|
||||
base = base.removesuffix("/v1")
|
||||
return base
|
||||
|
||||
def _resolve_size(aspect_ratio: str, image_size: str) -> str:
|
||||
"""Resolve the target resolution for a ratio and logical size preset.
|
||||
|
||||
Args:
|
||||
aspect_ratio (str): The aspect ratio string. Supported values: '1:1', '3:4', '4:3', '9:16', '16:9'.
|
||||
image_size (str): The logical size preset. Supported values: '512px', '1K', '2K', '4K'.
|
||||
"""
|
||||
normalized = normalize_image_size(image_size)
|
||||
size = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not size:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for ModelScope backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return size
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the ModelScope backend."""
|
||||
size = _resolve_size(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)+'/v1/images/generations'
|
||||
common_headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size.replace("*", "x"),
|
||||
|
||||
}
|
||||
|
||||
print("[ModelScope Models]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers={**common_headers,"X-ModelScope-Async-Mode": "true"}, json=payload, timeout=300)
|
||||
|
||||
if (response.status_code != 200):
|
||||
raise http_error(response, "ModelScope image generation")
|
||||
|
||||
task_id = response.json()["task_id"]
|
||||
data = poll_json(
|
||||
url=f"{_resolve_url(base_url)}/v1/tasks/{task_id}",
|
||||
headers={**common_headers, "X-ModelScope-Task-Type": "image_generation"},
|
||||
status_label="task_status",
|
||||
ready_values=["SUCCEED"],
|
||||
failed_values=["FAILED"],
|
||||
)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(data["output_images"][0], path)
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the ModelScope backend."""
|
||||
api_key = require_api_key(
|
||||
"MODELSCOPE_API_KEY",
|
||||
message="No API key found. Set MODELSCOPE_API_KEY in the current environment or the project-root .env.",
|
||||
)
|
||||
base_url = os.environ.get("MODELSCOPE_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("MODELSCOPE_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
|
||||
@@ -0,0 +1,688 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
OpenAI Compatible Image Generation Backend
|
||||
|
||||
Generates images via OpenAI-compatible APIs (OpenAI, local models like Qwen-Image, etc.).
|
||||
Used by image_gen.py as a backend module.
|
||||
|
||||
Configuration keys:
|
||||
OPENAI_API_KEY (required) API key
|
||||
OPENAI_BASE_URL (optional) Custom API endpoint (e.g. http://127.0.0.1:3000/v1)
|
||||
OPENAI_MODEL (optional) Model name (default: gpt-image-2)
|
||||
OPENAI_SIZE_PRESET (optional) auto, legacy, gpt-image, gpt-image-2, or dall-e-2
|
||||
OPENAI_RESPONSE_FORMAT (optional) auto, b64_json, url, or omit
|
||||
OPENAI_QUALITY (optional) auto, omit, low, medium, high, standard, or hd
|
||||
OPENAI_OUTPUT_FORMAT (optional) png, jpeg, or webp for GPT image models
|
||||
OPENAI_OUTPUT_COMPRESSION (optional) 0-100, only for jpeg/webp GPT image output
|
||||
OPENAI_BACKGROUND (optional) auto or opaque for gpt-image-2
|
||||
OPENAI_MODERATION (optional) auto or low for GPT image models
|
||||
|
||||
Image editing (image-to-image):
|
||||
When image_gen.py passes reference_image=<path> (single-image CLI only,
|
||||
via --reference-image), this backend calls /v1/images/edits with the source
|
||||
image + the prompt as an edit instruction, instead of /v1/images/generations.
|
||||
|
||||
Dependencies:
|
||||
pip install requests Pillow
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend openai")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
|
||||
import requests
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
save_image_bytes,
|
||||
)
|
||||
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Constants ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
# Aspect ratio -> DALL-E 3 / legacy compatible size mapping.
|
||||
# Unknown OpenAI-compatible models use this table to preserve old behavior.
|
||||
LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE = {
|
||||
"1:1": "1024x1024",
|
||||
"16:9": "1792x1024",
|
||||
"9:16": "1024x1792",
|
||||
"3:2": "1536x1024",
|
||||
"2:3": "1024x1536",
|
||||
"4:3": "1536x1024", # closest available
|
||||
"3:4": "1024x1536", # closest available
|
||||
"4:5": "1024x1024", # fallback to square
|
||||
"5:4": "1024x1024", # fallback to square
|
||||
"21:9": "1792x1024", # closest wide format
|
||||
}
|
||||
|
||||
# GPT Image 1/1.5/mini officially support only square, landscape, portrait, or auto.
|
||||
GPT_IMAGE_LEGACY_ASPECT_RATIO_TO_SIZE = {
|
||||
"1:1": "1024x1024",
|
||||
"16:9": "1536x1024",
|
||||
"9:16": "1024x1536",
|
||||
"3:2": "1536x1024",
|
||||
"2:3": "1024x1536",
|
||||
"4:3": "1536x1024",
|
||||
"3:4": "1024x1536",
|
||||
"4:5": "1024x1536",
|
||||
"5:4": "1536x1024",
|
||||
"21:9": "1536x1024",
|
||||
}
|
||||
|
||||
# GPT Image 2 supports flexible sizes when both edges are multiples of 16,
|
||||
# the edge ratio is <= 3:1, and the total pixels are within model limits.
|
||||
GPT_IMAGE_2_SIZES = {
|
||||
"512px": {
|
||||
"1:1": "1024x1024", "16:9": "1280x720", "9:16": "720x1280",
|
||||
"3:2": "1248x832", "2:3": "832x1248", "4:3": "1024x768",
|
||||
"3:4": "768x1024", "4:5": "896x1120", "5:4": "1120x896",
|
||||
"21:9": "1280x544",
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1024x1024", "16:9": "1280x720", "9:16": "720x1280",
|
||||
"3:2": "1248x832", "2:3": "832x1248", "4:3": "1024x768",
|
||||
"3:4": "768x1024", "4:5": "896x1120", "5:4": "1120x896",
|
||||
"21:9": "1280x544",
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "2048x2048", "16:9": "2048x1152", "9:16": "1152x2048",
|
||||
"3:2": "2016x1344", "2:3": "1344x2016", "4:3": "1920x1440",
|
||||
"3:4": "1440x1920", "4:5": "1600x2000", "5:4": "2000x1600",
|
||||
"21:9": "2560x1088",
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "2880x2880", "16:9": "3840x2160", "9:16": "2160x3840",
|
||||
"3:2": "3520x2352", "2:3": "2352x3520", "4:3": "3264x2448",
|
||||
"3:4": "2448x3264", "4:5": "2560x3200", "5:4": "3200x2560",
|
||||
"21:9": "3840x1648",
|
||||
},
|
||||
}
|
||||
|
||||
DALL_E_2_SIZE_BY_IMAGE_SIZE = {
|
||||
"512px": "512x512",
|
||||
"1K": "1024x1024",
|
||||
"2K": "1024x1024",
|
||||
"4K": "1024x1024",
|
||||
}
|
||||
|
||||
VALID_ASPECT_RATIOS = list(LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE.keys())
|
||||
|
||||
# image_size -> quality mapping
|
||||
IMAGE_SIZE_TO_QUALITY = {
|
||||
"512px": "low",
|
||||
"1K": "medium",
|
||||
"2K": "high",
|
||||
"4K": "high",
|
||||
}
|
||||
|
||||
DEFAULT_MODEL = "gpt-image-2"
|
||||
|
||||
GPT_IMAGE_2_MIN_PIXELS = 655_360
|
||||
GPT_IMAGE_2_MAX_PIXELS = 8_294_400
|
||||
GPT_IMAGE_2_MAX_EDGE = 3840
|
||||
GPT_IMAGE_2_MAX_RATIO = 3
|
||||
|
||||
GPT_IMAGE_OUTPUT_FORMATS = {"png", "jpeg", "webp"}
|
||||
GPT_IMAGE_OUTPUT_EXTENSIONS = {
|
||||
"png": ".png",
|
||||
"jpeg": ".jpg",
|
||||
"webp": ".webp",
|
||||
}
|
||||
OPENAI_SIZE_PRESETS = {
|
||||
"auto",
|
||||
"legacy",
|
||||
"gpt-image",
|
||||
"gpt-image-legacy",
|
||||
"gpt-image-2",
|
||||
"dall-e-2",
|
||||
"dalle-2",
|
||||
}
|
||||
OPENAI_RESPONSE_FORMATS = {"auto", "b64_json", "url", "omit"}
|
||||
OPENAI_QUALITY_VALUES = {
|
||||
"auto",
|
||||
"omit",
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"standard",
|
||||
"hd",
|
||||
}
|
||||
GPT_IMAGE_BACKGROUNDS = {"auto", "opaque", "transparent"}
|
||||
GPT_IMAGE_MODERATION_VALUES = {"auto", "low"}
|
||||
DEFAULT_BASE_URL = "https://api.openai.com/v1"
|
||||
|
||||
# Signals to image_gen.py that this backend can accept a reference_image
|
||||
# (image-to-image edit via /v1/images/edits). Other backends omit this marker.
|
||||
SUPPORTS_REFERENCE_IMAGE = True
|
||||
|
||||
|
||||
def _field(value, name: str):
|
||||
"""Read a response field from either an SDK object or a dict."""
|
||||
if isinstance(value, Mapping):
|
||||
return value.get(name)
|
||||
return getattr(value, name, None)
|
||||
|
||||
|
||||
def _normalized_model(model: str) -> str:
|
||||
return (model or "").strip().lower()
|
||||
|
||||
|
||||
def _is_gpt_image_model(model: str) -> bool:
|
||||
return _normalized_model(model).startswith("gpt-image-")
|
||||
|
||||
|
||||
def _is_gpt_image_2(model: str) -> bool:
|
||||
return _normalized_model(model).startswith("gpt-image-2")
|
||||
|
||||
|
||||
def _is_dall_e_2(model: str) -> bool:
|
||||
return _normalized_model(model) == "dall-e-2"
|
||||
|
||||
|
||||
def _parse_size(size: str) -> tuple[int, int]:
|
||||
try:
|
||||
width_s, height_s = size.lower().split("x", 1)
|
||||
return int(width_s), int(height_s)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Invalid image size '{size}'. Expected WIDTHxHEIGHT.") from exc
|
||||
|
||||
|
||||
def _validate_gpt_image_2_size(size: str) -> None:
|
||||
width, height = _parse_size(size)
|
||||
total_pixels = width * height
|
||||
long_edge = max(width, height)
|
||||
short_edge = min(width, height)
|
||||
|
||||
errors = []
|
||||
if long_edge > GPT_IMAGE_2_MAX_EDGE:
|
||||
errors.append(f"max edge {long_edge}px exceeds {GPT_IMAGE_2_MAX_EDGE}px")
|
||||
if width % 16 != 0 or height % 16 != 0:
|
||||
errors.append("both edges must be multiples of 16px")
|
||||
if long_edge / short_edge > GPT_IMAGE_2_MAX_RATIO:
|
||||
errors.append("long:short edge ratio must not exceed 3:1")
|
||||
if not (GPT_IMAGE_2_MIN_PIXELS <= total_pixels <= GPT_IMAGE_2_MAX_PIXELS):
|
||||
errors.append(
|
||||
f"total pixels {total_pixels:,} must be between "
|
||||
f"{GPT_IMAGE_2_MIN_PIXELS:,} and {GPT_IMAGE_2_MAX_PIXELS:,}"
|
||||
)
|
||||
if errors:
|
||||
raise ValueError(f"Invalid gpt-image-2 size '{size}': {', '.join(errors)}")
|
||||
|
||||
|
||||
def _select_size(
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
image_size: str,
|
||||
size_preset: str | None = None,
|
||||
) -> str:
|
||||
"""Select a model-compatible size while preserving legacy fallbacks."""
|
||||
preset = size_preset or "auto"
|
||||
if preset in {"gpt-image-2"} or (preset == "auto" and _is_gpt_image_2(model)):
|
||||
size = GPT_IMAGE_2_SIZES[image_size][aspect_ratio]
|
||||
_validate_gpt_image_2_size(size)
|
||||
return size
|
||||
if preset in {"gpt-image", "gpt-image-legacy"} or (
|
||||
preset == "auto" and _is_gpt_image_model(model)
|
||||
):
|
||||
return GPT_IMAGE_LEGACY_ASPECT_RATIO_TO_SIZE[aspect_ratio]
|
||||
if preset in {"dall-e-2", "dalle-2"} or (preset == "auto" and _is_dall_e_2(model)):
|
||||
return DALL_E_2_SIZE_BY_IMAGE_SIZE[image_size]
|
||||
return LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE[aspect_ratio]
|
||||
|
||||
|
||||
def _supports_response_format(model: str) -> bool:
|
||||
"""GPT Image models always return base64; DALL-E/compatible models may need this."""
|
||||
return not _is_gpt_image_model(model)
|
||||
|
||||
|
||||
def _read_env_choice(name: str, allowed: set[str]) -> str | None:
|
||||
value = os.environ.get(name)
|
||||
if value is None or not value.strip():
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
if normalized not in allowed:
|
||||
allowed_list = ", ".join(sorted(allowed))
|
||||
raise ValueError(f"Invalid {name}='{value}'. Supported: {allowed_list}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _read_env_int(name: str, minimum: int, maximum: int) -> int | None:
|
||||
value = os.environ.get(name)
|
||||
if value is None or not value.strip():
|
||||
return None
|
||||
try:
|
||||
parsed = int(value)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Invalid {name}='{value}'. Expected integer {minimum}-{maximum}.") from exc
|
||||
if not (minimum <= parsed <= maximum):
|
||||
raise ValueError(f"Invalid {name}={parsed}. Expected integer {minimum}-{maximum}.")
|
||||
return parsed
|
||||
|
||||
|
||||
def _gpt_image_options(model: str) -> tuple[dict, str]:
|
||||
"""Read optional GPT Image request parameters from environment."""
|
||||
output_format = _read_env_choice("OPENAI_OUTPUT_FORMAT", GPT_IMAGE_OUTPUT_FORMATS)
|
||||
output_ext = GPT_IMAGE_OUTPUT_EXTENSIONS[output_format] if output_format else ".png"
|
||||
options = {}
|
||||
if output_format:
|
||||
options["output_format"] = output_format
|
||||
|
||||
output_compression = _read_env_int("OPENAI_OUTPUT_COMPRESSION", 0, 100)
|
||||
if output_compression is not None:
|
||||
if output_format not in {"jpeg", "webp"}:
|
||||
raise ValueError(
|
||||
"OPENAI_OUTPUT_COMPRESSION is only supported when "
|
||||
"OPENAI_OUTPUT_FORMAT is jpeg or webp."
|
||||
)
|
||||
options["output_compression"] = output_compression
|
||||
|
||||
background = _read_env_choice("OPENAI_BACKGROUND", GPT_IMAGE_BACKGROUNDS)
|
||||
if background:
|
||||
if _is_gpt_image_2(model) and background == "transparent":
|
||||
raise ValueError("gpt-image-2 does not support OPENAI_BACKGROUND=transparent.")
|
||||
options["background"] = background
|
||||
|
||||
moderation = _read_env_choice("OPENAI_MODERATION", GPT_IMAGE_MODERATION_VALUES)
|
||||
if moderation:
|
||||
options["moderation"] = moderation
|
||||
|
||||
return options, output_ext
|
||||
|
||||
|
||||
def _image_generations_url(base_url: str | None) -> str:
|
||||
base = (base_url or DEFAULT_BASE_URL).rstrip("/")
|
||||
if base.endswith("/images/generations"):
|
||||
return base
|
||||
return f"{base}/images/generations"
|
||||
|
||||
|
||||
def _image_edits_url(base_url: str | None) -> str:
|
||||
base = (base_url or DEFAULT_BASE_URL).rstrip("/")
|
||||
if base.endswith("/images/edits"):
|
||||
return base
|
||||
if base.endswith("/images/generations"):
|
||||
# Swap the sibling endpoint rather than appending to a full URL.
|
||||
return base[: -len("/generations")] + "/edits"
|
||||
return f"{base}/images/edits"
|
||||
|
||||
|
||||
def _read_size_preset() -> str | None:
|
||||
"""Read the optional size mapping preset for OpenAI-compatible providers."""
|
||||
return _read_env_choice("OPENAI_SIZE_PRESET", OPENAI_SIZE_PRESETS)
|
||||
|
||||
|
||||
def _read_response_format() -> str | None:
|
||||
"""Read the optional response_format override."""
|
||||
return _read_env_choice("OPENAI_RESPONSE_FORMAT", OPENAI_RESPONSE_FORMATS)
|
||||
|
||||
|
||||
def _read_quality(image_size: str) -> str | None:
|
||||
"""Resolve the quality field for OpenAI-compatible requests."""
|
||||
quality = _read_env_choice("OPENAI_QUALITY", OPENAI_QUALITY_VALUES)
|
||||
if quality == "omit":
|
||||
return None
|
||||
if quality and quality != "auto":
|
||||
return quality
|
||||
return IMAGE_SIZE_TO_QUALITY.get(image_size, "auto")
|
||||
|
||||
|
||||
def _apply_response_format(request: dict, model: str) -> None:
|
||||
"""Apply response_format while preserving the existing default behavior."""
|
||||
response_format = _read_response_format()
|
||||
if response_format == "omit":
|
||||
return
|
||||
if response_format in {"b64_json", "url"}:
|
||||
request["response_format"] = response_format
|
||||
return
|
||||
if _supports_response_format(model):
|
||||
request["response_format"] = "b64_json"
|
||||
|
||||
|
||||
def _post_image_generation(api_key: str, base_url: str | None, request: dict) -> dict:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
response = requests.post(
|
||||
_image_generations_url(base_url),
|
||||
headers=headers,
|
||||
json=request,
|
||||
timeout=300,
|
||||
)
|
||||
if not response.ok:
|
||||
raise http_error(response, "OpenAI image generation")
|
||||
try:
|
||||
return response.json()
|
||||
except ValueError as exc:
|
||||
raise RuntimeError("OpenAI image generation returned invalid JSON.") from exc
|
||||
|
||||
|
||||
def _post_image_edit(api_key: str, base_url: str | None,
|
||||
data: dict, image_path: str) -> dict:
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
# GPT Image models take the image list field 'image[]'; dall-e-2 and other
|
||||
# OpenAI-compatible edit models use the singular 'image'.
|
||||
model = str(data.get("model", ""))
|
||||
field = "image[]" if _is_gpt_image_model(model) else "image"
|
||||
mime_type = mimetypes.guess_type(image_path)[0] or "application/octet-stream"
|
||||
# Let requests build the multipart/form-data body (and its Content-Type
|
||||
# boundary) from files=; do not set Content-Type by hand.
|
||||
with open(image_path, "rb") as image_file:
|
||||
response = requests.post(
|
||||
_image_edits_url(base_url),
|
||||
headers=headers,
|
||||
data=data,
|
||||
files=[(field, (Path(image_path).name, image_file, mime_type))],
|
||||
timeout=300,
|
||||
)
|
||||
if not response.ok:
|
||||
raise http_error(response, "OpenAI image edit")
|
||||
try:
|
||||
return response.json()
|
||||
except ValueError as exc:
|
||||
raise RuntimeError("OpenAI image edit returned invalid JSON.") from exc
|
||||
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Image Generation ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = None) -> str:
|
||||
"""
|
||||
Image generation via OpenAI-compatible API.
|
||||
|
||||
Maps aspect_ratio to OpenAI's size parameter, and image_size to quality.
|
||||
|
||||
Returns:
|
||||
Path of the saved image file
|
||||
|
||||
Raises:
|
||||
RuntimeError: When generation fails
|
||||
"""
|
||||
# Map parameters
|
||||
size_preset = _read_size_preset()
|
||||
size = _select_size(model, aspect_ratio, image_size, size_preset)
|
||||
quality = _read_quality(image_size)
|
||||
output_ext = ".png"
|
||||
request = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"size": size,
|
||||
"n": 1,
|
||||
}
|
||||
if quality is not None:
|
||||
request["quality"] = quality
|
||||
if _is_gpt_image_model(model):
|
||||
gpt_options, output_ext = _gpt_image_options(model)
|
||||
request.update(gpt_options)
|
||||
_apply_response_format(request, model)
|
||||
|
||||
mode_label = f"Proxy: {base_url}" if base_url else "OpenAI API"
|
||||
print(f"[OpenAI - {mode_label}]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Size: {size} (from aspect_ratio={aspect_ratio})")
|
||||
if size_preset and size_preset != "auto":
|
||||
print(f" Size Preset: {size_preset}")
|
||||
if quality is not None:
|
||||
print(f" Quality: {quality} (from image_size={image_size})")
|
||||
else:
|
||||
print(" Quality: omitted")
|
||||
if request.get("response_format"):
|
||||
print(f" Response: {request['response_format']}")
|
||||
elif _read_response_format() == "omit":
|
||||
print(" Response: omitted")
|
||||
if request.get("output_format"):
|
||||
print(f" Format: {request['output_format']}")
|
||||
if request.get("output_compression") is not None:
|
||||
print(f" Compression: {request['output_compression']}")
|
||||
if request.get("background"):
|
||||
print(f" Background: {request['background']}")
|
||||
if request.get("moderation"):
|
||||
print(f" Moderation: {request['moderation']}")
|
||||
print()
|
||||
|
||||
start_time = time.time()
|
||||
print(f" [..] Generating...", end="", flush=True)
|
||||
|
||||
# Heartbeat thread
|
||||
heartbeat_stop = threading.Event()
|
||||
|
||||
def _heartbeat():
|
||||
while not heartbeat_stop.is_set():
|
||||
heartbeat_stop.wait(5)
|
||||
if not heartbeat_stop.is_set():
|
||||
elapsed = time.time() - start_time
|
||||
print(f" {elapsed:.0f}s...", end="", flush=True)
|
||||
|
||||
hb_thread = threading.Thread(target=_heartbeat, daemon=True)
|
||||
hb_thread.start()
|
||||
|
||||
try:
|
||||
resp = _post_image_generation(api_key, base_url, request)
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
hb_thread.join(timeout=1)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n [DONE] Image generated ({elapsed:.1f}s)")
|
||||
|
||||
data = _field(resp, "data") if resp is not None else None
|
||||
if data:
|
||||
path = resolve_output_path(prompt, output_dir, filename, output_ext)
|
||||
first_image = data[0]
|
||||
b64_json = _field(first_image, "b64_json")
|
||||
image_url = _field(first_image, "url")
|
||||
if b64_json:
|
||||
image_data = base64.b64decode(b64_json)
|
||||
return save_image_bytes(image_data, path)
|
||||
if image_url:
|
||||
return download_image(image_url, path)
|
||||
|
||||
raise RuntimeError("No image was generated. The server may have refused the request.")
|
||||
|
||||
|
||||
def _edit_image(api_key: str, prompt: str, reference_image: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = None) -> str:
|
||||
"""
|
||||
Image-to-image edit via the OpenAI-compatible /v1/images/edits endpoint.
|
||||
|
||||
Sends the reference image plus the prompt (used as the edit instruction).
|
||||
Mirrors _generate_image's size/quality/format handling but posts
|
||||
multipart/form-data instead of JSON.
|
||||
|
||||
Returns:
|
||||
Path of the saved image file
|
||||
"""
|
||||
size_preset = _read_size_preset()
|
||||
size = _select_size(model, aspect_ratio, image_size, size_preset)
|
||||
quality = None if _is_dall_e_2(model) else _read_quality(image_size)
|
||||
output_ext = ".png"
|
||||
request = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"size": size,
|
||||
"n": 1,
|
||||
}
|
||||
if quality is not None:
|
||||
request["quality"] = quality
|
||||
if _is_gpt_image_model(model):
|
||||
gpt_options, output_ext = _gpt_image_options(model)
|
||||
request.update(gpt_options)
|
||||
_apply_response_format(request, model)
|
||||
|
||||
mode_label = f"Proxy: {base_url}" if base_url else "OpenAI API"
|
||||
print(f"[OpenAI - {mode_label}]")
|
||||
print(f" Mode: edit (image-to-image)")
|
||||
print(f" Model: {model}")
|
||||
print(f" Reference: {reference_image}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Size: {size} (from aspect_ratio={aspect_ratio})")
|
||||
if size_preset and size_preset != "auto":
|
||||
print(f" Size Preset: {size_preset}")
|
||||
if quality is not None:
|
||||
print(f" Quality: {quality} (from image_size={image_size})")
|
||||
else:
|
||||
print(" Quality: omitted")
|
||||
if request.get("response_format"):
|
||||
print(f" Response: {request['response_format']}")
|
||||
elif _read_response_format() == "omit":
|
||||
print(" Response: omitted")
|
||||
if request.get("output_format"):
|
||||
print(f" Format: {request['output_format']}")
|
||||
if request.get("output_compression") is not None:
|
||||
print(f" Compression: {request['output_compression']}")
|
||||
if request.get("background"):
|
||||
print(f" Background: {request['background']}")
|
||||
if request.get("moderation"):
|
||||
print(f" Moderation: {request['moderation']}")
|
||||
print()
|
||||
|
||||
start_time = time.time()
|
||||
print(f" [..] Editing...", end="", flush=True)
|
||||
|
||||
heartbeat_stop = threading.Event()
|
||||
|
||||
def _heartbeat():
|
||||
while not heartbeat_stop.is_set():
|
||||
heartbeat_stop.wait(5)
|
||||
if not heartbeat_stop.is_set():
|
||||
elapsed = time.time() - start_time
|
||||
print(f" {elapsed:.0f}s...", end="", flush=True)
|
||||
|
||||
hb_thread = threading.Thread(target=_heartbeat, daemon=True)
|
||||
hb_thread.start()
|
||||
|
||||
try:
|
||||
resp = _post_image_edit(api_key, base_url, request, reference_image)
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
hb_thread.join(timeout=1)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n [DONE] Image edited ({elapsed:.1f}s)")
|
||||
|
||||
data = _field(resp, "data") if resp is not None else None
|
||||
if data:
|
||||
path = resolve_output_path(prompt, output_dir, filename, output_ext)
|
||||
first_image = data[0]
|
||||
b64_json = _field(first_image, "b64_json")
|
||||
image_url = _field(first_image, "url")
|
||||
if b64_json:
|
||||
image_data = base64.b64decode(b64_json)
|
||||
return save_image_bytes(image_data, path)
|
||||
if image_url:
|
||||
return download_image(image_url, path)
|
||||
|
||||
raise RuntimeError("No image was returned. The server may have refused the edit request.")
|
||||
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Public Entry Point ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES,
|
||||
reference_image: str = None) -> str:
|
||||
"""
|
||||
OpenAI-compatible image generation (or image-to-image edit) with retry.
|
||||
|
||||
Reads credentials from the current process environment or a `.env` file:
|
||||
OPENAI_API_KEY
|
||||
OPENAI_BASE_URL
|
||||
OPENAI_MODEL (optional override)
|
||||
|
||||
Args:
|
||||
prompt: Prompt text (edit instruction when reference_image is set)
|
||||
aspect_ratio: Aspect ratio, mapped to OpenAI size
|
||||
image_size: Image size, mapped to OpenAI quality
|
||||
output_dir: Output directory
|
||||
filename: Output filename (without extension)
|
||||
model: Model name (default: gpt-image-2)
|
||||
max_retries: Maximum number of retries
|
||||
reference_image: Optional source image path. When set, the request
|
||||
goes to /v1/images/edits (image-to-image) instead of
|
||||
/v1/images/generations.
|
||||
|
||||
Returns:
|
||||
Path of the saved image file
|
||||
"""
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
base_url = os.environ.get("OPENAI_BASE_URL")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"No API key found. Set OPENAI_API_KEY in the current environment or a .env file."
|
||||
)
|
||||
|
||||
if model is None:
|
||||
model = os.environ.get("OPENAI_MODEL") or DEFAULT_MODEL
|
||||
|
||||
image_size = normalize_image_size(image_size)
|
||||
|
||||
if aspect_ratio not in LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE:
|
||||
supported = list(LEGACY_COMPAT_ASPECT_RATIO_TO_SIZE.keys())
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for OpenAI backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
if reference_image is not None:
|
||||
return _edit_image(api_key, prompt, reference_image,
|
||||
aspect_ratio, image_size, output_dir,
|
||||
filename, model, base_url)
|
||||
return _generate_image(api_key, prompt,
|
||||
aspect_ratio, image_size, output_dir,
|
||||
filename, model, base_url)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if attempt < max_retries and is_rate_limit_error(e):
|
||||
delay = retry_delay(attempt, rate_limited=True)
|
||||
print(f"\n [WARN] Rate limit hit (attempt {attempt + 1}/{max_retries + 1}). "
|
||||
f"Waiting {delay}s before retry...")
|
||||
time.sleep(delay)
|
||||
elif attempt < max_retries:
|
||||
delay = retry_delay(attempt, rate_limited=False)
|
||||
print(f"\n [WARN] Error (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
else:
|
||||
break
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,214 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
OpenRouter image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
OPENROUTER_API_KEY (required)
|
||||
OPENROUTER_BASE_URL (optional)
|
||||
OPENROUTER_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend openrouter")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
decode_data_uri,
|
||||
find_data_uri,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
save_image_bytes,
|
||||
)
|
||||
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Constants ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
VALID_ASPECT_RATIOS = [
|
||||
"1:1", "1:4", "1:8",
|
||||
"2:3", "3:2", "3:4", "4:1", "4:3",
|
||||
"4:5", "5:4", "8:1", "9:16", "16:9", "21:9"
|
||||
]
|
||||
|
||||
VALID_IMAGE_SIZES = ["1K", "2K", "4K", "0.5K"]
|
||||
|
||||
DEFAULT_MODEL = "google/gemini-3.1-flash-image-preview"
|
||||
DEFAULT_ENDPOINT = "https://openrouter.ai/api/v1"
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Image Generation ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the OpenRouter generation endpoint."""
|
||||
return base_url.rstrip("/") + "/chat/completions"
|
||||
|
||||
def _message_image_uri(message: dict) -> str | None:
|
||||
"""
|
||||
Locate the generated image in a chat completion message.
|
||||
|
||||
OpenRouter returns it in a dedicated `images` array; other OpenAI-compatible endpoints
|
||||
reachable through OPENROUTER_BASE_URL inline it in the message content instead.
|
||||
"""
|
||||
images = message.get("images")
|
||||
if images:
|
||||
url = images[0].get("image_url")
|
||||
if isinstance(url, dict):
|
||||
url = url.get("url")
|
||||
if url:
|
||||
return url
|
||||
|
||||
return find_data_uri(message.get("content"))
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""
|
||||
Image generation via OpenRouter's API.
|
||||
"""
|
||||
|
||||
url = _resolve_url(base_url)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt
|
||||
}
|
||||
],
|
||||
"modalities": ["image", "text"],
|
||||
"image_config": {
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"image_size": image_size
|
||||
}
|
||||
}
|
||||
|
||||
print(f"[OpenRouter]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Image Size: {image_size}")
|
||||
print()
|
||||
|
||||
start_time = time.time()
|
||||
print(f" [..] Generating...", end="", flush=True)
|
||||
|
||||
# Heartbeat thread
|
||||
heartbeat_stop = threading.Event()
|
||||
|
||||
def _heartbeat():
|
||||
while not heartbeat_stop.is_set():
|
||||
heartbeat_stop.wait(5)
|
||||
if not heartbeat_stop.is_set():
|
||||
elapsed = time.time() - start_time
|
||||
print(f" {elapsed:.0f}s...", end="", flush=True)
|
||||
|
||||
hb_thread = threading.Thread(target=_heartbeat, daemon=True)
|
||||
hb_thread.start()
|
||||
|
||||
try:
|
||||
result = requests.post(url, headers=headers, json=payload, timeout=300).json()
|
||||
finally:
|
||||
heartbeat_stop.set()
|
||||
hb_thread.join(timeout=1)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n [DONE] Image generated ({elapsed:.1f}s)")
|
||||
|
||||
if result.get("choices"):
|
||||
message = result["choices"][0]["message"]
|
||||
image_uri = _message_image_uri(message)
|
||||
if image_uri:
|
||||
image_data, content_type = decode_data_uri(image_uri)
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return save_image_bytes(image_data, path, content_type)
|
||||
|
||||
raise RuntimeError("No image was generated. The server may have refused the request.")
|
||||
|
||||
|
||||
# ╔══════════════════════════════════════════════════════════════════╗
|
||||
# ║ Public Entry Point ║
|
||||
# ╚══════════════════════════════════════════════════════════════════╝
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""
|
||||
OpenRouter image generation with automatic retry.
|
||||
|
||||
Reads credentials from the current process environment or a `.env` file:
|
||||
OPENROUTER_API_KEY
|
||||
OPENROUTER_BASE_URL
|
||||
OPENROUTER_MODEL (optional override)
|
||||
"""
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
base_url = os.environ.get("OPENROUTER_BASE_URL")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"No API key found. Set OPENROUTER_API_KEY in the current environment or a .env file."
|
||||
)
|
||||
|
||||
if model is None:
|
||||
model = os.environ.get("OPENROUTER_MODEL") or DEFAULT_MODEL
|
||||
|
||||
image_size = normalize_image_size(image_size)
|
||||
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(f"Invalid aspect ratio '{aspect_ratio}'. Valid: {VALID_ASPECT_RATIOS}")
|
||||
|
||||
if image_size not in VALID_IMAGE_SIZES:
|
||||
raise ValueError(f"Invalid image size '{image_size}'. Valid: {VALID_IMAGE_SIZES}")
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(api_key, prompt,
|
||||
aspect_ratio, image_size, output_dir,
|
||||
filename, model, base_url)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if attempt < max_retries and is_rate_limit_error(e):
|
||||
delay = retry_delay(attempt, rate_limited=True)
|
||||
print(f"\n [WARN] Rate limit hit (attempt {attempt + 1}/{max_retries + 1}). "
|
||||
f"Waiting {delay}s before retry...")
|
||||
time.sleep(delay)
|
||||
elif attempt < max_retries:
|
||||
delay = retry_delay(attempt, rate_limited=False)
|
||||
print(f"\n [WARN] Error (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
else:
|
||||
break
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,210 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Alibaba Cloud Qwen image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
QWEN_API_KEY / DASHSCOPE_API_KEY (required)
|
||||
QWEN_BASE_URL (optional)
|
||||
QWEN_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend qwen")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||||
DEFAULT_MODEL = "qwen-image-2.0-pro"
|
||||
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": "1024*1024",
|
||||
"2:3": "768*1152",
|
||||
"3:2": "1152*768",
|
||||
"3:4": "864*1152",
|
||||
"4:3": "1152*864",
|
||||
"4:5": "896*1120",
|
||||
"5:4": "1120*896",
|
||||
"9:16": "720*1280",
|
||||
"16:9": "1280*720",
|
||||
"21:9": "1344*576",
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1536*1536",
|
||||
"2:3": "1024*1536",
|
||||
"3:2": "1536*1024",
|
||||
"3:4": "1152*1536",
|
||||
"4:3": "1536*1152",
|
||||
"4:5": "1216*1536",
|
||||
"5:4": "1536*1216",
|
||||
"9:16": "896*1600",
|
||||
"16:9": "1600*896",
|
||||
"21:9": "1792*768",
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "2048*2048",
|
||||
"2:3": "1536*2048",
|
||||
"3:2": "2048*1536",
|
||||
"3:4": "1728*2368",
|
||||
"4:3": "2368*1728",
|
||||
"4:5": "1792*2240",
|
||||
"5:4": "2240*1792",
|
||||
"9:16": "1536*2688",
|
||||
"16:9": "2688*1536",
|
||||
"21:9": "2688*1152",
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "2048*2048",
|
||||
"2:3": "1536*2048",
|
||||
"3:2": "2048*1536",
|
||||
"3:4": "1728*2368",
|
||||
"4:3": "2368*1728",
|
||||
"4:5": "1792*2240",
|
||||
"5:4": "2240*1792",
|
||||
"9:16": "1536*2688",
|
||||
"16:9": "2688*1536",
|
||||
"21:9": "2688*1152",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the Qwen generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/generation"):
|
||||
return base
|
||||
return base + "/api/v1/services/aigc/multimodal-generation/generation"
|
||||
|
||||
|
||||
def _resolve_size(aspect_ratio: str, image_size: str) -> str:
|
||||
"""Resolve the target resolution for a ratio and logical size preset."""
|
||||
normalized = normalize_image_size(image_size)
|
||||
size = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not size:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Qwen backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return size
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the Qwen backend."""
|
||||
size = _resolve_size(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"text": prompt}],
|
||||
}
|
||||
]
|
||||
},
|
||||
"parameters": {
|
||||
"size": size,
|
||||
"prompt_extend": True,
|
||||
"watermark": False,
|
||||
},
|
||||
}
|
||||
|
||||
print("[Alibaba Qwen Image]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "Qwen image generation")
|
||||
|
||||
data = response.json()
|
||||
choices = ((data.get("output") or {}).get("choices") or [])
|
||||
contents = (((choices[0] if choices else {}).get("message") or {}).get("content") or [])
|
||||
image_url = contents[0].get("image") if contents else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"Qwen response missing image URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Qwen backend."""
|
||||
api_key = require_api_key(
|
||||
"QWEN_API_KEY",
|
||||
"DASHSCOPE_API_KEY",
|
||||
message="No API key found. Set QWEN_API_KEY or DASHSCOPE_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("QWEN_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("QWEN_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,182 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Replicate image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
REPLICATE_API_KEY / REPLICATE_API_TOKEN (required)
|
||||
REPLICATE_BASE_URL (optional)
|
||||
REPLICATE_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend replicate")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
poll_json,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
VALID_ASPECT_RATIOS = ["1:1", "16:9", "9:16", "3:4", "4:3", "3:2", "2:3", "4:5", "5:4", "21:9"]
|
||||
DEFAULT_BASE_URL = "https://api.replicate.com/v1"
|
||||
DEFAULT_MODEL = "black-forest-labs/flux-1.1-pro"
|
||||
|
||||
|
||||
def _split_model(model: str) -> tuple[str, str]:
|
||||
"""Split a Replicate model reference into owner and name."""
|
||||
parts = [part for part in model.strip().split("/") if part]
|
||||
if len(parts) != 2:
|
||||
raise ValueError(
|
||||
f"Replicate model must be in 'owner/name' format, got '{model}'."
|
||||
)
|
||||
return parts[0], parts[1]
|
||||
|
||||
|
||||
def _extract_output_url(payload: dict) -> str | None:
|
||||
"""Extract an output URL from a Replicate prediction payload."""
|
||||
output = payload.get("output")
|
||||
if isinstance(output, str):
|
||||
return output
|
||||
if isinstance(output, list) and output:
|
||||
first = output[0]
|
||||
if isinstance(first, str):
|
||||
return first
|
||||
if isinstance(first, dict):
|
||||
return first.get("url")
|
||||
if isinstance(output, dict):
|
||||
return output.get("url")
|
||||
return None
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_BASE_URL) -> str:
|
||||
"""Generate one image with the Replicate backend."""
|
||||
del image_size
|
||||
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Replicate backend. "
|
||||
f"Supported: {VALID_ASPECT_RATIOS}"
|
||||
)
|
||||
|
||||
owner, name = _split_model(model)
|
||||
url = f"{base_url.rstrip('/')}/models/{owner}/{name}/predictions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Prefer": "wait=60",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"input": {
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"output_format": "png",
|
||||
}
|
||||
}
|
||||
|
||||
print("[Replicate]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=180)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Initial response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code not in (200, 201):
|
||||
raise http_error(response, "Replicate generation request")
|
||||
|
||||
data = response.json()
|
||||
status = str(data.get("status", "")).lower()
|
||||
if status != "succeeded":
|
||||
poll_url = ((data.get("urls") or {}).get("get"))
|
||||
if not poll_url:
|
||||
prediction_id = data.get("id")
|
||||
if prediction_id:
|
||||
poll_url = f"{base_url.rstrip('/')}/predictions/{prediction_id}"
|
||||
if not poll_url:
|
||||
raise RuntimeError(f"Replicate response missing poll URL: {data}")
|
||||
|
||||
print(" [..] Polling result...")
|
||||
data = poll_json(
|
||||
poll_url,
|
||||
{"Authorization": f"Bearer {api_key}"},
|
||||
status_label="status",
|
||||
ready_values=["succeeded"],
|
||||
failed_values=["failed", "canceled"],
|
||||
)
|
||||
|
||||
image_url = _extract_output_url(data)
|
||||
if not image_url:
|
||||
raise RuntimeError(f"Replicate response missing output URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Replicate backend."""
|
||||
api_key = require_api_key(
|
||||
"REPLICATE_API_KEY",
|
||||
"REPLICATE_API_TOKEN",
|
||||
message="No API key found. Set REPLICATE_API_KEY or REPLICATE_API_TOKEN in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("REPLICATE_BASE_URL") or DEFAULT_BASE_URL
|
||||
resolved_model = model or os.environ.get("REPLICATE_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,193 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SiliconFlow image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
SILICONFLOW_API_KEY (required)
|
||||
SILICONFLOW_BASE_URL (optional)
|
||||
SILICONFLOW_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend siliconflow")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT = "https://api.siliconflow.cn/v1/images/generations"
|
||||
DEFAULT_MODEL = "Qwen/Qwen-Image"
|
||||
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": "1024x1024",
|
||||
"2:3": "1056x1584",
|
||||
"3:2": "1584x1056",
|
||||
"3:4": "1140x1472",
|
||||
"4:3": "1472x1140",
|
||||
"4:5": "1140x1425",
|
||||
"5:4": "1425x1140",
|
||||
"9:16": "928x1664",
|
||||
"16:9": "1664x928",
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1328x1328",
|
||||
"2:3": "1056x1584",
|
||||
"3:2": "1584x1056",
|
||||
"3:4": "1140x1472",
|
||||
"4:3": "1472x1140",
|
||||
"4:5": "1140x1425",
|
||||
"5:4": "1425x1140",
|
||||
"9:16": "928x1664",
|
||||
"16:9": "1664x928",
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "2048x2048",
|
||||
"2:3": "1536x2048",
|
||||
"3:2": "2048x1536",
|
||||
"3:4": "1536x2048",
|
||||
"4:3": "2048x1536",
|
||||
"4:5": "1638x2048",
|
||||
"5:4": "2048x1638",
|
||||
"9:16": "1152x2048",
|
||||
"16:9": "2048x1152",
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "2048x2048",
|
||||
"2:3": "1536x2048",
|
||||
"3:2": "2048x1536",
|
||||
"3:4": "1536x2048",
|
||||
"4:3": "2048x1536",
|
||||
"4:5": "1638x2048",
|
||||
"5:4": "2048x1638",
|
||||
"9:16": "1152x2048",
|
||||
"16:9": "2048x1152",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the SiliconFlow generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/images/generations"):
|
||||
return base
|
||||
return base + "/v1/images/generations"
|
||||
|
||||
|
||||
def _resolve_size(aspect_ratio: str, image_size: str) -> str:
|
||||
"""Resolve the target resolution for a ratio and logical size preset."""
|
||||
normalized = normalize_image_size(image_size)
|
||||
size = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not size:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for SiliconFlow backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return size
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the SiliconFlow backend."""
|
||||
size = _resolve_size(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"image_size": size,
|
||||
}
|
||||
|
||||
print("[SiliconFlow]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "SiliconFlow image generation")
|
||||
|
||||
data = response.json()
|
||||
images = data.get("images") or []
|
||||
image_url = images[0].get("url") if images else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"SiliconFlow response missing image URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the SiliconFlow backend."""
|
||||
api_key = require_api_key(
|
||||
"SILICONFLOW_API_KEY",
|
||||
message="No API key found. Set SILICONFLOW_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("SILICONFLOW_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("SILICONFLOW_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,156 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Stability AI image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
STABILITY_API_KEY (required)
|
||||
STABILITY_BASE_URL (optional)
|
||||
STABILITY_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend stability")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
report_resolution,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
VALID_ASPECT_RATIOS = [
|
||||
"1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
]
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.stability.ai"
|
||||
DEFAULT_MODEL = "stable-image-core"
|
||||
|
||||
MODEL_ENDPOINTS = {
|
||||
"core": "/v2beta/stable-image/generate/core",
|
||||
"stable-image-core": "/v2beta/stable-image/generate/core",
|
||||
"ultra": "/v2beta/stable-image/generate/ultra",
|
||||
"stable-image-ultra": "/v2beta/stable-image/generate/ultra",
|
||||
}
|
||||
|
||||
|
||||
def _resolve_endpoint(model: str, image_size: str, base_url: str) -> tuple[str, str]:
|
||||
"""Resolve the Stability model alias and endpoint URL."""
|
||||
resolved_model = model or DEFAULT_MODEL
|
||||
if not model and image_size.upper() in ("2K", "4K"):
|
||||
resolved_model = "stable-image-ultra"
|
||||
|
||||
normalized_model = resolved_model.lower()
|
||||
endpoint = MODEL_ENDPOINTS.get(normalized_model)
|
||||
if not endpoint:
|
||||
supported = sorted(MODEL_ENDPOINTS)
|
||||
raise ValueError(
|
||||
f"Unsupported Stability model '{resolved_model}'. Supported aliases: {supported}"
|
||||
)
|
||||
return normalized_model, base_url.rstrip("/") + endpoint
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_BASE_URL) -> str:
|
||||
"""Generate one image with the Stability backend."""
|
||||
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Stability backend. "
|
||||
f"Supported: {VALID_ASPECT_RATIOS}"
|
||||
)
|
||||
|
||||
resolved_model, url = _resolve_endpoint(model, image_size, base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "image/*",
|
||||
}
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"output_format": "png",
|
||||
}
|
||||
|
||||
print("[Stability AI]")
|
||||
print(f" Model: {resolved_model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Preset Size: {image_size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
|
||||
response = requests.post(url, headers=headers, data=data, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "Stability generation")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
with open(path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f" File saved to: {path}")
|
||||
report_resolution(path)
|
||||
return path
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Stability backend."""
|
||||
api_key = require_api_key(
|
||||
"STABILITY_API_KEY",
|
||||
message="No API key found. Set STABILITY_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("STABILITY_BASE_URL") or DEFAULT_BASE_URL
|
||||
resolved_model = model or os.environ.get("STABILITY_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Volcengine Seedream image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
VOLCENGINE_API_KEY / ARK_API_KEY (required)
|
||||
VOLCENGINE_BASE_URL (optional)
|
||||
VOLCENGINE_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend volcengine")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT = "https://operator.las.cn-beijing.volces.com/api/v1/images/generations"
|
||||
DEFAULT_MODEL = "doubao-seedream-4-5-251128"
|
||||
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": "1024x1024",
|
||||
"2:3": "1024x1536",
|
||||
"3:2": "1536x1024",
|
||||
"3:4": "1024x1365",
|
||||
"4:3": "1365x1024",
|
||||
"4:5": "1024x1280",
|
||||
"5:4": "1280x1024",
|
||||
"9:16": "1024x1820",
|
||||
"16:9": "1820x1024",
|
||||
"21:9": "2048x878",
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1536x1536",
|
||||
"2:3": "1344x2016",
|
||||
"3:2": "2016x1344",
|
||||
"3:4": "1440x1920",
|
||||
"4:3": "1920x1440",
|
||||
"4:5": "1536x1920",
|
||||
"5:4": "1920x1536",
|
||||
"9:16": "1152x2048",
|
||||
"16:9": "2048x1152",
|
||||
"21:9": "2048x878",
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "2048x2048",
|
||||
"2:3": "1536x2048",
|
||||
"3:2": "2048x1536",
|
||||
"3:4": "1536x2048",
|
||||
"4:3": "2048x1536",
|
||||
"4:5": "1638x2048",
|
||||
"5:4": "2048x1638",
|
||||
"9:16": "1152x2048",
|
||||
"16:9": "2048x1152",
|
||||
"21:9": "2048x878",
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "2048x2048",
|
||||
"2:3": "1536x2048",
|
||||
"3:2": "2048x1536",
|
||||
"3:4": "1536x2048",
|
||||
"4:3": "2048x1536",
|
||||
"4:5": "1638x2048",
|
||||
"5:4": "2048x1638",
|
||||
"9:16": "1152x2048",
|
||||
"16:9": "2048x1152",
|
||||
"21:9": "2048x878",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the Volcengine generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/images/generations"):
|
||||
return base
|
||||
return base + "/api/v1/images/generations"
|
||||
|
||||
|
||||
def _resolve_size(aspect_ratio: str, image_size: str) -> str:
|
||||
"""Resolve the target resolution for a ratio and logical size preset."""
|
||||
normalized = normalize_image_size(image_size)
|
||||
size = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not size:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Volcengine backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return size
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the Volcengine backend."""
|
||||
size = _resolve_size(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
"response_format": "url",
|
||||
"watermark": False,
|
||||
}
|
||||
|
||||
print("[Volcengine Seedream]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "Volcengine image generation")
|
||||
|
||||
data = response.json()
|
||||
items = data.get("data") or []
|
||||
image_url = items[0].get("url") if items else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"Volcengine response missing image URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".jpeg")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Volcengine backend."""
|
||||
api_key = require_api_key(
|
||||
"VOLCENGINE_API_KEY",
|
||||
"ARK_API_KEY",
|
||||
message="No API key found. Set VOLCENGINE_API_KEY or ARK_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("VOLCENGINE_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("VOLCENGINE_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
@@ -0,0 +1,199 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Zhipu GLM-Image generation backend.
|
||||
|
||||
Configuration keys:
|
||||
ZHIPU_API_KEY / BIGMODEL_API_KEY (required)
|
||||
ZHIPU_BASE_URL (optional)
|
||||
ZHIPU_MODEL (optional)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from console_encoding import configure_utf8_stdio # noqa: E402
|
||||
|
||||
configure_utf8_stdio()
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(__doc__)
|
||||
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend zhipu")
|
||||
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from image_backends.backend_common import (
|
||||
MAX_RETRIES,
|
||||
download_image,
|
||||
http_error,
|
||||
is_rate_limit_error,
|
||||
normalize_image_size,
|
||||
require_api_key,
|
||||
resolve_output_path,
|
||||
retry_delay,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_ENDPOINT = "https://open.bigmodel.cn/api/paas/v4/images/generations"
|
||||
DEFAULT_MODEL = "glm-image"
|
||||
|
||||
ASPECT_RATIO_SIZE_MAP = {
|
||||
"512px": {
|
||||
"1:1": "1024x1024",
|
||||
"2:3": "768x1152",
|
||||
"3:2": "1152x768",
|
||||
"3:4": "864x1152",
|
||||
"4:3": "1152x864",
|
||||
"4:5": "1024x1280",
|
||||
"5:4": "1280x1024",
|
||||
"9:16": "720x1440",
|
||||
"16:9": "1440x720",
|
||||
"21:9": "1536x640",
|
||||
},
|
||||
"1K": {
|
||||
"1:1": "1280x1280",
|
||||
"2:3": "960x1440",
|
||||
"3:2": "1440x960",
|
||||
"3:4": "1024x1365",
|
||||
"4:3": "1365x1024",
|
||||
"4:5": "1024x1280",
|
||||
"5:4": "1280x1024",
|
||||
"9:16": "768x1344",
|
||||
"16:9": "1344x768",
|
||||
"21:9": "1536x640",
|
||||
},
|
||||
"2K": {
|
||||
"1:1": "1440x1440",
|
||||
"2:3": "1152x1728",
|
||||
"3:2": "1728x1152",
|
||||
"3:4": "1152x1536",
|
||||
"4:3": "1536x1152",
|
||||
"4:5": "1280x1600",
|
||||
"5:4": "1600x1280",
|
||||
"9:16": "720x1440",
|
||||
"16:9": "1440x720",
|
||||
"21:9": "1792x768",
|
||||
},
|
||||
"4K": {
|
||||
"1:1": "1440x1440",
|
||||
"2:3": "1152x1728",
|
||||
"3:2": "1728x1152",
|
||||
"3:4": "1152x1536",
|
||||
"4:3": "1536x1152",
|
||||
"4:5": "1280x1600",
|
||||
"5:4": "1600x1280",
|
||||
"9:16": "720x1440",
|
||||
"16:9": "1440x720",
|
||||
"21:9": "1792x768",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _resolve_url(base_url: str) -> str:
|
||||
"""Resolve the Zhipu generation endpoint."""
|
||||
base = base_url.rstrip("/")
|
||||
if base.endswith("/images/generations"):
|
||||
return base
|
||||
return base + "/api/paas/v4/images/generations"
|
||||
|
||||
|
||||
def _resolve_size(aspect_ratio: str, image_size: str) -> str:
|
||||
"""Resolve the target resolution for a ratio and logical size preset."""
|
||||
normalized = normalize_image_size(image_size)
|
||||
size = (ASPECT_RATIO_SIZE_MAP.get(normalized) or {}).get(aspect_ratio)
|
||||
if not size:
|
||||
supported = sorted(ASPECT_RATIO_SIZE_MAP["1K"])
|
||||
raise ValueError(
|
||||
f"Unsupported aspect ratio '{aspect_ratio}' for Zhipu backend. "
|
||||
f"Supported: {supported}"
|
||||
)
|
||||
return size
|
||||
|
||||
|
||||
def _generate_image(api_key: str, prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = DEFAULT_MODEL, base_url: str = DEFAULT_ENDPOINT) -> str:
|
||||
"""Generate one image with the Zhipu backend."""
|
||||
size = _resolve_size(aspect_ratio, image_size)
|
||||
url = _resolve_url(base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
}
|
||||
|
||||
print("[Zhipu GLM-Image]")
|
||||
print(f" Model: {model}")
|
||||
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
||||
print(f" Aspect Ratio: {aspect_ratio}")
|
||||
print(f" Resolution: {size}")
|
||||
print()
|
||||
print(" [..] Generating...", end="", flush=True)
|
||||
start = time.time()
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=300)
|
||||
elapsed = time.time() - start
|
||||
print(f"\n [DONE] Response received ({elapsed:.1f}s)")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise http_error(response, "Zhipu image generation")
|
||||
|
||||
data = response.json()
|
||||
items = data.get("data") or []
|
||||
image_url = items[0].get("url") if items else None
|
||||
if not image_url:
|
||||
raise RuntimeError(f"Zhipu response missing image URL: {data}")
|
||||
|
||||
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
||||
return download_image(image_url, path)
|
||||
|
||||
|
||||
def generate(prompt: str,
|
||||
aspect_ratio: str = "1:1", image_size: str = "1K",
|
||||
output_dir: str = None, filename: str = None,
|
||||
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
||||
"""Generate an image with retries using the Zhipu backend."""
|
||||
api_key = require_api_key(
|
||||
"ZHIPU_API_KEY",
|
||||
"BIGMODEL_API_KEY",
|
||||
message="No API key found. Set ZHIPU_API_KEY or BIGMODEL_API_KEY in the current environment or a .env file.",
|
||||
)
|
||||
base_url = os.environ.get("ZHIPU_BASE_URL") or DEFAULT_ENDPOINT
|
||||
resolved_model = model or os.environ.get("ZHIPU_MODEL") or DEFAULT_MODEL
|
||||
|
||||
last_error = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return _generate_image(
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
output_dir=output_dir,
|
||||
filename=filename,
|
||||
model=resolved_model,
|
||||
base_url=base_url,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_error = exc
|
||||
if attempt >= max_retries:
|
||||
break
|
||||
limited = is_rate_limit_error(exc)
|
||||
delay = retry_delay(attempt, rate_limited=limited)
|
||||
label = "Rate limit hit" if limited else f"Error: {exc}"
|
||||
print(f"\n [WARN] {label}. Retrying in {delay}s...")
|
||||
time.sleep(delay)
|
||||
|
||||
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|
||||
Reference in New Issue
Block a user