101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
"""Agnes 原生协议适配器。
|
||
|
||
模型发现： GET {base_url}/models
|
||
出图： POST {base_url}/images/generations
|
||
尺寸： size 档位（1K/2K/3K/4K）+ ratio 宽高比
|
||
返回： data[0].b64_json 或 data[0].url
|
||
注意： response_format 必须放在 extra_body 内（顶层不允许）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
from .base import Capabilities, ImageResult, ProviderAdapter, ProviderError
|
||
|
||
RATIOS = ("1:1", "3:4", "4:3", "16:9", "9:16", "2:3", "3:2", "21:9")
|
||
|
||
# 统一像素尺寸 → Agnes 档位
|
||
_TIER_BY_MAX = (
|
||
(1024, "1K"),
|
||
(1536, "1K"),
|
||
(2048, "2K"),
|
||
(3072, "3K"),
|
||
(4096, "4K"),
|
||
)
|
||
|
||
|
||
def size_to_tier_and_ratio(size: str, ratio: str) -> tuple[str, str]:
|
||
"""把统一的 WxH（可选带 ratio）翻译成 Agnes 的档位 + 宽高比。"""
|
||
m = re.match(r"^\s*(\d+)\s*[xX×]\s*(\d+)\s*$", size or "")
|
||
tier = "1K"
|
||
derived = ""
|
||
if m:
|
||
w, h = int(m.group(1)), int(m.group(2))
|
||
long_side = max(w, h)
|
||
for limit, t in _TIER_BY_MAX:
|
||
if long_side <= limit:
|
||
tier = t
|
||
break
|
||
else:
|
||
tier = "4K"
|
||
derived = _nearest_ratio(w, h)
|
||
elif size and size.strip().upper() in ("1K", "2K", "3K", "4K"):
|
||
tier = size.strip().upper()
|
||
if ratio:
|
||
r = ratio.strip()
|
||
if r in RATIOS:
|
||
return tier, r
|
||
raise ProviderError("不支持的宽高比: " + r, status=400)
|
||
return tier, (derived or "1:1")
|
||
|
||
|
||
def _nearest_ratio(w: int, h: int) -> str:
|
||
target = w / h if h else 1.0
|
||
best, best_diff = "1:1", 1e9
|
||
for r in RATIOS:
|
||
a, b = r.split(":")
|
||
value = int(a) / int(b)
|
||
diff = abs(value - target)
|
||
if diff < best_diff:
|
||
best, best_diff = r, diff
|
||
return best
|
||
|
||
|
||
class AgnesAdapter(ProviderAdapter):
|
||
name = "agnes"
|
||
capabilities = Capabilities(
|
||
supports_image_input=True,
|
||
size_hint="1K~4K / 8 比例",
|
||
kind="image",
|
||
)
|
||
|
||
def generate(
|
||
self,
|
||
model: str,
|
||
prompt: str,
|
||
size: str,
|
||
ratio: str,
|
||
n: int,
|
||
images: list[str],
|
||
) -> list[ImageResult]:
|
||
tier, r = size_to_tier_and_ratio(size, ratio)
|
||
extra: dict = {"response_format": "b64_json"}
|
||
if images:
|
||
extra["image"] = images
|
||
body: dict = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"size": tier,
|
||
"ratio": r,
|
||
"n": n,
|
||
"extra_body": extra,
|
||
}
|
||
resp = self._request("POST", "/images/generations", json_body=body, timeout=600.0)
|
||
payload = self._json(resp)
|
||
results = self.collect_results(payload)
|
||
out: list[ImageResult] = []
|
||
for item in results:
|
||
out.append(self.download(item.url) if item.url else item)
|
||
return out
|