Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
+126
-23
@@ -48,7 +48,7 @@ _RETRY_DELAY = 5
|
||||
|
||||
class GrokImageClient:
|
||||
|
||||
def __init__(self, route: str = "全球加速"):
|
||||
def __init__(self, route: str = None):
|
||||
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
self.base_url = get_base_url_by_route(route)
|
||||
|
||||
@@ -73,8 +73,13 @@ class GrokImageClient:
|
||||
step = 0
|
||||
while len(png_bytes) > max_bytes:
|
||||
scale = 0.894
|
||||
w = max(1, int(w * scale))
|
||||
h = max(1, int(h * scale))
|
||||
next_w = max(1, int(w * scale))
|
||||
next_h = max(1, int(h * scale))
|
||||
if (next_w, next_h) == (w, h):
|
||||
raise RuntimeError(
|
||||
f"Grok Image 无法将图像缩小到 {max_bytes} 字节以内"
|
||||
)
|
||||
w, h = next_w, next_h
|
||||
img = img.resize((w, h), Image.LANCZOS)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
@@ -99,6 +104,116 @@ class GrokImageClient:
|
||||
tensors.append(torch.from_numpy(arr))
|
||||
return torch.stack(tensors, dim=0)
|
||||
|
||||
@classmethod
|
||||
def _build_edit_body(
|
||||
cls,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
image_list: List[torch.Tensor],
|
||||
) -> dict:
|
||||
"""构建最多三张参考图的编辑请求，并确保完整 JSON 不超过 20MB。"""
|
||||
if len(image_list) > 3:
|
||||
raise ValueError("Grok Image 最多支持 3 张参考图")
|
||||
|
||||
reference_images = []
|
||||
for index, tensor in enumerate(image_list, start=1):
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
if not pil_images:
|
||||
raise ValueError(f"无法读取参考图{index}")
|
||||
reference_images.append(pil_images[0])
|
||||
if not reference_images:
|
||||
raise ValueError("图像编辑至少需要 1 张参考图")
|
||||
|
||||
body: dict = {
|
||||
"model": _MODEL_NAME_MAP.get(model, model),
|
||||
"prompt": prompt,
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "auto":
|
||||
body["aspect_ratio"] = aspect_ratio
|
||||
if resolution:
|
||||
body["resolution"] = resolution
|
||||
|
||||
# base64 大约是原始字节的 4/3。预留 1MB 给提示词和 JSON 字段，
|
||||
# 同时保留旧逻辑的单图 PNG 最大 10MB 上限。
|
||||
base_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
available = _MAX_BODY_BYTES - base_size - 1024 * 1024
|
||||
if available <= 0:
|
||||
raise ValueError("提示词和请求参数已超过 Grok Image 20MB 请求体限制")
|
||||
per_image_limit = min(
|
||||
_MAX_BODY_BYTES // 2,
|
||||
max(1, int(available * 0.75) // len(reference_images)),
|
||||
)
|
||||
|
||||
png_images = []
|
||||
for index, image in enumerate(reference_images, start=1):
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
png_images.append(
|
||||
cls._shrink_png_to_limit(
|
||||
buffer.getvalue(),
|
||||
per_image_limit,
|
||||
label=f"参考图{index}",
|
||||
)
|
||||
)
|
||||
|
||||
def _set_images() -> None:
|
||||
encoded = [base64.b64encode(data).decode("ascii") for data in png_images]
|
||||
if len(encoded) == 1:
|
||||
# 单图保持当前 o1key 兼容格式，不改变既有请求行为。
|
||||
body.pop("images", None)
|
||||
body["image"] = encoded[0]
|
||||
else:
|
||||
body.pop("image", None)
|
||||
body["images"] = [
|
||||
{
|
||||
"type": "image_url",
|
||||
"url": f"data:image/png;base64,{value}",
|
||||
}
|
||||
for value in encoded
|
||||
]
|
||||
|
||||
_set_images()
|
||||
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
for _ in range(8):
|
||||
if body_size <= _MAX_BODY_BYTES:
|
||||
return body
|
||||
shrink_ratio = max(0.1, (_MAX_BODY_BYTES / body_size) * 0.95)
|
||||
png_images = [
|
||||
cls._shrink_png_to_limit(
|
||||
data,
|
||||
max(1, int(len(data) * shrink_ratio)),
|
||||
label=f"参考图{index}",
|
||||
)
|
||||
for index, data in enumerate(png_images, start=1)
|
||||
]
|
||||
_set_images()
|
||||
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
|
||||
raise RuntimeError(
|
||||
f"Grok Image 参考图缩放后请求体仍超过 20MB：{body_size / 1024 / 1024:.2f}MB"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _redact_edit_body(body: dict) -> dict:
|
||||
result = dict(body)
|
||||
image = result.get("image")
|
||||
if isinstance(image, str) and len(image) > 50:
|
||||
result["image"] = image[:50] + "..."
|
||||
images = result.get("images")
|
||||
if isinstance(images, list):
|
||||
result["images"] = [
|
||||
{
|
||||
**item,
|
||||
"url": item.get("url", "")[:50] + "...",
|
||||
}
|
||||
if isinstance(item, dict) else item
|
||||
for item in images
|
||||
]
|
||||
return result
|
||||
|
||||
# ── 中断轮询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
@@ -202,28 +317,16 @@ class GrokImageClient:
|
||||
n: int,
|
||||
image_list: List[torch.Tensor],
|
||||
) -> List[Image.Image]:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
body: dict = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "auto":
|
||||
body["aspect_ratio"] = aspect_ratio
|
||||
if resolution:
|
||||
body["resolution"] = resolution
|
||||
|
||||
# 参考图转 base64 字符串
|
||||
pil_images = tensor_to_pil(image_list[0])
|
||||
img = pil_images[0]
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
png_bytes = buf.getvalue()
|
||||
png_bytes = self._shrink_png_to_limit(png_bytes, _MAX_BODY_BYTES // 2)
|
||||
body["image"] = base64.b64encode(png_bytes).decode("utf-8")
|
||||
body = self._build_edit_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
image_list=image_list,
|
||||
)
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_EDITS}"
|
||||
log_body = {k: (v[:50] + "..." if k == "image" and len(v) > 50 else v) for k, v in body.items()}
|
||||
log_body = self._redact_edit_body(body)
|
||||
print(f"[o1key Grok Image] 请求 URL: {url}")
|
||||
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user