"""MJ image jobs. pip install requests pillow. Never resubmit timed-out jobs.""" import argparse import hashlib import io import json import math import os import time from pathlib import Path from urllib.parse import quote import requests from PIL import Image def api(method, path, key, body=None, read_timeout=90): response = requests.request(method, "https://www.yunqiai.chat" + path, headers={"Authorization": "Bearer " + key, "Accept": "application/json"}, json=body, allow_redirects=False, timeout=(min(15, read_timeout), read_timeout)) if not 200 <= response.status_code < 300: error = RuntimeError("HTTP " + str(response.status_code) + ": " + response.text[:4000]) error.status_code = response.status_code error.retry_after = response.headers.get("Retry-After") raise error data = response.json() if data.get("error"): raise RuntimeError("API error: " + json.dumps(data["error"], ensure_ascii=False)) return data def wait_job(job_id, key, output, seconds=360): if not isinstance(seconds, (int, float)) or not math.isfinite(seconds) or seconds <= 0: raise ValueError("Wait duration must be positive and finite") deadline = time.monotonic() + seconds temporary_errors = 0 while time.monotonic() < deadline: try: data = api("GET", "/v1/tob/job/" + quote(job_id, safe=""), key, read_timeout=min(90, max(0.01, deadline - time.monotonic()))) temporary_errors = 0 except (RuntimeError, requests.exceptions.Timeout, requests.exceptions.ConnectionError) as error: if isinstance(error, requests.exceptions.SSLError): raise network_error = isinstance(error, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)) if not network_error and getattr(error, "status_code", 0) not in [429, 500, 502, 503, 504]: raise temporary_errors += 1 if temporary_errors > 6: raise delay = getattr(error, "retry_after", None) pause = min(30, max(1, int(delay))) if delay and delay.isdigit() else min(30, 5 * temporary_errors) time.sleep(min(pause, max(0, deadline - time.monotonic()))) continue (output / "job.json").write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") if data.get("status") == 2: return data if data.get("status") == 3: raise RuntimeError("Job failed; see job.json for comment and response") if data.get("status") != 1: raise RuntimeError("Unknown job status; see job.json") time.sleep(min(5, max(0, deadline - time.monotonic()))) raise TimeoutError("Local wait expired. Resume using --job " + job_id + "; do not submit again") def download_all(data, job_id, output): results = [] for index, url in enumerate(data.get("urls", [])): try: if not isinstance(url, str) or not url: raise ValueError("Empty image URL") response = requests.get(url, timeout=(15, 180)) response.raise_for_status() with Image.open(io.BytesIO(response.content)) as image: image.load() target = output / ("mj-imagine-" + job_id + "-" + str(index) + ".png") image.save(target, "PNG") results.append({"imageNo": index, "file": target.name, "size": list(image.size)}) except Exception as error: results.append({"imageNo": index, "error": str(error)}) (output / "downloads.json").write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8") if not results or any("error" in row for row in results): raise RuntimeError("Some images were not downloaded; resume the original job, see downloads.json") return results def main(): parser = argparse.ArgumentParser() parser.add_argument("--text") parser.add_argument("--job", help="Resume an existing job without resubmission") parser.add_argument("--action", default="diffusion", choices=["diffusion", "variation", "upscale", "reroll", "pan", "outpaint", "inpaint", "remix", "edit", "upload-paint", "retexture", "remove-background", "enhance"]) parser.add_argument("--body", help="JSON request body for a secondary action") parser.add_argument("--out", default="mj-results") parser.add_argument("--wait", type=int, default=360, help="Local polling duration in seconds; does not cancel the remote job") args = parser.parse_args() if args.wait <= 0: raise ValueError("--wait must be positive") key = os.environ["YUNQIAI_API_KEY"] output = Path(args.out) output.mkdir(parents=True, exist_ok=True) if args.job: job_id = args.job else: if (output / "submitted.json").exists() or (output / "submission-attempt.json").exists(): raise RuntimeError("This directory already contains a submission attempt. Resume its known job; if no id was received, confirm its state before any new submission") if args.action == "diffusion": model = os.environ["YUNQIAI_MODEL"] if model not in ["mj_imagine", "youchuan-image"] or not args.text: raise ValueError("Choose an exact MJ model and provide --text") body = {"model": model, "text": args.text} else: if not args.body: raise ValueError("Provide --body for this secondary action") body = json.loads(Path(args.body).read_text(encoding="utf-8")) route = "/v1/tob/" + args.action # An interrupted POST may already have created a billable remote task. with (output / "submission-attempt.json").open("x", encoding="utf-8") as journal: json.dump({"path": route, "request_sha256": hashlib.sha256(json.dumps(body, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest(), "started_at": time.time()}, journal) try: submitted = api("POST", route, key, body) except requests.exceptions.RequestException as error: raise RuntimeError("Submission state is unknown. Do not resubmit; use platform request records to recover the job id") from error (output / "submitted.json").write_text(json.dumps(submitted, ensure_ascii=False, indent=2), encoding="utf-8") job_id = submitted.get("id") if not isinstance(job_id, str) or not job_id: raise RuntimeError("Submission returned no job id; do not automatically resubmit") print("Job ID:", job_id, flush=True) if any(x not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for x in job_id): raise ValueError("Unsafe job ID for local filenames") data = wait_job(job_id, key, output, seconds=args.wait) print(json.dumps(download_all(data, job_id, output), ensure_ascii=False)) if __name__ == "__main__": main()