From dc643074ded929b2d9c4061bd881365d8053a89b Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:25:01 +0800 Subject: [PATCH] feat: add optional Atlas Cloud image generation Signed-off-by: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> --- skills/marketing-image-generation/SKILL.md | 18 ++ .../scripts/atlas_image.py | 283 ++++++++++++++++++ .../scripts/test_atlas_image.py | 90 ++++++ 3 files changed, 391 insertions(+) create mode 100755 skills/marketing-image-generation/scripts/atlas_image.py create mode 100644 skills/marketing-image-generation/scripts/test_atlas_image.py diff --git a/skills/marketing-image-generation/SKILL.md b/skills/marketing-image-generation/SKILL.md index 1d1e7ef..b986c71 100644 --- a/skills/marketing-image-generation/SKILL.md +++ b/skills/marketing-image-generation/SKILL.md @@ -14,6 +14,24 @@ description: Generate high-quality marketing images, ad creatives, launch visual 5. Use the image generation tool only when the user has explicitly asked for an image asset or visual mockup. 6. Review the result against the brief, then iterate with targeted changes instead of rewriting the whole concept. +## Optional Atlas Cloud Provider + +When the user explicitly requests Atlas Cloud, or no native image-generation tool is available, use the bundled provider helper. Atlas Cloud remains optional; do not replace a working native image tool without the user's request. + +Set the API key in the environment, discover a current image model, then generate: + +```bash +export ATLASCLOUD_API_KEY="your-api-key" +python3 scripts/atlas_image.py list --search image +python3 scripts/atlas_image.py generate \ + --model "MODEL_ID_FROM_LIST" \ + --prompt "Create a product launch image with a single focal point" \ + --size 1024x1024 \ + --output launch.png +``` + +The helper reads the live Atlas model catalog and the selected model's schema before submission. It submits generation exactly once, polls the result with a bounded loop, and can either print the output URL or download it with `--output`. Only pass `--size`, `--quality`, or `--output-format` values supported by the live schema. + ## Marketing Defaults - Prefer one strong focal point over collages. diff --git a/skills/marketing-image-generation/scripts/atlas_image.py b/skills/marketing-image-generation/scripts/atlas_image.py new file mode 100755 index 0000000..990d27e --- /dev/null +++ b/skills/marketing-image-generation/scripts/atlas_image.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +"""Generate marketing images through the optional Atlas Cloud provider.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + + +DEFAULT_BASE_URL = "https://api.atlascloud.ai" + + +class AtlasError(RuntimeError): + def __init__(self, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.retryable = retryable + + +def request_json( + method: str, + url: str, + *, + api_key: str | None = None, + payload: dict[str, Any] | None = None, + timeout: float = 30, +) -> dict[str, Any]: + headers = { + "Accept": "application/json", + "User-Agent": "atlas-cloud-marketing-image-skill/1.0", + } + data = None + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + if payload is not None: + headers["Content-Type"] = "application/json" + data = json.dumps(payload).encode("utf-8") + + request = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + body = response.read().decode("utf-8") + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace") + raise AtlasError( + f"Atlas HTTP {exc.code}: {detail}", + retryable=exc.code == 429 or exc.code >= 500, + ) from exc + except urllib.error.URLError as exc: + raise AtlasError(f"Atlas request failed: {exc.reason}", retryable=True) from exc + + try: + result = json.loads(body) + except json.JSONDecodeError as exc: + raise AtlasError("Atlas returned invalid JSON") from exc + if not isinstance(result, dict): + raise AtlasError("Atlas returned an unexpected response shape") + return result + + +def unwrap(payload: dict[str, Any]) -> dict[str, Any]: + code = payload.get("code") + if code is not None and str(code) not in {"0", "200"}: + raise AtlasError(str(payload.get("message") or f"Atlas API error {code}")) + data = payload.get("data", payload) + if not isinstance(data, dict): + raise AtlasError("Atlas response data is not an object") + return data + + +def fetch_models(base_url: str) -> list[dict[str, Any]]: + payload = request_json("GET", f"{base_url}/api/v1/models") + code = payload.get("code") + if code is not None and str(code) not in {"0", "200"}: + raise AtlasError(str(payload.get("message") or f"Atlas API error {code}")) + models = payload.get("data") + if not isinstance(models, list): + raise AtlasError("Atlas model catalog is not a list") + return [model for model in models if isinstance(model, dict)] + + +def visible_image_models(models: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + model + for model in models + if (model.get("display_console") is True or model.get("displayConsole") is True) + and str(model.get("type", "")).lower() == "image" + ] + + +def find_model(models: list[dict[str, Any]], model_id: str) -> dict[str, Any]: + for model in visible_image_models(models): + if model.get("model") == model_id: + return model + raise AtlasError(f"Model is not a visible Atlas image model: {model_id}") + + +def fetch_input_schema(model: dict[str, Any]) -> dict[str, Any]: + schema_url = model.get("schema") + if not isinstance(schema_url, str) or not schema_url.startswith("https://"): + raise AtlasError("Selected model does not expose an HTTPS schema URL") + document = request_json("GET", schema_url) + try: + schema = document["components"]["schemas"]["Input"] + except (KeyError, TypeError) as exc: + raise AtlasError("Selected model schema has no components.schemas.Input") from exc + if not isinstance(schema, dict): + raise AtlasError("Selected model input schema is invalid") + return schema + + +def build_payload(args: argparse.Namespace, schema: dict[str, Any]) -> dict[str, Any]: + properties = schema.get("properties", {}) + if not isinstance(properties, dict) or "prompt" not in properties: + raise AtlasError("Selected model schema does not support prompt") + + payload: dict[str, Any] = {"model": args.model, "prompt": args.prompt} + optional = { + "size": args.size, + "quality": args.quality, + "output_format": args.output_format, + } + for name, value in optional.items(): + if value is None: + continue + if name not in properties: + raise AtlasError(f"Selected model schema does not support {name}") + allowed = properties[name].get("enum") + if isinstance(allowed, list) and value not in allowed: + raise AtlasError(f"Invalid {name}: {value}. Allowed: {', '.join(map(str, allowed))}") + payload[name] = value + if "enable_sync_mode" in properties: + payload["enable_sync_mode"] = False + if "enable_base64_output" in properties: + payload["enable_base64_output"] = False + return payload + + +def submit_generation(base_url: str, api_key: str, payload: dict[str, Any]) -> dict[str, Any]: + # Intentionally one POST only. A failed submission is never retried automatically. + response = request_json( + "POST", + f"{base_url}/api/v1/model/generateImage", + api_key=api_key, + payload=payload, + timeout=60, + ) + return unwrap(response) + + +def poll_generation( + base_url: str, + api_key: str, + prediction_id: str, + *, + attempts: int, + interval: float, +) -> dict[str, Any]: + for attempt in range(attempts): + try: + prediction = unwrap( + request_json( + "GET", + f"{base_url}/api/v1/model/result/{prediction_id}", + api_key=api_key, + ) + ) + except AtlasError as exc: + if not exc.retryable or attempt + 1 >= attempts: + raise + time.sleep(min(interval * (2**attempt), 8)) + continue + + status = str(prediction.get("status", "")).lower() + if status in {"completed", "succeeded"}: + return prediction + if status in {"failed", "canceled", "cancelled"}: + raise AtlasError(str(prediction.get("error") or f"Prediction {status}")) + if attempt + 1 < attempts: + time.sleep(interval) + raise AtlasError(f"Prediction did not complete after {attempts} polls") + + +def extract_output_url(prediction: dict[str, Any]) -> str: + outputs = prediction.get("outputs") + if isinstance(outputs, list) and outputs and isinstance(outputs[0], str): + return outputs[0] + output = prediction.get("output") + if isinstance(output, str): + return output + raise AtlasError("Completed prediction has no output URL") + + +def download(url: str, output: Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + try: + with urllib.request.urlopen(url, timeout=60) as response: + output.write_bytes(response.read()) + except urllib.error.URLError as exc: + raise AtlasError(f"Could not download generated image: {exc.reason}") from exc + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--base-url", + default=os.environ.get("ATLASCLOUD_BASE_URL", DEFAULT_BASE_URL), + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + list_parser = subparsers.add_parser("list", help="List current visible image models") + list_parser.add_argument("--search", default="") + + generate = subparsers.add_parser("generate", help="Generate one image") + generate.add_argument("--model", required=True) + generate.add_argument("--prompt", required=True) + generate.add_argument("--size") + generate.add_argument("--quality") + generate.add_argument("--output-format") + generate.add_argument("--output", type=Path) + generate.add_argument("--poll-attempts", type=int, default=100) + generate.add_argument("--poll-interval", type=float, default=3) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + base_url = args.base_url.rstrip("/") + try: + models = fetch_models(base_url) + if args.command == "list": + query = args.search.lower() + for model in visible_image_models(models): + model_id = str(model.get("model", "")) + display_name = str(model.get("displayName", "")) + if query and query not in f"{model_id} {display_name}".lower(): + continue + print(f"{model_id}\t{display_name}") + return 0 + + if args.poll_attempts < 1 or args.poll_interval < 0: + raise AtlasError("Polling limits must be positive") + api_key = os.environ.get("ATLASCLOUD_API_KEY") + if not api_key: + raise AtlasError("ATLASCLOUD_API_KEY is required for generation") + + model = find_model(models, args.model) + schema = fetch_input_schema(model) + payload = build_payload(args, schema) + prediction = submit_generation(base_url, api_key, payload) + status = str(prediction.get("status", "")).lower() + if status not in {"completed", "succeeded"}: + prediction_id = prediction.get("id") + if not isinstance(prediction_id, str) or not prediction_id: + raise AtlasError("Generation response has no prediction ID") + prediction = poll_generation( + base_url, + api_key, + prediction_id, + attempts=args.poll_attempts, + interval=args.poll_interval, + ) + + output_url = extract_output_url(prediction) + if args.output: + download(output_url, args.output) + print(args.output) + else: + print(output_url) + return 0 + except AtlasError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/marketing-image-generation/scripts/test_atlas_image.py b/skills/marketing-image-generation/scripts/test_atlas_image.py new file mode 100644 index 0000000..d7ee0b0 --- /dev/null +++ b/skills/marketing-image-generation/scripts/test_atlas_image.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import argparse +import importlib.util +import unittest +from pathlib import Path +from unittest.mock import patch + + +MODULE_PATH = Path(__file__).with_name("atlas_image.py") +SPEC = importlib.util.spec_from_file_location("atlas_image", MODULE_PATH) +assert SPEC and SPEC.loader +atlas_image = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(atlas_image) + + +class AtlasImageTest(unittest.TestCase): + def test_filters_visible_image_models(self) -> None: + models = [ + {"model": "visible", "type": "Image", "display_console": True}, + {"model": "hidden", "type": "Image", "display_console": False}, + {"model": "video", "type": "Video", "display_console": True}, + ] + self.assertEqual(["visible"], [m["model"] for m in atlas_image.visible_image_models(models)]) + + def test_build_payload_uses_only_schema_supported_options(self) -> None: + args = argparse.Namespace( + model="provider/model", + prompt="A launch image", + size="1024x1024", + quality="high", + output_format="png", + ) + schema = { + "properties": { + "prompt": {"type": "string"}, + "size": {"enum": ["1024x1024"]}, + "quality": {"enum": ["medium", "high"]}, + "output_format": {"enum": ["jpeg", "png"]}, + "enable_sync_mode": {"type": "boolean"}, + } + } + self.assertEqual( + { + "model": "provider/model", + "prompt": "A launch image", + "size": "1024x1024", + "quality": "high", + "output_format": "png", + "enable_sync_mode": False, + }, + atlas_image.build_payload(args, schema), + ) + + def test_submit_generation_makes_one_post(self) -> None: + with patch.object( + atlas_image, + "request_json", + return_value={"code": 200, "data": {"id": "prediction-1", "status": "starting"}}, + ) as request: + result = atlas_image.submit_generation("https://example.test", "secret", {"prompt": "x"}) + self.assertEqual("prediction-1", result["id"]) + self.assertEqual(1, request.call_count) + self.assertEqual("POST", request.call_args.args[0]) + + def test_polling_is_bounded_and_returns_completed_prediction(self) -> None: + responses = [ + {"code": 200, "data": {"id": "prediction-1", "status": "processing"}}, + { + "code": 200, + "data": { + "id": "prediction-1", + "status": "completed", + "outputs": ["https://cdn.example.test/image.png"], + }, + }, + ] + with ( + patch.object(atlas_image, "request_json", side_effect=responses) as request, + patch.object(atlas_image.time, "sleep"), + ): + result = atlas_image.poll_generation( + "https://example.test", "secret", "prediction-1", attempts=3, interval=0 + ) + self.assertEqual("completed", result["status"]) + self.assertEqual(2, request.call_count) + + +if __name__ == "__main__": + unittest.main()