Spaces:
Running
Running
File size: 2,092 Bytes
c47ec10 | 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 53 54 55 56 57 58 59 | from __future__ import annotations
import json
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from uuid import uuid4
BASE_URL = "http://127.0.0.1:8000"
PROMPT = "按原图位置拆分海报元素并合成可编辑 PSD，同时输出每个图层素材 zip。"
BASE64_IMAGES: list[str] = []
TIMEOUT_SECS = 600
POLL_INTERVAL_SECS = 5
def request_json(method: str, path: str, payload: dict | None = None) -> dict:
api_key = json.loads((Path(__file__).resolve().parents[1] / "config.json").read_text(encoding="utf-8"))["auth-key"]
if not api_key.strip():
raise ValueError("API_KEY is empty")
data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode()
request = urllib.request.Request(
BASE_URL.rstrip("/") + path,
data=data,
method=method,
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
)
try:
with urllib.request.urlopen(request, timeout=60) as response:
return json.loads(response.read().decode())
except urllib.error.HTTPError as exc:
raise RuntimeError(exc.read().decode("utf-8", "replace")) from exc
def main() -> None:
task = request_json("POST", "/v1/psd/generations", {
"client_task_id": str(uuid4()),
"prompt": PROMPT,
"base64_images": BASE64_IMAGES,
})
task_id = str(task.get("taskId") or task.get("id") or "")
if not task_id:
raise RuntimeError(f"missing taskId: {task}")
print(json.dumps(task, ensure_ascii=False, indent=2))
deadline = time.time() + TIMEOUT_SECS
while time.time() < deadline:
time.sleep(POLL_INTERVAL_SECS)
status = request_json("GET", "/v1/editable-file-tasks?ids=" + urllib.parse.quote(task_id))
print(json.dumps(status, ensure_ascii=False, indent=2))
item = (status.get("items") or [{}])[0]
if item.get("status") in {"success", "error"}:
return
raise TimeoutError(f"task timeout: {task_id}")
if __name__ == "__main__":
main()
|