diff --git a/worldfoundry/data/benchmarks/assets/wrbench/natural25/README.md b/worldfoundry/data/benchmarks/assets/wrbench/natural25/README.md index 67a2c66b0..c9ab9f65c 100644 --- a/worldfoundry/data/benchmarks/assets/wrbench/natural25/README.md +++ b/worldfoundry/data/benchmarks/assets/wrbench/natural25/README.md @@ -22,6 +22,22 @@ also regenerate first-frame PNGs per family via `wrbench firstframe` (optional extra) or substitute images from your own T2I pipeline. Human annotation verdicts are released separately from this repository. +Atlas Cloud is available as an optional first-frame provider. For example, a +live model whose current schema accepts `size` and PNG output can be run with: + +```bash +wrbench firstframe \ + --family-id demo --prompt "A simple tabletop scene" --out outputs/first-frames \ + --provider atlascloud --model bytedance/seedream-v5.0-lite \ + --api-key "$ATLASCLOUD_API_KEY" \ + --endpoint https://api.atlascloud.ai/api/v1 \ + --size '2048*2048' --n 1 --overwrite-existing +``` + +The provider submits each image task once, polls the bounded prediction +endpoint, and downloads the completed output. Existing bundled frames and +provider defaults are unchanged. + `variants.jsonl` keeps `ti2v_prompt` as the first-frame-anchored prompt of record. Text-only runs should explicitly materialize a prompt profile such as `t2v_layout_anchor`, where the initial layout and event tails are maintained in diff --git a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/cli.py b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/cli.py index 965357587..cd5aafa06 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/cli.py +++ b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/cli.py @@ -1018,7 +1018,7 @@ def _build_parser() -> argparse.ArgumentParser: p_ff.add_argument("--families-jsonl", dest="families_jsonl", metavar="PATH") p_ff.add_argument("--family-id", dest="family_id", metavar="ID") p_ff.add_argument("--prompt", metavar="TEXT", help="T2I prompt (with --family-id).") - p_ff.add_argument("--provider", required=True, help="T2I provider: dashscope, mock.") + p_ff.add_argument("--provider", required=True, help="T2I provider: atlascloud, dashscope, mock.") p_ff.add_argument("--model", required=True, help="T2I model name.") p_ff.add_argument("--api-key", dest="api_key", help="T2I API key.") p_ff.add_argument("--endpoint", required=True, help="T2I API endpoint.") diff --git a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/__init__.py b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/__init__.py index 9e7083ac9..3c8231844 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/__init__.py +++ b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/__init__.py @@ -1,6 +1,7 @@ """First-frame image generation.""" from wrbench.firstframe.generate import ( + AtlasCloudT2IProvider, DashScopeT2IProvider, FirstFrameManifest, MockT2IProvider, @@ -11,6 +12,7 @@ ) __all__ = [ + "AtlasCloudT2IProvider", "DashScopeT2IProvider", "FirstFrameManifest", "MockT2IProvider", diff --git a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/generate.py b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/generate.py index 5317676fa..d3efe8681 100644 --- a/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/generate.py +++ b/worldfoundry/evaluation/tasks/execution/runners/wrbench/runtime/wrbench/firstframe/generate.py @@ -4,6 +4,7 @@ import base64 import json +import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, Protocol @@ -148,6 +149,107 @@ def generate(self, *, prompt: str, family_id: str, out_path: Path) -> dict[str, return {"source": "url", "url": urls[0], "provider": self.provider_name, "model": self.model} +class AtlasCloudT2IProvider: + """Generate first frames through the Atlas Cloud asynchronous image API.""" + + provider_name = "atlascloud" + + def __init__( + self, + *, + model: str | None = None, + api_key: str | None = None, + endpoint: str | None = None, + size: str | None = None, + n: int | str | None = None, + timeout: float = 180.0, + poll_interval: float = 2.0, + ) -> None: + _require_httpx() + import httpx + + self.model = _require_config_value(model, label="Atlas Cloud first-frame model") + self.api_key = _require_config_value(api_key, label="Atlas Cloud API key") + self.endpoint = _require_config_value( + endpoint, + label="Atlas Cloud API endpoint", + ).rstrip("/") + self.size = _require_config_value(size, label="Atlas Cloud first-frame size") + self.n = _require_config_int(n, label="Atlas Cloud first-frame n") + if self.n != 1: + raise RuntimeError("Atlas Cloud first-frame n must be 1") + self.timeout = timeout + self.poll_interval = poll_interval + self._client = httpx.Client(timeout=timeout) + + def _headers(self) -> dict[str, str]: + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + def _submit_url(self) -> str: + if self.endpoint.endswith("/model/generateImage"): + return self.endpoint + return f"{self.endpoint}/model/generateImage" + + def _prediction_url(self, prediction_id: str) -> str: + api_root = self._submit_url().removesuffix("/model/generateImage") + return f"{api_root}/model/prediction/{prediction_id}" + + @staticmethod + def _prediction_data(payload: dict[str, Any]) -> dict[str, Any]: + data = payload.get("data") + return data if isinstance(data, dict) else payload + + def generate(self, *, prompt: str, family_id: str, out_path: Path) -> dict[str, Any]: + del family_id + payload = { + "model": self.model, + "prompt": prompt, + "size": self.size, + "output_format": "png", + } + response = self._client.post(self._submit_url(), headers=self._headers(), json=payload) + response.raise_for_status() + submitted = self._prediction_data(response.json()) + prediction_id = submitted.get("id") + if not isinstance(prediction_id, str) or not prediction_id: + raise RuntimeError("Atlas Cloud submission returned no prediction id") + + deadline = time.monotonic() + self.timeout + while time.monotonic() < deadline: + result = self._client.get( + self._prediction_url(prediction_id), + headers=self._headers(), + ) + result.raise_for_status() + prediction = self._prediction_data(result.json()) + status = str(prediction.get("status") or "").lower() + if status in {"completed", "succeeded"}: + outputs = prediction.get("outputs") or [] + first = outputs[0] if outputs else None + output_url = first if isinstance(first, str) else first.get("url") if isinstance(first, dict) else None + if not output_url: + raise RuntimeError("Atlas Cloud prediction returned no output URL") + image = self._client.get(output_url) + image.raise_for_status() + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_bytes(image.content) + return { + "source": "url", + "url": output_url, + "provider": self.provider_name, + "model": self.model, + "prediction_id": prediction_id, + } + if status in {"failed", "canceled"}: + raise RuntimeError(f"Atlas Cloud prediction {status}: {prediction.get('error') or 'unknown error'}") + time.sleep(self.poll_interval) + + raise TimeoutError(f"Atlas Cloud prediction {prediction_id} timed out") + + class MockT2IProvider: """Write a minimal PNG placeholder for tests (1x1 transparent).""" provider_name = "mock" @@ -173,7 +275,9 @@ def get_t2i_provider(name: str | None = None, **kwargs: Any) -> T2IProvider: return MockT2IProvider(model=model) if provider in {"dashscope", "wan"}: return DashScopeT2IProvider(**kwargs) - raise ValueError(f"Unknown T2I provider {provider!r}; expected dashscope or mock") + if provider in {"atlascloud", "atlas_cloud"}: + return AtlasCloudT2IProvider(**kwargs) + raise ValueError(f"Unknown T2I provider {provider!r}; expected atlascloud, dashscope, or mock") def generate_first_frame( diff --git a/worldfoundry/evaluation/tasks/execution/runners/wrbench/test_atlas_cloud_firstframe.py b/worldfoundry/evaluation/tasks/execution/runners/wrbench/test_atlas_cloud_firstframe.py new file mode 100644 index 000000000..1e9a0a495 --- /dev/null +++ b/worldfoundry/evaluation/tasks/execution/runners/wrbench/test_atlas_cloud_firstframe.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).parent / "runtime")) + +from wrbench.firstframe import AtlasCloudT2IProvider # noqa: E402 + + +class _Response: + def __init__(self, payload=None, content: bytes = b"") -> None: + self._payload = payload + self.content = content + + def raise_for_status(self) -> None: + return None + + def json(self): + return self._payload + + +class _Client: + def __init__(self) -> None: + self.posts = [] + self.gets = [] + + def post(self, url, **kwargs): + self.posts.append((url, kwargs)) + return _Response({"data": {"id": "prediction-1"}}) + + def get(self, url, **kwargs): + self.gets.append((url, kwargs)) + if url.endswith("/prediction/prediction-1"): + return _Response({"data": {"status": "completed", "outputs": ["https://cdn.example/frame.png"]}}) + return _Response(content=b"png-bytes") + + +def test_atlas_cloud_firstframe_submits_once_and_polls(tmp_path): + provider = AtlasCloudT2IProvider( + model="bytedance/seedream-v5.0-lite", + api_key="test-key", + endpoint="https://api.atlascloud.ai/api/v1", + size="2048*2048", + n=1, + poll_interval=0, + ) + client = _Client() + provider._client = client + output = tmp_path / "frame.png" + + metadata = provider.generate(prompt="A clean test frame", family_id="test", out_path=output) + + assert len(client.posts) == 1 + assert client.posts[0][0] == "https://api.atlascloud.ai/api/v1/model/generateImage" + assert client.posts[0][1]["json"] == { + "model": "bytedance/seedream-v5.0-lite", + "prompt": "A clean test frame", + "size": "2048*2048", + "output_format": "png", + } + assert client.gets[0][0] == "https://api.atlascloud.ai/api/v1/model/prediction/prediction-1" + assert output.read_bytes() == b"png-bytes" + assert metadata["provider"] == "atlascloud" + + +def test_atlas_cloud_firstframe_rejects_multiple_outputs(): + with pytest.raises(RuntimeError, match="n must be 1"): + AtlasCloudT2IProvider( + model="bytedance/seedream-v5.0-lite", + api_key="test-key", + endpoint="https://api.atlascloud.ai/api/v1", + size="2048*2048", + n=2, + )