#!/usr/bin/env python3 """Python 3.10+, standard library only. Submit once; resume with --task.""" import argparse import base64 import json import math import mimetypes import os import re import sys import time import urllib.error import urllib.parse import urllib.request from pathlib import Path class NoRedirect(urllib.request.HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): return None # Never forward API credentials through a redirect. class APIError(RuntimeError): def __init__(self, status, message, retry_after=None): super().__init__(f'HTTP {status}: {message}') self.status, self.retry_after = status, retry_after def interval(value, default=3): try: number = float(value) return min(30, max(1, number)) if math.isfinite(number) else default except (ValueError, TypeError): return default class ImageClient: def __init__(self, key, base_url='https://api.manxiaobai.online'): self.key = key self.base = base_url.rstrip('/') parsed = urllib.parse.urlsplit(self.base) if parsed.scheme != 'https' or parsed.username or parsed.password or parsed.path or parsed.query or parsed.fragment: raise ValueError('base_url must be an HTTPS origin, without /v1 or credentials') self.opener = urllib.request.build_opener(NoRedirect()) def request(self, path, payload=None, timeout=30): if not path.startswith('/v1/') or path.startswith('//'): raise ValueError('Only fixed /v1/ API paths are accepted') body = None if payload is None else json.dumps(payload, ensure_ascii=False).encode('utf-8') request = urllib.request.Request(self.base + path, data=body, headers={ 'Authorization': 'Bearer ' + self.key, 'Content-Type': 'application/json', }, method='GET' if payload is None else 'POST') try: with self.opener.open(request, timeout=timeout) as response: result = json.load(response) except urllib.error.HTTPError as error: raw = error.read(2048).decode('utf-8', errors='replace') raise APIError(error.code, raw, error.headers.get('Retry-After')) from error if not isinstance(result, dict): raise RuntimeError('API response must be a JSON object') return result def submit(self, model, prompt, size='1024x1024', reference=None): body = {'model': model, 'prompt': prompt, 'size': size, 'n': 1} path = '/v1/images/generations/async' if reference: file = Path(reference) mime = mimetypes.guess_type(file.name)[0] if mime not in ('image/png', 'image/jpeg', 'image/webp', 'image/bmp'): raise ValueError('Use a real PNG/JPEG/WebP/BMP reference image') raw = file.read_bytes() if not raw or len(raw) > 25 * 1024 * 1024: raise ValueError('Reference image must be nonempty and at most 25 MiB; channel limits may be lower') body['images'] = [{'image_url': f'data:{mime};base64,' + base64.b64encode(raw).decode('ascii')}] path = '/v1/images/edits/async' # Intentionally exactly one POST. A timeout may mean submission succeeded. result = self.request(path, body, timeout=60) task_id = result.get('task_id') or result.get('id') if not isinstance(task_id, str) or not re.fullmatch(r'[A-Za-z0-9_-]+', task_id): raise RuntimeError('No valid task ID returned. Check logs before submitting again.') return task_id, result def wait(self, task_id, timeout=1800): if not re.fullmatch(r'[A-Za-z0-9_-]+', task_id): raise ValueError('Invalid task ID') deadline, delay = time.monotonic() + timeout, 3 while time.monotonic() < deadline: try: task = self.request('/v1/images/tasks/' + task_id) except APIError as error: if error.status not in (429, 500, 502, 503, 504): raise time.sleep(interval(error.retry_after, delay)) delay = min(30, delay * 1.5) continue except (urllib.error.URLError, TimeoutError): time.sleep(delay) delay = min(30, delay * 1.5) continue status = str(task.get('status') or task.get('state') or '').lower() if task.get('submission_uncertain'): raise RuntimeError(f'Submission uncertain; check task {task_id} and billing before any new POST') if status in ('completed', 'succeeded', 'success', 'done'): nested = task.get('result') data = task.get('data') or (nested.get('data') if isinstance(nested, dict) else None) if not isinstance(data, list) or not data or not isinstance(data[0], dict): raise RuntimeError(f'Task {task_id} completed without valid image data; do not resubmit automatically') return data, task if status in ('failed', 'failure', 'error', 'cancelled', 'canceled', 'expired'): raise RuntimeError(f'Task {task_id} failed: {json.dumps(task.get("error"), ensure_ascii=False)}') delay = interval(task.get('poll_after_ms', 3000) / 1000 if isinstance(task.get('poll_after_ms', 3000), (int, float)) else 3) time.sleep(delay) raise TimeoutError(f'Stopped waiting. Resume with --task {task_id}; this does not cancel the server task.') def download_image(item, destination): if item.get('url'): url = str(item['url']) if urllib.parse.urlsplit(url).scheme not in ('https', 'http'): raise ValueError('Unsupported result URL scheme') # Separate unauthenticated request; never send the API Key to image hosts. with urllib.request.urlopen(url, timeout=120) as response: if not response.headers.get('Content-Type', '').lower().startswith('image/'): raise RuntimeError('Download did not return an image Content-Type') data = response.read() elif item.get('b64_json'): data = base64.b64decode(item['b64_json'], validate=True) else: raise RuntimeError('Image result has neither URL nor Base64') if not data: raise RuntimeError('Empty image result') Path(destination).write_bytes(data) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--model', default='gpt-image-2.5') parser.add_argument('--prompt') parser.add_argument('--size', default='1024x1024') parser.add_argument('--reference') parser.add_argument('--task', help='Resume an existing task; never creates another task') parser.add_argument('--timeout', type=int, default=1800) parser.add_argument('--output', default='result.png') args = parser.parse_args() key = os.environ.get('MXB_API_KEY') if not key: parser.error('Set MXB_API_KEY in the server environment') if args.timeout <= 0: parser.error('--timeout must be positive') if not args.task and not args.prompt: parser.error('--prompt is required when creating a task') client = ImageClient(key) task_id = args.task if not task_id: try: task_id, response = client.submit(args.model, args.prompt, args.size, args.reference) except Exception: print('Submission was not retried. If transmission started, check logs before another POST.', file=sys.stderr) raise print('Task:', task_id, flush=True) Path(f'{task_id}.json').write_text(json.dumps(response, ensure_ascii=False, indent=2), encoding='utf-8') data, final = client.wait(task_id, args.timeout) Path(f'{task_id}.result.json').write_text(json.dumps(final, ensure_ascii=False, indent=2), encoding='utf-8') download_image(data[0], args.output) print('Saved:', args.output, '(file extension does not convert the image format)') if __name__ == '__main__': main()