291 lines
11 KiB
Python
291 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Submit a Chorus text-to-image task, poll it, and download its first image."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ipaddress
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import urljoin, urlsplit
|
|
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener
|
|
|
|
|
|
class ChorusError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def origin(url: str) -> tuple[str, str, int | None]:
|
|
parsed = urlsplit(url)
|
|
port = parsed.port
|
|
if port is None:
|
|
port = 443 if parsed.scheme.lower() == "https" else 80 if parsed.scheme.lower() == "http" else None
|
|
return parsed.scheme.lower(), (parsed.hostname or "").lower(), port
|
|
|
|
|
|
class SameOriginRedirectHandler(HTTPRedirectHandler):
|
|
def redirect_request(self, request, file_pointer, code, message, headers, new_url):
|
|
if origin(request.full_url) != origin(new_url):
|
|
raise ChorusError("Chorus attempted a cross-origin redirect")
|
|
return super().redirect_request(request, file_pointer, code, message, headers, new_url)
|
|
|
|
|
|
def open_url(request: Request, timeout: float):
|
|
host = urlsplit(request.full_url).hostname or ""
|
|
try:
|
|
loopback = ipaddress.ip_address(host).is_loopback
|
|
except ValueError:
|
|
loopback = host.lower() == "localhost"
|
|
if loopback:
|
|
return build_opener(ProxyHandler({}), SameOriginRedirectHandler()).open(request, timeout=timeout)
|
|
return build_opener(SameOriginRedirectHandler()).open(request, timeout=timeout)
|
|
|
|
|
|
def read_env(path: Path) -> dict[str, str]:
|
|
if not path.is_file():
|
|
raise ChorusError(f"environment file not found: {path}")
|
|
values: dict[str, str] = {}
|
|
for number, raw_line in enumerate(path.read_text(encoding="utf-8-sig").splitlines(), 1):
|
|
line = raw_line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
if line.lower().startswith("export "):
|
|
line = line[7:].lstrip()
|
|
if "=" not in line:
|
|
raise ChorusError(f"invalid environment entry at line {number}")
|
|
key, value = line.split("=", 1)
|
|
key = key.strip().upper()
|
|
value = value.strip()
|
|
if len(value) >= 2 and value[0] == value[-1] and value[0] in "\"'":
|
|
value = value[1:-1]
|
|
values[key] = value
|
|
return values
|
|
|
|
|
|
def required_setting(values: dict[str, str], names: tuple[str, ...], label: str) -> str:
|
|
for name in names:
|
|
value = values.get(name, "").strip()
|
|
if value:
|
|
return value
|
|
raise ChorusError(f"{label} is missing from chorus.env")
|
|
|
|
|
|
def output_path(args: argparse.Namespace) -> Path:
|
|
if args.output:
|
|
if args.output_dir or args.filename:
|
|
raise ChorusError("use --output or --output-dir with --filename, not both")
|
|
return Path(args.output).expanduser().resolve()
|
|
if not args.output_dir or not args.filename:
|
|
raise ChorusError("provide --output or both --output-dir and --filename")
|
|
if Path(args.filename).name != args.filename:
|
|
raise ChorusError("--filename must not contain directory components")
|
|
return (Path(args.output_dir).expanduser() / args.filename).resolve()
|
|
|
|
|
|
def multipart_body(prompt: str) -> tuple[bytes, str]:
|
|
boundary = f"chorus-{uuid.uuid4().hex}"
|
|
metadata = json.dumps(
|
|
{"capability": "image_generate", "role_rule": "", "images": []},
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
parts: list[bytes] = []
|
|
for name, value, content_type in (
|
|
("prompt", prompt, "text/plain; charset=utf-8"),
|
|
("metadata", metadata, "application/json; charset=utf-8"),
|
|
):
|
|
parts.extend(
|
|
[
|
|
f"--{boundary}\r\n".encode(),
|
|
f'Content-Disposition: form-data; name="{name}"\r\n'.encode(),
|
|
f"Content-Type: {content_type}\r\n\r\n".encode(),
|
|
value.encode("utf-8"),
|
|
b"\r\n",
|
|
]
|
|
)
|
|
parts.append(f"--{boundary}--\r\n".encode())
|
|
return b"".join(parts), boundary
|
|
|
|
|
|
def error_message(payload: bytes, fallback: str) -> str:
|
|
try:
|
|
parsed = json.loads(payload.decode("utf-8"))
|
|
error = parsed.get("error", {}) if isinstance(parsed, dict) else {}
|
|
code = error.get("code")
|
|
message = error.get("message")
|
|
if code and message:
|
|
return f"{code}: {message}"
|
|
if message:
|
|
return str(message)
|
|
except (UnicodeDecodeError, json.JSONDecodeError, AttributeError):
|
|
pass
|
|
return fallback
|
|
|
|
|
|
def request_json(
|
|
url: str,
|
|
api_key: str,
|
|
method: str = "GET",
|
|
body: bytes | None = None,
|
|
headers: dict[str, str] | None = None,
|
|
timeout: float = 30,
|
|
) -> dict:
|
|
request_headers = {"Accept": "application/json", "Authorization": f"Bearer {api_key}"}
|
|
if headers:
|
|
request_headers.update(headers)
|
|
request = Request(url, data=body, headers=request_headers, method=method)
|
|
try:
|
|
with open_url(request, timeout=timeout) as response:
|
|
payload = response.read()
|
|
except HTTPError as exc:
|
|
payload = exc.read()
|
|
raise ChorusError(f"Chorus HTTP {exc.code}: {error_message(payload, exc.reason)}") from None
|
|
except URLError as exc:
|
|
raise ChorusError(f"cannot reach Chorus: {exc.reason}") from None
|
|
try:
|
|
parsed = json.loads(payload.decode("utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise ChorusError("Chorus returned invalid JSON") from exc
|
|
if not isinstance(parsed, dict):
|
|
raise ChorusError("Chorus returned an unexpected response")
|
|
return parsed
|
|
|
|
|
|
def submit(base_url: str, api_key: str, prompt: str, idempotency_key: str) -> dict:
|
|
body, boundary = multipart_body(prompt)
|
|
return request_json(
|
|
f"{base_url}/openapi/v1/generations/image",
|
|
api_key,
|
|
method="POST",
|
|
body=body,
|
|
headers={
|
|
"Content-Type": f"multipart/form-data; boundary={boundary}",
|
|
"Idempotency-Key": idempotency_key,
|
|
},
|
|
)
|
|
|
|
|
|
def wait_for_result(
|
|
base_url: str,
|
|
api_key: str,
|
|
generation_id: int,
|
|
timeout: float,
|
|
poll_interval: float,
|
|
) -> dict:
|
|
deadline = time.monotonic() + timeout
|
|
while True:
|
|
generation = request_json(
|
|
f"{base_url}/openapi/v1/generations/{generation_id}", api_key
|
|
)
|
|
status = generation.get("status")
|
|
if status == "succeeded":
|
|
return generation
|
|
if status == "failed":
|
|
code = generation.get("error_code") or "generation_failed"
|
|
message = generation.get("error_message") or "generation failed"
|
|
raise ChorusError(f"{code}: {message}")
|
|
if status not in {"pending", "running"}:
|
|
raise ChorusError(f"unexpected generation status: {status!r}")
|
|
if time.monotonic() >= deadline:
|
|
raise ChorusError(f"generation {generation_id} did not finish within {timeout:g} seconds")
|
|
time.sleep(min(poll_interval, max(0.0, deadline - time.monotonic())))
|
|
|
|
|
|
def choose_image(generation: dict) -> str:
|
|
outputs = generation.get("outputs")
|
|
if not isinstance(outputs, list):
|
|
raise ChorusError("completed generation has no outputs")
|
|
for output in outputs:
|
|
if isinstance(output, dict) and output.get("kind") == "image" and output.get("url"):
|
|
return str(output["url"])
|
|
raise ChorusError("completed generation has no image output")
|
|
|
|
|
|
def download(url: str, api_key: str, destination: Path, overwrite: bool) -> None:
|
|
if destination.exists() and not overwrite:
|
|
raise ChorusError(f"output already exists: {destination}")
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = destination.with_name(f".{destination.name}.{uuid.uuid4().hex}.part")
|
|
request = Request(url, headers={"Accept": "image/*", "Authorization": f"Bearer {api_key}"})
|
|
try:
|
|
with open_url(request, timeout=60) as response, temporary.open("wb") as output:
|
|
content_type = response.headers.get_content_type()
|
|
if not content_type.startswith("image/"):
|
|
raise ChorusError(f"output is not an image: {content_type}")
|
|
while chunk := response.read(1024 * 1024):
|
|
output.write(chunk)
|
|
os.replace(temporary, destination)
|
|
except HTTPError as exc:
|
|
raise ChorusError(f"image download HTTP {exc.code}: {error_message(exc.read(), exc.reason)}") from None
|
|
except URLError as exc:
|
|
raise ChorusError(f"cannot download image: {exc.reason}") from None
|
|
finally:
|
|
if temporary.exists():
|
|
temporary.unlink()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--env", default="chorus.env", help="path to chorus.env")
|
|
parser.add_argument("--prompt", required=True)
|
|
parser.add_argument("--output")
|
|
parser.add_argument("--output-dir")
|
|
parser.add_argument("--filename")
|
|
parser.add_argument("--timeout", type=float)
|
|
parser.add_argument("--poll-interval", type=float)
|
|
parser.add_argument("--idempotency-key")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
try:
|
|
values = read_env(Path(args.env).expanduser().resolve())
|
|
base_url = required_setting(
|
|
values, ("CHORUS_URL", "CHORUS_API_URL", "URL"), "Chorus URL"
|
|
).rstrip("/")
|
|
api_key = required_setting(
|
|
values,
|
|
("CHORUS_API_KEY", "CHORUS_APIKEY", "APIKEY", "API_KEY"),
|
|
"Chorus API key",
|
|
)
|
|
destination = output_path(args)
|
|
if destination.exists() and not args.overwrite:
|
|
raise ChorusError(f"output already exists: {destination}")
|
|
timeout = args.timeout if args.timeout is not None else float(
|
|
values.get("CHORUS_TIMEOUT_SECONDS", "600")
|
|
)
|
|
poll_interval = args.poll_interval if args.poll_interval is not None else float(
|
|
values.get("CHORUS_POLL_INTERVAL_SECONDS", "2")
|
|
)
|
|
if not args.prompt.strip():
|
|
raise ChorusError("prompt must not be empty")
|
|
if timeout <= 0 or poll_interval <= 0:
|
|
raise ChorusError("timeout and poll interval must be positive")
|
|
idempotency_key = args.idempotency_key or f"chorus-image-{uuid.uuid4().hex}"
|
|
generation = submit(base_url, api_key, args.prompt, idempotency_key)
|
|
generation_id = generation.get("id")
|
|
if not isinstance(generation_id, int) or generation_id <= 0:
|
|
raise ChorusError("submission response has no valid generation ID")
|
|
print(f"Chorus generation {generation_id} submitted", flush=True)
|
|
generation = wait_for_result(base_url, api_key, generation_id, timeout, poll_interval)
|
|
image_url = urljoin(f"{base_url}/", choose_image(generation))
|
|
if origin(image_url) != origin(base_url):
|
|
raise ChorusError("Chorus returned a cross-origin image URL")
|
|
download(image_url, api_key, destination, args.overwrite)
|
|
print(f"Saved Chorus generation {generation_id} to {destination}")
|
|
return 0
|
|
except (ChorusError, OSError, ValueError) as exc:
|
|
print(f"chorus: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|