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,418 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import threading
|
||||
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_download_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_download_module()
|
||||
|
||||
|
||||
def _load_batch_module():
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.Tensor = type("Tensor", (), {})
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
numpy_module = types.ModuleType("numpy")
|
||||
numpy_module.random = types.SimpleNamespace(seed=lambda _seed: None)
|
||||
sys.modules["numpy"] = numpy_module
|
||||
|
||||
comfy_api = types.ModuleType("comfy_api")
|
||||
comfy_latest = types.ModuleType("comfy_api.latest")
|
||||
comfy_latest.io = types.SimpleNamespace(ComfyNode=object, NodeOutput=object)
|
||||
sys.modules["comfy_api"] = comfy_api
|
||||
sys.modules["comfy_api.latest"] = comfy_latest
|
||||
|
||||
comfy = types.ModuleType("comfy")
|
||||
comfy.__path__ = []
|
||||
comfy_utils = types.ModuleType("comfy.utils")
|
||||
comfy_utils.ProgressBar = object
|
||||
comfy_model_management = types.ModuleType("comfy.model_management")
|
||||
comfy_model_management.processing_interrupted = lambda: False
|
||||
comfy_model_management.InterruptProcessingException = type(
|
||||
"InterruptProcessingException",
|
||||
(RuntimeError,),
|
||||
{},
|
||||
)
|
||||
sys.modules["comfy"] = comfy
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
sys.modules["comfy.model_management"] = comfy_model_management
|
||||
|
||||
folder_paths = types.ModuleType("folder_paths")
|
||||
folder_paths.get_output_directory = lambda: str(ROOT)
|
||||
sys.modules["folder_paths"] = folder_paths
|
||||
|
||||
psutil = types.ModuleType("psutil")
|
||||
psutil.Process = object
|
||||
sys.modules["psutil"] = psutil
|
||||
|
||||
image_utils = types.ModuleType("comfyui_o1key.utils.image_utils")
|
||||
image_utils.tensor_to_pil = lambda _value: []
|
||||
image_utils.pil_to_tensor = lambda value: value
|
||||
image_utils.parse_batch_prompts = lambda _value: []
|
||||
sys.modules[image_utils.__name__] = image_utils
|
||||
|
||||
file_utils = types.ModuleType("comfyui_o1key.utils.file_utils")
|
||||
file_utils.ImageInfo = type("ImageInfo", (), {})
|
||||
for name in (
|
||||
"load_images_from_folder",
|
||||
"pair_images_indexed",
|
||||
"pair_images_by_name",
|
||||
"pair_images_cartesian",
|
||||
"generate_timestamp_filename",
|
||||
"save_image",
|
||||
):
|
||||
setattr(file_utils, name, lambda *_args, **_kwargs: [])
|
||||
sys.modules[file_utils.__name__] = file_utils
|
||||
|
||||
config = types.ModuleType("comfyui_o1key.utils.config")
|
||||
config.NETWORK_ROUTE_OPTIONS = []
|
||||
config.get_base_url_by_route = lambda _route: "https://example.invalid"
|
||||
config.get_api_key_or_raise = lambda _name: "test"
|
||||
sys.modules[config.__name__] = config
|
||||
|
||||
models_config = types.ModuleType("comfyui_o1key.models_config")
|
||||
models_config.get_model_supported_aspect_ratios = lambda _model: []
|
||||
models_config.get_all_supported_aspect_ratios = lambda: []
|
||||
models_config.get_model_supported_resolutions = lambda _model: []
|
||||
models_config.get_all_supported_resolutions = lambda: []
|
||||
sys.modules[models_config.__name__] = models_config
|
||||
|
||||
module_spec = importlib.util.spec_from_file_location(
|
||||
"comfyui_o1key.nodes.batch_nano_banana",
|
||||
ROOT / "nodes" / "batch_nano_banana.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(module_spec)
|
||||
sys.modules[module_spec.name] = module
|
||||
module_spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
BATCH_NODE = _load_batch_module()
|
||||
|
||||
|
||||
def _png_bytes():
|
||||
buffer = io.BytesIO()
|
||||
Image.new("RGB", (2, 2), (12, 34, 56)).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class _FakeContent:
|
||||
def __init__(self, chunks, delay=0):
|
||||
self._chunks = chunks
|
||||
self._delay = delay
|
||||
|
||||
async def iter_chunked(self, _chunk_size):
|
||||
for chunk in self._chunks:
|
||||
if self._delay:
|
||||
await asyncio.sleep(self._delay)
|
||||
yield chunk
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, session, chunks, headers=None, delay=0, status=200):
|
||||
self._session = session
|
||||
self.status = status
|
||||
self.headers = headers or {}
|
||||
self.content = _FakeContent(chunks, delay=delay)
|
||||
|
||||
async def __aenter__(self):
|
||||
self._session.active += 1
|
||||
self._session.max_active = max(self._session.max_active, self._session.active)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||||
self._session.active -= 1
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, chunks, headers=None, delay=0, status=200):
|
||||
self._chunks = chunks
|
||||
self._headers = headers
|
||||
self._delay = delay
|
||||
self._status = status
|
||||
self.calls = []
|
||||
self.active = 0
|
||||
self.max_active = 0
|
||||
|
||||
def get(self, url, **_kwargs):
|
||||
self.calls.append(url)
|
||||
return _FakeResponse(
|
||||
self,
|
||||
self._chunks,
|
||||
headers=self._headers,
|
||||
delay=self._delay,
|
||||
status=self._status,
|
||||
)
|
||||
|
||||
|
||||
class DownloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_successful_task_query_keeps_transport_trace_silent(self):
|
||||
body = json.dumps(
|
||||
{"task_id": "nano-trace", "status": "SUCCESS", "data": {"images": []}},
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
session = _FakeSession(
|
||||
[body],
|
||||
headers={"Content-Length": str(len(body))},
|
||||
)
|
||||
|
||||
with patch("builtins.print") as print_mock:
|
||||
payload = await NANO_ASYNC._poll_task(
|
||||
session,
|
||||
"https://example.invalid",
|
||||
"test-key",
|
||||
"nano-trace",
|
||||
"Nano Banana",
|
||||
log_success=False,
|
||||
initial_delay=False,
|
||||
)
|
||||
|
||||
self.assertEqual(payload["status"], "SUCCESS")
|
||||
rendered = " ".join(str(value) for call in print_mock.call_args_list for value in call.args)
|
||||
self.assertNotIn("任务查询传输追踪", rendered)
|
||||
|
||||
async def test_task_query_retries_when_content_length_is_short(self):
|
||||
body = b'{"task_id":"nano-short","status":"SUCCESS"}'
|
||||
session = _FakeSession(
|
||||
[body],
|
||||
headers={"Content-Length": str(len(body) + 12)},
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
NANO_ASYNC,
|
||||
"_interruptible_sleep",
|
||||
new=AsyncMock(),
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
rf"Content-Length={len(body) + 12}.*received={len(body)}B.*length_check=mismatch",
|
||||
):
|
||||
await NANO_ASYNC._poll_task(
|
||||
session,
|
||||
"https://example.invalid",
|
||||
"test-key",
|
||||
"nano-short",
|
||||
"Nano Banana",
|
||||
log_success=False,
|
||||
initial_delay=False,
|
||||
)
|
||||
|
||||
self.assertEqual(len(session.calls), 4)
|
||||
|
||||
def test_unparseable_response_log_never_prints_partial_base64(self):
|
||||
partial_secret = "A" * 4097
|
||||
with patch("builtins.print") as print_mock:
|
||||
NANO_ASYNC._log_body(
|
||||
"task response",
|
||||
'{"data":{"images":[{"b64_json":"' + partial_secret,
|
||||
)
|
||||
|
||||
rendered = " ".join(str(value) for call in print_mock.call_args_list for value in call.args)
|
||||
self.assertNotIn(partial_secret, rendered)
|
||||
self.assertIn("content omitted", rendered)
|
||||
|
||||
async def test_downloaded_result_is_decoded_only_once(self):
|
||||
image_bytes = _png_bytes()
|
||||
session = _FakeSession([image_bytes])
|
||||
original_open = NANO_ASYNC._open_result_image
|
||||
|
||||
with patch.object(NANO_ASYNC, "_open_result_image", wraps=original_open) as open_image:
|
||||
image = await NANO_ASYNC._image_from_url_or_data(
|
||||
"https://example.invalid/result.png",
|
||||
session,
|
||||
)
|
||||
|
||||
self.assertEqual(image.size, (2, 2))
|
||||
self.assertEqual(open_image.call_count, 1)
|
||||
|
||||
async def test_deduplicates_urls_and_caps_parallel_downloads_at_50(self):
|
||||
image_bytes = _png_bytes()
|
||||
session = _FakeSession(
|
||||
[image_bytes],
|
||||
delay=0.01,
|
||||
)
|
||||
urls = [f"https://example.invalid/{index}.png" for index in range(60)]
|
||||
payload = {
|
||||
"images": urls,
|
||||
"result": {"image_url": urls[0]},
|
||||
}
|
||||
|
||||
images = await NANO_ASYNC._parse_direct_images(
|
||||
payload,
|
||||
session,
|
||||
download_semaphore=asyncio.Semaphore(50),
|
||||
)
|
||||
|
||||
self.assertEqual(len(images), 60)
|
||||
self.assertEqual(len(session.calls), 60)
|
||||
self.assertEqual(session.max_active, 50)
|
||||
|
||||
async def test_retries_when_downloaded_bytes_are_not_a_complete_image(self):
|
||||
session = _FakeSession([b"not a valid image"])
|
||||
original_delays = NANO_ASYNC._DOWNLOAD_RETRY_DELAYS
|
||||
NANO_ASYNC._DOWNLOAD_RETRY_DELAYS = (0, 0)
|
||||
try:
|
||||
with self.assertRaisesRegex(RuntimeError, "transport failed"):
|
||||
await NANO_ASYNC._image_from_url_or_data(
|
||||
"https://example.invalid/truncated.png",
|
||||
session,
|
||||
)
|
||||
finally:
|
||||
NANO_ASYNC._DOWNLOAD_RETRY_DELAYS = original_delays
|
||||
|
||||
self.assertEqual(len(session.calls), 3)
|
||||
|
||||
async def test_interrupt_check_cancels_during_stream(self):
|
||||
session = _FakeSession([b"1234", b"5678"], delay=0.01)
|
||||
checks = 0
|
||||
|
||||
def check_interrupt():
|
||||
nonlocal checks
|
||||
checks += 1
|
||||
if checks >= 2:
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await NANO_ASYNC._download_image_bytes(
|
||||
"https://example.invalid/cancel.png",
|
||||
session,
|
||||
check_interrupt=check_interrupt,
|
||||
)
|
||||
|
||||
async def test_refetches_same_task_when_inline_base64_is_incomplete(self):
|
||||
invalid_payload = {
|
||||
"status": "SUCCESS",
|
||||
"data": {"images": [{"b64_json": "truncated-base64"}]},
|
||||
}
|
||||
valid_payload = {
|
||||
"status": "SUCCESS",
|
||||
"data": {"images": [{
|
||||
"b64_json": base64.b64encode(_png_bytes()).decode("ascii"),
|
||||
}]},
|
||||
}
|
||||
session = object()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"_poll_task",
|
||||
new=AsyncMock(return_value=valid_payload),
|
||||
) as poll_task,
|
||||
patch.object(
|
||||
NANO_ASYNC,
|
||||
"_interruptible_sleep",
|
||||
new=AsyncMock(),
|
||||
) as retry_sleep,
|
||||
):
|
||||
final_payload, parsed = await NANO_ASYNC._parse_completed_task_images_with_retry(
|
||||
invalid_payload,
|
||||
session=session,
|
||||
base_url="https://example.invalid",
|
||||
api_key="test-key",
|
||||
task_id="task-1",
|
||||
node_label="Nano Banana",
|
||||
)
|
||||
|
||||
images, metrics = parsed
|
||||
try:
|
||||
self.assertIs(final_payload, valid_payload)
|
||||
self.assertEqual(len(images), 1)
|
||||
self.assertEqual(images[0].size, (2, 2))
|
||||
self.assertEqual(metrics["inline_images"], 1)
|
||||
retry_sleep.assert_awaited_once()
|
||||
poll_task.assert_awaited_once_with(
|
||||
session,
|
||||
"https://example.invalid",
|
||||
"test-key",
|
||||
"task-1",
|
||||
"Nano Banana",
|
||||
check_interrupt=None,
|
||||
log_body_enabled=False,
|
||||
progress_callback=None,
|
||||
log_success=False,
|
||||
initial_delay=False,
|
||||
)
|
||||
finally:
|
||||
for image in images:
|
||||
image.close()
|
||||
|
||||
|
||||
class BatchCancellationTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_timeout_scales_per_50_task_batch(self):
|
||||
per_batch = BATCH_NODE._PER_BATCH_TIMEOUT_SECONDS
|
||||
grace = BATCH_NODE._BATCH_TIMEOUT_GRACE_SECONDS
|
||||
self.assertEqual(BATCH_NODE._batch_timeout_seconds(1), per_batch + grace)
|
||||
self.assertEqual(BATCH_NODE._batch_timeout_seconds(50), per_batch + grace)
|
||||
self.assertEqual(BATCH_NODE._batch_timeout_seconds(51), per_batch * 2 + grace)
|
||||
|
||||
async def test_stop_event_cancels_active_coroutine(self):
|
||||
stop_event = threading.Event()
|
||||
started = asyncio.Event()
|
||||
cleaned_up = asyncio.Event()
|
||||
|
||||
async def active_work():
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
finally:
|
||||
cleaned_up.set()
|
||||
|
||||
task = asyncio.create_task(
|
||||
BATCH_NODE._run_with_interrupt(active_work(), stop_event)
|
||||
)
|
||||
await started.wait()
|
||||
stop_event.set()
|
||||
|
||||
with self.assertRaises(BATCH_NODE._BatchStopRequested):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
self.assertTrue(cleaned_up.is_set())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user