Spaces:
Sleeping
Sleeping
File size: 3,706 Bytes
5c1dcf0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Shared OpenRouter image-generation adapter (dedicated /api/v1/images API).
Same safety posture as ai_client: secrets only from env, bounded retries for
transient failures, telemetry without prompt content, user-safe errors.
LOREIFY_FAKE_AI=1 returns an embedded placeholder image (local testing only).
"""
import base64
import logging
import os
import time
import httpx
from .ai_client import AIUnavailable, USER_ERROR, NO_KEY_ERROR
log = logging.getLogger("lf.image")
# 8x8 solid-gold PNG placeholder for LOREIFY_FAKE_AI local pipeline tests.
_FAKE_PNG = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAgAAAAICAYAAADED76LAAAAFklEQVR4nGP8v4rhPwMRgIkY"
"RaMKKVMIAJVzAx2Vr/xnAAAAAElFTkSuQmCC"
)
def _cfg():
return {
"api_key": os.environ.get("OPENROUTER_API_KEY", ""),
"base_url": os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"),
"model": os.environ.get("OPENROUTER_IMAGE_MODEL", "krea/krea-2-large"),
"timeout": float(os.environ.get("OPENROUTER_IMAGE_TIMEOUT_SECONDS", "120")),
}
def generate_image(prompt: str) -> tuple[bytes, str]:
"""Generate one image. Returns (image_bytes, media_type).
Raises AIUnavailable with a user-safe message on any failure.
"""
if os.environ.get("LOREIFY_FAKE_AI") == "1":
time.sleep(0.4)
return _FAKE_PNG, "image/png"
cfg = _cfg()
if not cfg["api_key"]:
raise AIUnavailable(NO_KEY_ERROR)
headers = {
"Authorization": f"Bearer {cfg['api_key']}",
"HTTP-Referer": "https://loreify.ai",
"X-Title": "Loreify D&D Tools",
"Content-Type": "application/json",
}
# krea models advertise no extra supported_parameters — keep the body minimal.
body = {"model": cfg["model"], "prompt": prompt}
delay = 2.0
last_exc = None
for attempt in range(3):
start = time.time()
try:
resp = httpx.post(f"{cfg['base_url']}/images", headers=headers,
json=body, timeout=cfg["timeout"])
except (httpx.TimeoutException, httpx.TransportError) as exc:
last_exc = exc
log.info("image call attempt=%d error_class=%s", attempt, type(exc).__name__)
time.sleep(delay)
delay *= 2
continue
latency = time.time() - start
if resp.status_code == 429 or resp.status_code >= 500:
last_exc = AIUnavailable(USER_ERROR)
retry_after = resp.headers.get("retry-after")
wait = min(float(retry_after), 20.0) if retry_after and retry_after.isdigit() else delay
log.info("image call attempt=%d status=%d retry_in=%.1f", attempt, resp.status_code, wait)
time.sleep(wait)
delay *= 2
continue
if resp.status_code != 200:
log.info("image call status=%d error_class=http", resp.status_code)
raise AIUnavailable(USER_ERROR)
payload = resp.json()
data = payload.get("data") or []
if not data or not data[0].get("b64_json"):
log.info("image call ok but empty data")
raise AIUnavailable(USER_ERROR)
usage = payload.get("usage") or {}
log.info(
"image ok provider=openrouter model=%s latency=%.2fs tokens=%s cost=%s",
cfg["model"], latency, usage.get("total_tokens"), usage.get("cost"),
)
media_type = data[0].get("media_type") or "image/png"
try:
return base64.b64decode(data[0]["b64_json"]), media_type
except Exception:
raise AIUnavailable(USER_ERROR)
raise last_exc if isinstance(last_exc, AIUnavailable) else AIUnavailable(USER_ERROR)
|