Spaces:
Running on Zero
Running on Zero
File size: 1,459 Bytes
b701455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import asyncio
import contextlib
import pytest
import server
@pytest.mark.asyncio
async def test_img2img_requests_with_different_images_do_not_coalesce(monkeypatch):
buf = server.GenerationBuffer()
calls: list[list[str]] = []
async def fake_process_group(items):
calls.append([item.request_id for item in items])
for item in items:
item.future.set_result({"image": f"data:image/png;base64,{item.request_id}"})
async def noop_prefetch(_signature):
return None
monkeypatch.setattr(buf, "_process_group", fake_process_group)
monkeypatch.setattr(buf, "_look_ahead_and_prefetch", noop_prefetch)
req1 = server.GenerateRequest(
prompt="p1",
img2img_mode=True,
img2img_image="/tmp/input-a.png",
)
req2 = server.GenerateRequest(
prompt="p2",
img2img_mode=True,
img2img_image="/tmp/input-b.png",
)
pending1 = server.PendingRequest(req1, request_id="r1")
pending2 = server.PendingRequest(req2, request_id="r2")
async with buf._lock:
buf._pending.extend([pending1, pending2])
buf._new_request.set()
worker = asyncio.create_task(buf._worker())
try:
await asyncio.wait_for(asyncio.gather(pending1.future, pending2.future), timeout=1.0)
finally:
worker.cancel()
with contextlib.suppress(asyncio.CancelledError):
await worker
assert calls == [["r1"], ["r2"]]
|