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
@@ -0,0 +1,389 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_module():
|
||||
package = types.ModuleType("comfyui_o1key")
|
||||
package.__path__ = [str(ROOT)]
|
||||
utils_package = types.ModuleType("comfyui_o1key.utils")
|
||||
utils_package.__path__ = [str(ROOT / "utils")]
|
||||
clients_package = types.ModuleType("comfyui_o1key.clients")
|
||||
clients_package.__path__ = [str(ROOT / "clients")]
|
||||
|
||||
sys.modules[package.__name__] = package
|
||||
sys.modules[utils_package.__name__] = utils_package
|
||||
sys.modules[clients_package.__name__] = clients_package
|
||||
|
||||
http_error_spec = importlib.util.spec_from_file_location(
|
||||
"comfyui_o1key.utils.http_error",
|
||||
ROOT / "utils" / "http_error.py",
|
||||
)
|
||||
http_error = importlib.util.module_from_spec(http_error_spec)
|
||||
sys.modules[http_error_spec.name] = http_error
|
||||
http_error_spec.loader.exec_module(http_error)
|
||||
|
||||
gemini_module = types.ModuleType("comfyui_o1key.clients.gemini_client")
|
||||
gemini_module.GeminiAPIClient = object
|
||||
sys.modules[gemini_module.__name__] = gemini_module
|
||||
|
||||
module_spec = importlib.util.spec_from_file_location(
|
||||
"comfyui_o1key.utils.nano_banana_async",
|
||||
ROOT / "utils" / "nano_banana_async.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(module_spec)
|
||||
sys.modules[module_spec.name] = module
|
||||
module_spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
NANO_ASYNC = _load_module()
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status, payload, headers=None):
|
||||
self.status = status
|
||||
self._text = json.dumps(payload)
|
||||
self.headers = headers or {}
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||||
return False
|
||||
|
||||
async def text(self):
|
||||
return self._text
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
self.calls = []
|
||||
|
||||
def post(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.response
|
||||
|
||||
|
||||
def _noise_image(width=64, height=64):
|
||||
rng = random.Random(7)
|
||||
data = bytes(rng.randrange(256) for _ in range(width * height * 3))
|
||||
return Image.frombytes("RGB", (width, height), data)
|
||||
|
||||
|
||||
def _jpeg_image_bytes(image, quality=95):
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="JPEG", quality=quality)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class TempUploadTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_png_upload_payload_keeps_original_bytes(self):
|
||||
image = _noise_image(16, 12)
|
||||
original = NANO_ASYNC._png_bytes(image)
|
||||
image._o1key_original_format = "PNG"
|
||||
image._o1key_original_bytes = original
|
||||
|
||||
data, extension, content_type = NANO_ASYNC._image_to_upload_payload(image)
|
||||
|
||||
self.assertEqual(data, original)
|
||||
self.assertEqual(extension, ".png")
|
||||
self.assertEqual(content_type, "image/png")
|
||||
|
||||
def test_jpeg_upload_payload_keeps_original_bytes(self):
|
||||
image = _noise_image(32, 24)
|
||||
original = _jpeg_image_bytes(image, quality=95)
|
||||
image._o1key_original_format = "JPEG"
|
||||
image._o1key_original_bytes = original
|
||||
|
||||
data, extension, content_type = NANO_ASYNC._image_to_upload_payload(image)
|
||||
|
||||
self.assertEqual(data, original)
|
||||
self.assertEqual(extension, ".jpg")
|
||||
self.assertEqual(content_type, "image/jpeg")
|
||||
|
||||
async def test_upload_posts_multipart_file_and_returns_https_url(self):
|
||||
expected_url = "https://cf-api.o1key.com/tmp/input/reference.png"
|
||||
response = _FakeResponse(200, {
|
||||
"url": expected_url,
|
||||
"filename": "reference.png",
|
||||
"content_type": "image/png",
|
||||
"size": 100,
|
||||
"expires_at": 1_787_495_062,
|
||||
})
|
||||
session = _FakeSession(response)
|
||||
|
||||
urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
|
||||
session=session,
|
||||
base_url="https://cf-api.o1key.com/",
|
||||
api_key="secret",
|
||||
images=[Image.new("RGB", (2, 2), "red")],
|
||||
)
|
||||
|
||||
self.assertEqual(urls, [expected_url])
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
url, kwargs = session.calls[0]
|
||||
self.assertEqual(url, "https://cf-api.o1key.com/v1/o1key/uploads")
|
||||
self.assertEqual(kwargs["headers"], {"Authorization": "Bearer secret"})
|
||||
self.assertNotIn("Content-Type", kwargs["headers"])
|
||||
filename, field_value, content_type = kwargs["files"]["file"]
|
||||
self.assertEqual(filename, "reference_1.png")
|
||||
self.assertEqual(content_type, "image/png")
|
||||
self.assertTrue(field_value.startswith(b"\x89PNG"))
|
||||
|
||||
async def test_original_jpeg_upload_uses_jpeg_multipart(self):
|
||||
expected_url = "https://cf-api.o1key.com/tmp/input/reference.jpg"
|
||||
response = _FakeResponse(200, {"url": expected_url})
|
||||
session = _FakeSession(response)
|
||||
image = _noise_image(16, 12)
|
||||
image._o1key_original_format = "JPEG"
|
||||
image._o1key_original_bytes = _jpeg_image_bytes(image)
|
||||
|
||||
urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
|
||||
session=session,
|
||||
base_url="https://cf-api.o1key.com/",
|
||||
api_key="secret",
|
||||
images=[image],
|
||||
)
|
||||
|
||||
self.assertEqual(urls, [expected_url])
|
||||
_, kwargs = session.calls[0]
|
||||
filename, field_value, content_type = kwargs["files"]["file"]
|
||||
self.assertEqual(filename, "reference_1.jpg")
|
||||
self.assertEqual(content_type, "image/jpeg")
|
||||
self.assertTrue(field_value.startswith(b"\xff\xd8"))
|
||||
|
||||
async def test_concurrent_tasks_reuse_one_upload(self):
|
||||
image = Image.new("RGB", (2, 2), "blue")
|
||||
cache = {}
|
||||
calls = 0
|
||||
|
||||
async def fake_upload(**_kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
await asyncio.sleep(0)
|
||||
return "https://example.invalid/reference.png"
|
||||
|
||||
with patch.object(
|
||||
NANO_ASYNC,
|
||||
"_upload_nano_banana_temp_image",
|
||||
side_effect=fake_upload,
|
||||
):
|
||||
first, second = await asyncio.gather(
|
||||
NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
|
||||
session=object(),
|
||||
base_url="https://example.invalid",
|
||||
api_key="key",
|
||||
images=[image],
|
||||
upload_cache=cache,
|
||||
),
|
||||
NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
|
||||
session=object(),
|
||||
base_url="https://example.invalid",
|
||||
api_key="key",
|
||||
images=[image],
|
||||
upload_cache=cache,
|
||||
),
|
||||
)
|
||||
|
||||
self.assertEqual(calls, 1)
|
||||
self.assertEqual(first, second)
|
||||
|
||||
async def test_all_reference_uploads_run_concurrently_and_keep_order(self):
|
||||
active = 0
|
||||
max_active = 0
|
||||
|
||||
async def fake_upload(**kwargs):
|
||||
nonlocal active, max_active
|
||||
active += 1
|
||||
max_active = max(max_active, active)
|
||||
await asyncio.sleep(0.01)
|
||||
active -= 1
|
||||
return f"https://example.invalid/{kwargs['filename']}"
|
||||
|
||||
images = [Image.new("RGB", (2, 2), color) for color in ("red", "green", "blue")]
|
||||
with patch.object(
|
||||
NANO_ASYNC,
|
||||
"_upload_nano_banana_temp_image",
|
||||
side_effect=fake_upload,
|
||||
):
|
||||
urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
|
||||
session=object(),
|
||||
base_url="https://example.invalid",
|
||||
api_key="key",
|
||||
images=images,
|
||||
)
|
||||
|
||||
self.assertEqual(max_active, 3)
|
||||
self.assertEqual(urls, [
|
||||
"https://example.invalid/reference_1.png",
|
||||
"https://example.invalid/reference_2.png",
|
||||
"https://example.invalid/reference_3.png",
|
||||
])
|
||||
|
||||
async def test_generation_body_contains_inline_png_without_upload(self):
|
||||
image = Image.new("RGB", (2, 2), "green")
|
||||
result_image = Image.new("RGB", (2, 2), "white")
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"upload_nano_banana_images_to_temp_urls",
|
||||
new=AsyncMock(),
|
||||
) as upload,
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"_submit_task",
|
||||
new=AsyncMock(return_value="task-1"),
|
||||
) as submit,
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"_poll_task",
|
||||
new=AsyncMock(return_value={"status": "success"}),
|
||||
),
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"_parse_task_images",
|
||||
new=AsyncMock(return_value=[result_image]),
|
||||
),
|
||||
):
|
||||
images, _ = await NANO_ASYNC.generate_nano_banana_async(
|
||||
session=object(),
|
||||
base_url="https://example.invalid",
|
||||
api_key="key",
|
||||
prompt="test",
|
||||
model="model",
|
||||
resolution="1K",
|
||||
aspect_ratio="1:1",
|
||||
images=[image],
|
||||
)
|
||||
|
||||
self.assertEqual(images, [result_image])
|
||||
upload.assert_not_awaited()
|
||||
body = submit.await_args.args[3]
|
||||
self.assertEqual(len(body["images"]), 1)
|
||||
inline_data = body["images"][0]["inlineData"]
|
||||
self.assertEqual(inline_data["mimeType"], "image/png")
|
||||
self.assertTrue(base64.b64decode(inline_data["data"]).startswith(b"\x89PNG\r\n\x1a\n"))
|
||||
self.assertNotIn("data:image", inline_data["data"])
|
||||
sanitized = NANO_ASYNC._shorten_base64_for_log(body)
|
||||
self.assertEqual(
|
||||
sanitized["images"][0]["inlineData"]["data"],
|
||||
f"<base64 data, {len(inline_data['data'])} chars>",
|
||||
)
|
||||
|
||||
def test_submit_body_declares_exact_jpeg_mime_type(self):
|
||||
image = _noise_image(16, 12)
|
||||
original = _jpeg_image_bytes(image)
|
||||
image._o1key_original_format = "JPEG"
|
||||
image._o1key_original_bytes = original
|
||||
|
||||
body = NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="1:1",
|
||||
images=[image],
|
||||
)
|
||||
|
||||
inline_data = body["images"][0]["inlineData"]
|
||||
self.assertEqual(inline_data["mimeType"], "image/jpeg")
|
||||
self.assertEqual(base64.b64decode(inline_data["data"]), original)
|
||||
|
||||
def test_submit_body_only_includes_google_search_when_enabled(self):
|
||||
disabled = NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="1:1",
|
||||
)
|
||||
enabled = NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="1:1",
|
||||
google_search=True,
|
||||
)
|
||||
|
||||
self.assertNotIn("google_search", disabled)
|
||||
self.assertIs(enabled["google_search"], True)
|
||||
|
||||
def test_submit_body_omits_size_for_smart_resolution(self):
|
||||
body = NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="智能",
|
||||
aspect_ratio="智能",
|
||||
)
|
||||
|
||||
self.assertNotIn("size", body)
|
||||
|
||||
def test_submit_body_rejects_reference_urls(self):
|
||||
with self.assertRaisesRegex(ValueError, "inlineData"):
|
||||
NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="1:1",
|
||||
image_urls=["https://example.invalid/reference.png"],
|
||||
)
|
||||
|
||||
async def test_no_resize_rejects_the_exact_json_body_above_the_local_limit(self):
|
||||
image = _noise_image(256, 128)
|
||||
with patch.object(NANO_ASYNC, "NANO_BANANA_REQUEST_BODY_LIMIT_BYTES", 20 * 1024):
|
||||
with self.assertRaisesRegex(ValueError, "当前设置为“不缩放”"):
|
||||
await NANO_ASYNC.prepare_nano_banana_inline_images(
|
||||
[image],
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="2:1",
|
||||
resize_mode="不缩放",
|
||||
)
|
||||
|
||||
async def test_smart_resize_preserves_ratio_and_fits_the_exact_json_body(self):
|
||||
image = _noise_image(256, 128)
|
||||
with (
|
||||
patch.object(NANO_ASYNC, "NANO_BANANA_REQUEST_BODY_LIMIT_BYTES", 20 * 1024),
|
||||
patch.object(NANO_ASYNC, "_SMART_RESIZE_MIN_LONG_EDGE", 32),
|
||||
):
|
||||
inline_images = await NANO_ASYNC.prepare_nano_banana_inline_images(
|
||||
[image],
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="2:1",
|
||||
resize_mode="智能缩放",
|
||||
)
|
||||
body = NANO_ASYNC.build_nano_banana_submit_body(
|
||||
model="model",
|
||||
prompt="test",
|
||||
resolution="1K",
|
||||
aspect_ratio="2:1",
|
||||
inline_images=inline_images,
|
||||
)
|
||||
body_size = NANO_ASYNC.validate_nano_banana_request_body(body)
|
||||
|
||||
decoded = base64.b64decode(inline_images[0]["inlineData"]["data"])
|
||||
with Image.open(io.BytesIO(decoded)) as resized:
|
||||
self.assertLess(resized.width, image.width)
|
||||
self.assertAlmostEqual(resized.width / resized.height, 2.0, places=1)
|
||||
self.assertLessEqual(body_size, 20 * 1024)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user