feat: add image-list return format and quality controls
This commit is contained in:
Robert Wojciechowski
@@ -458,7 +458,14 @@ async def coordinator():
|
||||
ALLOW_MASTER = True
|
||||
config['settings']['master_delegate_only'] = False
|
||||
participation_counts = {}
|
||||
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
|
||||
for mode, selected, fmt in [('mixed', [w['id'] for w in workers], 'exact'),
|
||||
('master-only', [], 'exact'),
|
||||
('mixed-png', [w['id'] for w in workers], 'png'),
|
||||
('master-only-png', [], 'png')]:
|
||||
case_prompt = copy.deepcopy(prompt)
|
||||
if fmt != 'exact':
|
||||
case_prompt['4']['inputs'].update(transfer_format=fmt, transfer_quality=73)
|
||||
saved_case = copy.deepcopy(case_prompt)
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
|
||||
@@ -472,8 +479,8 @@ async def coordinator():
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', selected), 45)
|
||||
assert prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
case_prompt, None, 'test', selected), 45)
|
||||
assert case_prompt == saved_case and prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
@@ -499,7 +506,7 @@ async def coordinator():
|
||||
assert '5' not in submission['node_ids']
|
||||
remote_calls.extend(calls)
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
if mode.startswith('mixed'):
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
@@ -507,13 +514,18 @@ async def coordinator():
|
||||
images, masks, texts = CAPTURED[-1]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
|
||||
expected = row['image'] + row['strength']
|
||||
if fmt == 'png' and i in remote_calls:
|
||||
expected = expected.clone()
|
||||
expected[..., :3] = (expected[..., :3].clamp(0, 1) * 255).to(torch.uint8).float() / 255
|
||||
assert torch.equal(images[i], expected), (mode, i)
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
assert len(outer_executions) == 4 and all(row['captures'] == 1 for row in outer_executions)
|
||||
assert len(outer_executions) == 2 + len(participation_counts) and all(row['captures'] == 1 for row in outer_executions)
|
||||
print('IMAGE_LIST_TRANSFER_HTTP_OK mixed-png and master-only-png; masks/text/alpha/local images exact')
|
||||
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""End policy reaches cloud exports and initial worker compatibility preflight."""
|
||||
import asyncio
|
||||
import copy
|
||||
import unittest
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import test_image_list_serverless as cloud
|
||||
import test_image_list_master_selection as selection
|
||||
|
||||
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt', FORMATS)
|
||||
def test_cloud_returns_use_selected_image_format_but_inputs_remain_exact(fmt):
|
||||
async def check():
|
||||
prompt = cloud.source()
|
||||
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=73)
|
||||
original = copy.deepcopy(prompt)
|
||||
value = cloud.item(dtype=torch.float64)
|
||||
job = await cloud.prepare(prompt, value)
|
||||
assert prompt == original
|
||||
for binding in job['manifest']:
|
||||
controls = job['workflow'][binding['node_id']]['inputs']
|
||||
if binding['kind'] == 'IMAGE' and fmt != 'exact':
|
||||
assert controls['transfer_format'] == fmt
|
||||
assert controls['transfer_quality'] == 73
|
||||
else:
|
||||
assert 'transfer_format' not in controls and 'transfer_quality' not in controls
|
||||
for artifact in job['media']['artifacts']:
|
||||
assert artifact['descriptor']['encoding'] == 'tensor-le-v1'
|
||||
inputs = await cloud.decoded_inputs(job)
|
||||
for ref in job['media']['inputs']:
|
||||
assert torch.equal(inputs[ref['artifact_id']], value[ref['input_name']])
|
||||
cloud.handler._validate_request({k: job[k] for k in ('workflow', 'images', 'media')})
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('jpeg',True),('jpeg',101),('jpeg','90')])
|
||||
def test_cloud_rejects_invalid_policy(fmt, quality):
|
||||
prompt = cloud.source()
|
||||
prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=quality)
|
||||
with pytest.raises(ValueError):
|
||||
asyncio.run(cloud.prepare(prompt))
|
||||
|
||||
|
||||
class TransferPreflightTests(unittest.IsolatedAsyncioTestCase):
|
||||
asyncSetUp = selection.MasterSelectionTests.asyncSetUp
|
||||
queue = selection.MasterSelectionTests.queue
|
||||
|
||||
async def test_queue_requests_image_codec_capability(self):
|
||||
self.prompt['3']['inputs'].update(transfer_format='jpeg', transfer_quality=73)
|
||||
await self.queue()
|
||||
self.assertEqual(self.preflight.call_args.kwargs['required_transfer_formats'], ['jpeg'])
|
||||
|
||||
async def test_invalid_policy_fails_before_worker_probe_or_queue(self):
|
||||
self.prompt['3']['inputs'].update(transfer_format='bad-codec')
|
||||
with self.assertRaises(ValueError):
|
||||
await self.queue()
|
||||
self.q.select_active_workers.assert_not_awaited()
|
||||
self.preflight.assert_not_awaited()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_missing_capability_cannot_silently_fall_back_to_master(self):
|
||||
self.prompt['3']['inputs'].update(transfer_format='jpeg')
|
||||
self.preflight.side_effect = RuntimeError('upgrade worker for image return encoding')
|
||||
with self.assertRaisesRegex(RuntimeError, 'upgrade'):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
@@ -0,0 +1,265 @@
|
||||
"""CPU-only direct-worker IMAGE return codecs and capability negotiation."""
|
||||
import asyncio
|
||||
import base64
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from image_list import runtime as rt
|
||||
from serverless_bridge import image_codec
|
||||
from serverless_bridge.image_transfer import FORMATS
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def result(channels=4, dtype=torch.float32):
|
||||
image = torch.linspace(-0.1, 1.1, 17 * 19 * channels, dtype=dtype).reshape(1, 17, 19, channels)
|
||||
return {"image": image, "mask": torch.full((1, 17, 19), 0.123456789, dtype=torch.float64),
|
||||
"text": "exact text: \u00e9\n"}
|
||||
|
||||
|
||||
def legacy_value(value, kind):
|
||||
if isinstance(value, str):
|
||||
return {"type": kind, "value": value}
|
||||
raw = value.contiguous().view(torch.uint8).numpy().tobytes()
|
||||
return {"type": kind, "dtype": str(value.dtype).removeprefix("torch."),
|
||||
"shape": list(value.shape), "data": base64.b64encode(raw).decode("ascii")}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
def test_exact_result_and_inputs_are_byte_for_byte_legacy(dtype):
|
||||
source = result(dtype=dtype)
|
||||
expected = {name: legacy_value(source[name], kind)
|
||||
for name, kind in (("image", "IMAGE"), ("mask", "MASK"), ("text", "STRING"))}
|
||||
body = rt.encode_request_payload(expected)
|
||||
with patch.object(image_codec, "encode_image", side_effect=AssertionError("exact used codec")):
|
||||
assert rt.encode_request_payload(rt.encode_result(source)) == body
|
||||
assert rt.encode_request_payload(rt.encode_result(source, transfer_format="exact", transfer_quality=90)) == body
|
||||
decoded = rt.decode_result(expected)
|
||||
assert torch.equal(decoded["image"].view(torch.uint8), source["image"].view(torch.uint8))
|
||||
item = {"image": source["image"], "mask": source["mask"], "index": 0, "values": {"value_1": source["image"]}}
|
||||
encoded = rt.encode_item(item)
|
||||
assert encoded["image"] == expected["image"]
|
||||
assert encoded["mask"] == expected["mask"]
|
||||
assert encoded["values"]["value_1"] == legacy_value(source["image"], "TENSOR")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", FORMATS[1:])
|
||||
@pytest.mark.parametrize("channels", [3, 4])
|
||||
def test_all_compressed_formats_roundtrip_only_main_image(fmt, channels):
|
||||
source = result(channels)
|
||||
encoded = rt.encode_result(source, transfer_format=fmt, transfer_quality=90)
|
||||
image = encoded["image"]
|
||||
assert set(image) == {"type", "encoding", "descriptor", "data"}
|
||||
assert image["type"] == "IMAGE"
|
||||
assert image["encoding"] == "image-transfer-v1"
|
||||
assert image["descriptor"]["format"] == fmt
|
||||
assert image["descriptor"]["shape"] == [1, 17, 19, channels]
|
||||
assert image["descriptor"]["quality"] == 90
|
||||
assert base64.b64decode(image["data"], validate=True)
|
||||
decoded = rt.decode_result(json.loads(rt.encode_request_payload(encoded)))
|
||||
assert decoded["image"].shape == source["image"].shape
|
||||
assert decoded["image"].dtype == source["image"].dtype
|
||||
assert decoded["image"].device.type == "cpu"
|
||||
assert torch.isfinite(decoded["image"]).all()
|
||||
pixels = source["image"][..., :3].float().mul(255).clamp(0, 255).to(torch.uint8).float().div(255)
|
||||
if fmt in ("legacy_png", "png", "ffv1"):
|
||||
assert torch.equal(decoded["image"][..., :3], pixels)
|
||||
else:
|
||||
assert (decoded["image"][..., :3] - pixels).abs().mean() < 0.08
|
||||
if channels == 4:
|
||||
assert torch.equal(decoded["image"][..., 3], source["image"][..., 3])
|
||||
assert encoded["mask"] == legacy_value(source["mask"], "MASK")
|
||||
assert torch.equal(decoded["mask"], source["mask"])
|
||||
assert decoded["text"] == source["text"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt,quality", [("nope", 90), (["node", 0], 90), ("png", True), ("png", 0), ("exact", 101)])
|
||||
def test_invalid_options_fail_closed(fmt, quality):
|
||||
with pytest.raises(ValueError):
|
||||
rt.encode_result(result(), transfer_format=fmt, transfer_quality=quality)
|
||||
|
||||
|
||||
def test_compressed_inputs_masks_and_generic_tensors_are_rejected():
|
||||
source = result()
|
||||
encoded = rt.encode_result(source, transfer_format="png")
|
||||
for kind in ("IMAGE", "MASK", "TENSOR", "STRING"):
|
||||
with pytest.raises(ValueError):
|
||||
rt.decode_value({**encoded["image"], "type": kind}, kind)
|
||||
item = rt.encode_item({"image": source["image"], "index": 0})
|
||||
with pytest.raises(ValueError):
|
||||
rt.decode_item({**item, "image": encoded["image"]})
|
||||
for key, kind in (("mask", "MASK"), ("text", "STRING")):
|
||||
with pytest.raises(ValueError):
|
||||
rt.decode_result({**encoded, key: {**encoded["image"], "type": kind}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [[2, 17, 19, 3], [1, 17, 19, 1], [17, 19, 3], [True, 17, 19, 3], [1, 0, 19, 3]])
|
||||
def test_invalid_image_shapes_rejected_before_codec_decode(shape):
|
||||
encoded = rt.encode_result(result(), transfer_format="png")
|
||||
encoded["image"]["descriptor"]["shape"] = shape
|
||||
with patch.object(image_codec, "decode_image", side_effect=AssertionError("native codec reached")):
|
||||
with pytest.raises(ValueError):
|
||||
rt.decode_result(encoded)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", [
|
||||
{"encoding": "image-transfer-v2"}, {"type": "MASK"}, {"extra": 1}, {"data": "%%%"},
|
||||
{"descriptor": None}, {"descriptor": {"encoding": "tensor-le-v1", "shape": [1, 17, 19, 4]}},
|
||||
])
|
||||
def test_malformed_compressed_envelopes_rejected(change):
|
||||
encoded = rt.encode_result(result(), transfer_format="png")
|
||||
encoded["image"].update(change)
|
||||
with pytest.raises(ValueError):
|
||||
rt.decode_result(encoded)
|
||||
|
||||
|
||||
def test_single_frame_and_image_list_not_bridge_budgets(monkeypatch):
|
||||
with pytest.raises(ValueError):
|
||||
rt.encode_result({"image": torch.zeros(2, 17, 19, 3)}, transfer_format="png")
|
||||
with patch.object(image_codec, "encode_image", wraps=image_codec.encode_image) as encode:
|
||||
encoded = rt.encode_result(result(), transfer_format="png")
|
||||
assert encode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES
|
||||
with patch.object(image_codec, "decode_image", wraps=image_codec.decode_image) as decode:
|
||||
rt.decode_result(encoded)
|
||||
assert decode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES
|
||||
monkeypatch.setattr(rt, "MAX_PAYLOAD_BYTES", 64)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
rt.encode_result({"image": torch.empty(1, 4, 4, 3, device="meta")}, transfer_format="png")
|
||||
with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")):
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
rt.decode_result(encoded)
|
||||
|
||||
|
||||
def test_compressed_request_combined_budget(monkeypatch):
|
||||
encoded = rt.encode_result(result(), transfer_format="png")
|
||||
monkeypatch.setattr(rt, "MAX_REQUEST_BYTES", len(rt.encode_request_payload(encoded)) - 1)
|
||||
with pytest.raises(ValueError, match="request"):
|
||||
rt.encode_result(result(), transfer_format="png")
|
||||
with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")):
|
||||
with pytest.raises(ValueError, match="request"):
|
||||
rt.decode_result(encoded)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dispatch_module(monkeypatch):
|
||||
name = "_image_list_transfer_tests"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".utils"):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip(".").replace(".", "/"))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, name + ".utils.logging", types.SimpleNamespace(log=lambda *a: None, debug_log=lambda *a: None))
|
||||
monkeypatch.setitem(sys.modules, name + ".utils.network", types.SimpleNamespace(
|
||||
build_worker_url=lambda worker, path="": "http://" + worker["id"] + path,
|
||||
get_client_session=AsyncMock(), probe_worker=AsyncMock()))
|
||||
spec = importlib.util.spec_from_file_location(name + ".api.orchestration.dispatch", ROOT / "api/orchestration/dispatch.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def capability():
|
||||
return {"supported": True, "protocol_version": 2, "execution_targets": True,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"]}
|
||||
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
pass
|
||||
|
||||
async def json(self):
|
||||
return self.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("advertised", [None, [], ["exact"], "png", ["png", 3]])
|
||||
def test_nonexact_requires_advertised_capability_before_queue(dispatch_module, advertised):
|
||||
cap = capability()
|
||||
if advertised is not None:
|
||||
cap["image_transfer_formats"] = advertised
|
||||
session = types.SimpleNamespace(get=Mock(return_value=Response(cap)), post=Mock())
|
||||
dispatch_module.get_client_session = AsyncMock(return_value=session)
|
||||
with pytest.raises(RuntimeError, match="upgrade.*no work was submitted"):
|
||||
asyncio.run(dispatch_module.preflight_image_list_workers(
|
||||
[{"id": "worker"}], required_transfer_formats=["png"]))
|
||||
assert session.get.call_count == 1
|
||||
session.post.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("requested", [(), ("exact",), ("png", "jpeg")])
|
||||
def test_exact_legacy_or_advertised_nonexact_capability_succeeds(dispatch_module, requested):
|
||||
cap = capability()
|
||||
if "png" in requested:
|
||||
cap["image_transfer_formats"] = list(FORMATS)
|
||||
session = types.SimpleNamespace(get=Mock(side_effect=[Response(cap), Response({"queue_running": [], "queue_pending": []})]))
|
||||
dispatch_module.get_client_session = AsyncMock(return_value=session)
|
||||
assert asyncio.run(dispatch_module.preflight_image_list_workers(
|
||||
[{"id": "worker"}], required_transfer_formats=requested)) == {"worker": True}
|
||||
|
||||
|
||||
def test_capability_endpoint_advertises_shared_formats(monkeypatch):
|
||||
from aiohttp import web
|
||||
name = "_transfer_routes"
|
||||
monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
|
||||
monkeypatch.setitem(sys.modules, name + ".api", types.ModuleType(name + ".api"))
|
||||
monkeypatch.setitem(sys.modules, name + ".image_list", types.ModuleType(name + ".image_list"))
|
||||
monkeypatch.setitem(sys.modules, name + ".image_list.runtime", rt)
|
||||
monkeypatch.setitem(sys.modules, "server", types.SimpleNamespace(PromptServer=types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(routes=web.RouteTableDef()))))
|
||||
spec = importlib.util.spec_from_file_location(name + ".api.image_list_routes", ROOT / "api/image_list_routes.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
response = asyncio.run(module.image_list_capabilities_endpoint(None))
|
||||
assert json.loads(response.text)["image_transfer_formats"] == list(FORMATS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reject_call", [None, 1, 2, 3])
|
||||
def test_runtime_initial_and_each_nonexact_assignment_preflight(reject_call):
|
||||
from test_image_list_runtime_boundaries import RuntimeBoundaryTests
|
||||
|
||||
async def exercise():
|
||||
fixture = RuntimeBoundaryTests()
|
||||
await fixture.asyncSetUp()
|
||||
try:
|
||||
context = fixture.create(2)
|
||||
state = fixture.rt.get_master_job(context)
|
||||
state.original_prompt[state.end_id]["inputs"].update(transfer_format="png", transfer_quality=87)
|
||||
calls = []
|
||||
|
||||
async def preflight(workers, **kwargs):
|
||||
calls.append(kwargs)
|
||||
assert kwargs["required_transfer_formats"] == ("png",)
|
||||
if len(calls) == reject_call:
|
||||
raise RuntimeError("upgrade worker; no work was submitted")
|
||||
return {"worker": True}
|
||||
|
||||
sys.modules["_image_list_boundaries.api.orchestration.dispatch"].preflight_image_list_workers = preflight
|
||||
if reject_call:
|
||||
with pytest.raises(RuntimeError, match="upgrade"):
|
||||
await fixture.rt.run_image_list_job(context)
|
||||
assert len(fixture.uploads) == max(0, reject_call - 2)
|
||||
assert len(fixture.prompts) == max(0, reject_call - 2)
|
||||
else:
|
||||
await fixture.rt.run_image_list_job(context)
|
||||
assert len(calls) == 3
|
||||
assert len(fixture.uploads) == 2
|
||||
assert not fixture.rt._MASTER_JOBS
|
||||
finally:
|
||||
# Do not call doCleanups(): this borrowed IsolatedAsyncioTestCase has
|
||||
# no runner, but its sys.modules patch still must be restored.
|
||||
fixture.patch_modules.stop()
|
||||
|
||||
asyncio.run(exercise())
|
||||
Reference in New Issue
Block a user