Shaman
71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
"""Live tests against the real ComfyUI server (192.168.0.3) via a running adapter.
|
|
|
|
Manual run:
|
|
1. start adapter: uvicorn app.main:app --port 8000
|
|
2. set env: ADAPTER_LIVE=1 ADAPTER_URL=http://127.0.0.1:8000 ADAPTER_KEY=<ADAPTER_API_KEY>
|
|
3. pytest -m live
|
|
"""
|
|
import base64
|
|
import os
|
|
import struct
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
pytestmark = pytest.mark.live
|
|
|
|
URL = os.environ.get("ADAPTER_URL", "http://127.0.0.1:8000")
|
|
KEY = os.environ.get("ADAPTER_KEY", "")
|
|
|
|
pytestmark = [
|
|
pytest.mark.live,
|
|
pytest.mark.skipif(not os.environ.get("ADAPTER_LIVE"), reason="set ADAPTER_LIVE=1 to run live tests"),
|
|
]
|
|
|
|
|
|
def _client() -> httpx.AsyncClient:
|
|
return httpx.AsyncClient(base_url=URL, headers={"Authorization": f"Bearer {KEY}"}, timeout=600)
|
|
|
|
|
|
def _png_size(data: bytes) -> tuple[int, int]:
|
|
return struct.unpack(">II", data[16:24])
|
|
|
|
|
|
def _solid_png(w: int = 128, h: int = 96) -> bytes:
|
|
import zlib
|
|
|
|
def chunk(t, d):
|
|
return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d))
|
|
|
|
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
|
raw = b"".join(b"\x00" + bytes((200, 40, 40)) * w for _ in range(h))
|
|
return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")
|
|
|
|
|
|
async def test_generate_1024_with_seed():
|
|
async with _client() as c:
|
|
r = await c.post("/v1/images/generations", json={"prompt": "a red cube on white background", "size": "1024x1024", "seed": 20261004})
|
|
assert r.status_code == 200, r.text
|
|
img = base64.b64decode(r.json()["data"][0]["b64_json"])
|
|
assert img.startswith(b"\x89PNG")
|
|
assert _png_size(img) == (1024, 1024)
|
|
|
|
|
|
async def test_same_seed_same_image():
|
|
async with _client() as c:
|
|
payload = {"prompt": "a blue sphere", "size": "1024x1024", "seed": 777}
|
|
r1 = await c.post("/v1/images/generations", json=payload)
|
|
r2 = await c.post("/v1/images/generations", json=payload)
|
|
assert r1.json()["data"][0]["b64_json"] == r2.json()["data"][0]["b64_json"]
|
|
|
|
|
|
async def test_edit_output_is_about_1mp_with_input_aspect():
|
|
async with _client() as c:
|
|
files = {"image": ("in.png", _solid_png(128, 96), "image/png")}
|
|
r = await c.post("/v1/images/edits", files=files, data={"prompt": "make the background deep blue"})
|
|
assert r.status_code == 200, r.text
|
|
img = base64.b64decode(r.json()["data"][0]["b64_json"])
|
|
w, h = _png_size(img)
|
|
assert 0.9 <= w * h / 1_000_000 <= 1.15 # ~1 MP
|
|
assert abs(w / h - 128 / 96) < 0.05 # aspect preserved
|