@leejungkiin/awkit 1.6.6 → 1.7.0

package/scripts/openrouter_image_gen.py

@@ -0,0 +1,772 @@
+#!/usr/bin/env python3
+"""OpenRouter Image Generation CLI — ported from Codex image_gen.py.
+
+Supports generate (text-to-image) and edit (image-to-image) via OpenRouter
+chat/completions endpoint with GPT Image models.
+
+Features ported from Codex:
+- Structured prompt augmentation (--style, --lighting, --palette, etc.)
+- Auto-downscale for web assets
+- Retry with exponential backoff on transient/rate-limit errors
+- Async batch processing from JSONL
+- Dry-run mode for payload inspection
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+import re
+import sys
+import time
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+DEFAULT_MODEL = "openai/gpt-5.4-image-2"
+DEFAULT_SIZE = "auto"
+DEFAULT_QUALITY = "medium"
+DEFAULT_OUTPUT_FORMAT = "webp"
+DEFAULT_CONCURRENCY = 3
+DEFAULT_DOWNSCALE_SUFFIX = "-web"
+DEFAULT_MAX_TOKENS = 4096
+OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
+TMPFILES_UPLOAD_URL = "https://tmpfiles.org/api/v1/upload"
+MAX_IMAGE_BYTES = 50 * 1024 * 1024
+MAX_BATCH_JOBS = 100
+
+CRED_PATH = os.path.expanduser("~/.gemini/antigravity/.credentials.json")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _die(message: str, code: int = 1) -> None:
+    print(f"Error: {message}", file=sys.stderr)
+    raise SystemExit(code)
+
+
+def _warn(message: str) -> None:
+    print(f"Warning: {message}", file=sys.stderr)
+
+
+def _load_api_key() -> str:
+    key = os.getenv("OPENROUTER_API_KEY")
+    if key:
+        return key
+    try:
+        with open(CRED_PATH, "r") as f:
+            creds = json.load(f)
+            key = creds.get("openrouter_api_key")
+    except Exception:
+        pass
+    if not key:
+        _die(
+            "OpenRouter API key not found. Set OPENROUTER_API_KEY env var "
+            "or add openrouter_api_key to ~/.gemini/antigravity/.credentials.json"
+        )
+    return key
+
+
+def _ensure_requests():
+    try:
+        import requests  # noqa: F401
+        return requests
+    except ImportError:
+        _die("requests library is required. Install with: pip install requests")
+
+
+# ---------------------------------------------------------------------------
+# Image upload (local file → tmpfiles.org → public URL)
+# ---------------------------------------------------------------------------
+
+def _upload_to_tmpfiles(image_path: str) -> str:
+    requests = _ensure_requests()
+    print(f"📤 Uploading local image to tmpfiles.org...", file=sys.stderr)
+    try:
+        with open(image_path, "rb") as f:
+            resp = requests.post(TMPFILES_UPLOAD_URL, files={"file": f}, timeout=60)
+        resp.raise_for_status()
+        data = resp.json()
+        if "data" in data and "url" in data["data"]:
+            url = data["data"]["url"]
+            direct_url = url.replace("tmpfiles.org/", "tmpfiles.org/dl/")
+            print(f"✅ Uploaded: {direct_url}", file=sys.stderr)
+            return direct_url
+    except Exception as e:
+        _warn(f"tmpfiles.org upload failed ({e}), falling back to base64")
+
+    # Fallback: inline base64
+    import mimetypes
+    mime_type, _ = mimetypes.guess_type(image_path)
+    if not mime_type:
+        mime_type = "image/png"
+    with open(image_path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def _resolve_image_url(image_input: str) -> str:
+    if image_input.startswith(("http://", "https://", "data:")):
+        return image_input
+    path = Path(image_input)
+    if not path.exists():
+        _die(f"Image file not found: {path}")
+    if path.stat().st_size > MAX_IMAGE_BYTES:
+        _warn(f"Image exceeds 50MB: {path}")
+    return _upload_to_tmpfiles(str(path))
+
+
+# ---------------------------------------------------------------------------
+# Prompt augmentation (ported from Codex)
+# ---------------------------------------------------------------------------
+
+def _fields_from_args(args: argparse.Namespace) -> Dict[str, Optional[str]]:
+    return {
+        "use_case": getattr(args, "use_case", None),
+        "scene": getattr(args, "scene", None),
+        "subject": getattr(args, "subject", None),
+        "style": getattr(args, "style", None),
+        "composition": getattr(args, "composition", None),
+        "lighting": getattr(args, "lighting", None),
+        "palette": getattr(args, "palette", None),
+        "materials": getattr(args, "materials", None),
+        "text": getattr(args, "text", None),
+        "constraints": getattr(args, "constraints", None),
+        "negative": getattr(args, "negative", None),
+    }
+
+
+def _augment_prompt(augment: bool, prompt: str, fields: Dict[str, Optional[str]]) -> str:
+    if not augment:
+        return prompt
+    sections: List[str] = []
+    mapping = [
+        ("use_case", "Use case"),
+        ("scene", "Scene/background"),
+        ("subject", "Subject"),
+        ("style", "Style/medium"),
+        ("composition", "Composition/framing"),
+        ("lighting", "Lighting/mood"),
+        ("palette", "Color palette"),
+        ("materials", "Materials/textures"),
+        ("text", "Text (verbatim)"),
+        ("constraints", "Constraints"),
+        ("negative", "Avoid"),
+    ]
+    for key, label in mapping:
+        val = fields.get(key)
+        if val and key != "text":
+            sections.append(f"{label}: {val}")
+        elif val and key == "text":
+            sections.append(f'{label}: "{val}"')
+    if sections:
+        return f"Primary request: {prompt}\n" + "\n".join(sections)
+    return prompt
+
+
+# ---------------------------------------------------------------------------
+# Retry logic (ported from Codex)
+# ---------------------------------------------------------------------------
+
+def _is_transient(status_code: int) -> bool:
+    return status_code in (429, 500, 502, 503, 504)
+
+
+def _extract_retry_after(headers: dict) -> Optional[float]:
+    val = headers.get("retry-after") or headers.get("Retry-After")
+    if val:
+        try:
+            return float(val)
+        except ValueError:
+            pass
+    return None
+
+
+def _call_openrouter(
+    api_key: str,
+    payload: dict,
+    *,
+    max_attempts: int = 3,
+    timeout: int = 180,
+    label: str = "",
+) -> dict:
+    requests = _ensure_requests()
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+        "HTTP-Referer": "https://github.com/antigravity-awf",
+        "X-Title": "Antigravity Image Gen",
+    }
+    last_exc = None
+    for attempt in range(1, max_attempts + 1):
+        try:
+            resp = requests.post(
+                OPENROUTER_API_URL,
+                headers=headers,
+                json=payload,
+                timeout=timeout,
+            )
+            if resp.status_code == 200:
+                return resp.json()
+            if _is_transient(resp.status_code) and attempt < max_attempts:
+                sleep_s = _extract_retry_after(dict(resp.headers)) or min(60.0, 2.0 ** attempt)
+                print(
+                    f"{label} attempt {attempt}/{max_attempts} got {resp.status_code}; "
+                    f"retrying in {sleep_s:.1f}s",
+                    file=sys.stderr,
+                )
+                time.sleep(sleep_s)
+                continue
+            # Non-transient error or last attempt — log response body
+            try:
+                err_body = resp.text[:500]
+            except Exception:
+                err_body = "(unreadable)"
+            print(f"{label} API error {resp.status_code}: {err_body}", file=sys.stderr)
+            resp.raise_for_status()
+        except Exception as e:
+            last_exc = e
+            if attempt == max_attempts:
+                raise
+            sleep_s = min(60.0, 2.0 ** attempt)
+            print(
+                f"{label} attempt {attempt}/{max_attempts} failed ({e.__class__.__name__}); "
+                f"retrying in {sleep_s:.1f}s",
+                file=sys.stderr,
+            )
+            time.sleep(sleep_s)
+    raise last_exc or RuntimeError("unknown error")
+
+
+async def _call_openrouter_async(
+    api_key: str,
+    payload: dict,
+    *,
+    max_attempts: int = 3,
+    timeout: int = 180,
+    label: str = "",
+) -> dict:
+    """Async wrapper — runs sync call in executor to avoid blocking event loop."""
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        None,
+        lambda: _call_openrouter(
+            api_key, payload, max_attempts=max_attempts, timeout=timeout, label=label
+        ),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Response parsing — extract image from OpenRouter chat/completions response
+# ---------------------------------------------------------------------------
+
+def _extract_images_from_response(data: dict) -> List[bytes]:
+    """Extract base64 image bytes from OpenRouter response."""
+    images: List[bytes] = []
+    if "choices" not in data or not data["choices"]:
+        return images
+
+    message = data["choices"][0].get("message", {})
+
+    # Path 1: message.images[] (OpenAI native format via OpenRouter)
+    if "images" in message:
+        for img in message["images"]:
+            if isinstance(img, dict):
+                url = img.get("image_url", {}).get("url", "")
+            else:
+                url = str(img)
+            if url.startswith("data:image"):
+                _, encoded = url.split(",", 1)
+                images.append(base64.b64decode(encoded))
+            elif url.startswith("http"):
+                requests = _ensure_requests()
+                resp = requests.get(url, timeout=60)
+                resp.raise_for_status()
+                images.append(resp.content)
+
+    # Path 2: content contains markdown image or inline base64
+    if not images and "content" in message:
+        content = message["content"] or ""
+        # Check for base64 data URIs
+        b64_matches = re.findall(r'data:image/[^;]+;base64,([A-Za-z0-9+/=]+)', content)
+        for b64 in b64_matches:
+            try:
+                images.append(base64.b64decode(b64))
+            except Exception:
+                pass
+        # Check for HTTP URLs
+        if not images:
+            url_matches = re.findall(r'(https?://\S+\.(?:png|jpg|jpeg|webp|gif))', content)
+            if url_matches:
+                requests = _ensure_requests()
+                for url in url_matches[:1]:
+                    try:
+                        resp = requests.get(url, timeout=60)
+                        resp.raise_for_status()
+                        images.append(resp.content)
+                    except Exception:
+                        pass
+
+    return images
+
+
+# ---------------------------------------------------------------------------
+# Output helpers
+# ---------------------------------------------------------------------------
+
+def _build_output_paths(
+    out: str, output_format: str, count: int, out_dir: Optional[str]
+) -> List[Path]:
+    ext = "." + output_format
+    if out_dir:
+        base = Path(out_dir)
+        base.mkdir(parents=True, exist_ok=True)
+        return [base / f"image_{i}{ext}" for i in range(1, count + 1)]
+
+    out_path = Path(out)
+    if out_path.suffix == "":
+        out_path = out_path.with_suffix(ext)
+
+    if count == 1:
+        return [out_path]
+    return [
+        out_path.with_name(f"{out_path.stem}-{i}{out_path.suffix}")
+        for i in range(1, count + 1)
+    ]
+
+
+def _derive_downscale_path(path: Path, suffix: str) -> Path:
+    if suffix and not suffix.startswith(("-", "_")):
+        suffix = "-" + suffix
+    return path.with_name(f"{path.stem}{suffix}{path.suffix}")
+
+
+def _downscale_image_bytes(image_bytes: bytes, *, max_dim: int, output_format: str) -> bytes:
+    try:
+        from PIL import Image
+    except ImportError:
+        _die("Downscaling requires Pillow. Install with: pip install Pillow")
+
+    with Image.open(BytesIO(image_bytes)) as img:
+        img.load()
+        w, h = img.size
+        scale = min(1.0, float(max_dim) / float(max(w, h)))
+        target = (max(1, int(round(w * scale))), max(1, int(round(h * scale))))
+        resized = img if target == (w, h) else img.resize(target, Image.Resampling.LANCZOS)
+
+        fmt = output_format.lower()
+        if fmt in ("jpg", "jpeg"):
+            fmt = "jpeg"
+            if resized.mode in ("RGBA", "LA"):
+                bg = Image.new("RGB", resized.size, (255, 255, 255))
+                bg.paste(resized, mask=resized.split()[-1])
+                resized = bg
+            else:
+                resized = resized.convert("RGB")
+
+        out = BytesIO()
+        resized.save(out, format=fmt.upper())
+        return out.getvalue()
+
+
+def _convert_image_format(image_bytes: bytes, output_format: str) -> bytes:
+    try:
+        from PIL import Image
+    except ImportError:
+        _die("Format conversion requires Pillow. Install with: pip install Pillow")
+
+    with Image.open(BytesIO(image_bytes)) as img:
+        fmt = output_format.lower()
+        if fmt in ("jpg", "jpeg"):
+            fmt = "jpeg"
+            if img.mode in ("RGBA", "LA"):
+                bg = Image.new("RGB", img.size, (255, 255, 255))
+                bg.paste(img, mask=img.split()[-1])
+                img = bg
+            else:
+                img = img.convert("RGB")
+
+        out = BytesIO()
+        img.save(out, format=fmt.upper())
+        return out.getvalue()
+
+
+def _write_images(
+    images: List[bytes],
+    outputs: List[Path],
+    *,
+    force: bool = False,
+    downscale_max_dim: Optional[int] = None,
+    downscale_suffix: str = DEFAULT_DOWNSCALE_SUFFIX,
+    output_format: str = DEFAULT_OUTPUT_FORMAT,
+) -> None:
+    for idx, img_bytes in enumerate(images):
+        if idx >= len(outputs):
+            break
+            
+        fmt = output_format.lower()
+        is_webp = img_bytes.startswith(b"RIFF") and b"WEBP" in img_bytes[:16]
+        is_png = img_bytes.startswith(b"\x89PNG\r\n\x1a\n")
+        is_jpeg = img_bytes.startswith(b"\xff\xd8")
+        
+        needs_convert = False
+        if fmt == "webp" and not is_webp:
+            needs_convert = True
+        elif fmt in ("jpg", "jpeg") and not is_jpeg:
+            needs_convert = True
+        elif fmt == "png" and not is_png:
+            needs_convert = True
+            
+        if needs_convert:
+            try:
+                img_bytes = _convert_image_format(img_bytes, output_format)
+            except Exception as e:
+                _warn(f"Failed to convert image format to {output_format}: {e}")
+
+        out_path = outputs[idx]
+        if out_path.exists() and not force:
+            _die(f"Output already exists: {out_path} (use --force to overwrite)")
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_bytes(img_bytes)
+        print(f"Wrote {out_path}")
+
+        if downscale_max_dim is not None:
+            derived = _derive_downscale_path(out_path, downscale_suffix)
+            if derived.exists() and not force:
+                _die(f"Output already exists: {derived} (use --force)")
+            resized = _downscale_image_bytes(
+                img_bytes, max_dim=downscale_max_dim, output_format=output_format
+            )
+            derived.write_bytes(resized)
+            print(f"Wrote {derived} (downscaled to max {downscale_max_dim}px)")
+
+
+# ---------------------------------------------------------------------------
+# Build OpenRouter payload
+# ---------------------------------------------------------------------------
+
+def _build_messages(prompt: str, image_urls: Optional[List[str]] = None) -> list:
+    if image_urls:
+        content: list = [{"type": "text", "text": prompt}]
+        for url in image_urls:
+            content.append({"type": "image_url", "image_url": {"url": url}})
+        return [{"role": "user", "content": content}]
+    return [{"role": "user", "content": prompt}]
+
+
+def _build_payload(
+    model: str,
+    messages: list,
+    *,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+) -> dict:
+    return {
+        "model": model,
+        "messages": messages,
+        "max_tokens": max_tokens,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Commands: generate, edit, generate-batch
+# ---------------------------------------------------------------------------
+
+def _cmd_generate(args: argparse.Namespace) -> None:
+    api_key = _load_api_key()
+    prompt = _read_prompt(args.prompt, getattr(args, "prompt_file", None))
+    fields = _fields_from_args(args)
+    prompt = _augment_prompt(args.augment, prompt, fields)
+
+    messages = _build_messages(prompt)
+    payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)
+
+    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT
+    outputs = _build_output_paths(args.out, output_format, args.n, args.out_dir)
+
+    if args.dry_run:
+        print(json.dumps({"outputs": [str(p) for p in outputs], **payload}, indent=2))
+        return
+
+    print(f"🎨 [OpenRouter] Generating {args.n} image(s) with {args.model}...", file=sys.stderr)
+    print(f"📝 Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}", file=sys.stderr)
+
+    started = time.time()
+    images = []
+    for i in range(args.n):
+        if args.n > 1:
+            print(f"⏳ Generating image {i+1}/{args.n}...", file=sys.stderr)
+            
+        data = _call_openrouter(api_key, payload, max_attempts=args.max_attempts, label=f"[generate {i+1}]" if args.n > 1 else "[generate]")
+        batch_images = _extract_images_from_response(data)
+        if batch_images:
+            images.extend(batch_images)
+        else:
+            _die(f"No images found in response for image {i+1}.\nResponse: {json.dumps(data, indent=2)}")
+
+    elapsed = time.time() - started
+    print(f"✅ Generation completed in {elapsed:.1f}s.", file=sys.stderr)
+
+    _write_images(
+        images,
+        outputs,
+        force=args.force,
+        downscale_max_dim=args.downscale_max_dim,
+        downscale_suffix=args.downscale_suffix,
+        output_format=output_format,
+    )
+
+    for p in outputs:
+        if p.exists():
+            print(f"✨ file://{p}")
+
+
+def _cmd_edit(args: argparse.Namespace) -> None:
+    api_key = _load_api_key()
+    prompt = _read_prompt(args.prompt, getattr(args, "prompt_file", None))
+    fields = _fields_from_args(args)
+    prompt = _augment_prompt(args.augment, prompt, fields)
+
+    # Resolve all input images to URLs
+    image_urls = [_resolve_image_url(img) for img in args.image]
+
+    messages = _build_messages(prompt, image_urls)
+    payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)
+
+    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT
+    outputs = _build_output_paths(args.out, output_format, args.n, args.out_dir)
+
+    if args.dry_run:
+        print(json.dumps({"outputs": [str(p) for p in outputs], **payload}, indent=2))
+        return
+
+    print(f"🎨 [OpenRouter] Editing with {args.model} ({len(image_urls)} image(s)) x {args.n}...", file=sys.stderr)
+    print(f"📝 Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}", file=sys.stderr)
+
+    started = time.time()
+    images = []
+    for i in range(args.n):
+        if args.n > 1:
+            print(f"⏳ Editing image {i+1}/{args.n}...", file=sys.stderr)
+            
+        data = _call_openrouter(api_key, payload, max_attempts=args.max_attempts, label=f"[edit {i+1}]" if args.n > 1 else "[edit]")
+        batch_images = _extract_images_from_response(data)
+        if batch_images:
+            images.extend(batch_images)
+        else:
+            _die(f"No images found in response for edit {i+1}.\nResponse: {json.dumps(data, indent=2)}")
+
+    elapsed = time.time() - started
+    print(f"✅ Edit completed in {elapsed:.1f}s.", file=sys.stderr)
+
+    _write_images(
+        images,
+        outputs,
+        force=args.force,
+        downscale_max_dim=args.downscale_max_dim,
+        downscale_suffix=args.downscale_suffix,
+        output_format=output_format,
+    )
+
+    for p in outputs:
+        if p.exists():
+            print(f"✨ file://{p}")
+
+
+def _cmd_generate_batch(args: argparse.Namespace) -> None:
+    api_key = _load_api_key()
+    jobs = _read_jobs_jsonl(args.input)
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    base_fields = _fields_from_args(args)
+    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT
+
+    print(f"🎨 [OpenRouter Batch] {len(jobs)} jobs, concurrency={args.concurrency}", file=sys.stderr)
+
+    async def run_all():
+        sem = asyncio.Semaphore(args.concurrency)
+        failed = 0
+
+        async def run_one(i: int, job: dict):
+            nonlocal failed
+            prompt = str(job["prompt"]).strip()
+            job_fields = {k: job.get(k, base_fields.get(k)) for k in base_fields}
+            augmented = _augment_prompt(args.augment, prompt, job_fields)
+
+            image_urls = None
+            if "image" in job:
+                raw_images = job["image"] if isinstance(job["image"], list) else [job["image"]]
+                image_urls = [_resolve_image_url(img) for img in raw_images]
+
+            messages = _build_messages(augmented, image_urls)
+            payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)
+
+            ext = "." + output_format
+            out_path = out_dir / f"{i:03d}-{_slugify(prompt[:60])}{ext}"
+            label = f"[job {i}/{len(jobs)}]"
+
+            try:
+                async with sem:
+                    print(f"{label} starting", file=sys.stderr)
+                    started = time.time()
+                    data = await _call_openrouter_async(
+                        api_key, payload, max_attempts=args.max_attempts, label=label
+                    )
+                    elapsed = time.time() - started
+                    print(f"{label} completed in {elapsed:.1f}s", file=sys.stderr)
+
+                images = _extract_images_from_response(data)
+                if images:
+                    _write_images(images, [out_path], force=args.force, output_format=output_format)
+                else:
+                    print(f"{label} no images in response", file=sys.stderr)
+                    failed += 1
+            except Exception as e:
+                print(f"{label} failed: {e}", file=sys.stderr)
+                failed += 1
+                if args.fail_fast:
+                    raise
+
+        tasks = [asyncio.create_task(run_one(i, job)) for i, job in enumerate(jobs, 1)]
+        await asyncio.gather(*tasks, return_exceptions=not args.fail_fast)
+        return failed
+
+    failed = asyncio.run(run_all())
+    if failed:
+        raise SystemExit(1)
+
+
+# ---------------------------------------------------------------------------
+# Utility
+# ---------------------------------------------------------------------------
+
+def _read_prompt(prompt: Optional[str], prompt_file: Optional[str]) -> str:
+    if prompt and prompt_file:
+        _die("Use --prompt or --prompt-file, not both.")
+    if prompt_file:
+        p = Path(prompt_file)
+        if not p.exists():
+            _die(f"Prompt file not found: {p}")
+        return p.read_text(encoding="utf-8").strip()
+    if prompt:
+        return prompt.strip()
+    _die("Missing prompt. Use --prompt or --prompt-file.")
+    return ""
+
+
+def _slugify(value: str) -> str:
+    value = value.strip().lower()
+    value = re.sub(r"[^a-z0-9]+", "-", value)
+    value = re.sub(r"-{2,}", "-", value).strip("-")
+    return value[:60] if value else "output"
+
+
+def _read_jobs_jsonl(path: str) -> List[Dict[str, Any]]:
+    p = Path(path)
+    if not p.exists():
+        _die(f"Input file not found: {p}")
+    jobs: List[Dict[str, Any]] = []
+    for line_no, raw in enumerate(p.read_text(encoding="utf-8").splitlines(), start=1):
+        line = raw.strip()
+        if not line or line.startswith("#"):
+            continue
+        try:
+            if line.startswith("{"):
+                item = json.loads(line)
+            else:
+                item = {"prompt": line}
+            if "prompt" not in item or not str(item["prompt"]).strip():
+                _die(f"Missing prompt at line {line_no}")
+            jobs.append(item)
+        except json.JSONDecodeError as e:
+            _die(f"Invalid JSON on line {line_no}: {e}")
+    if not jobs:
+        _die("No jobs found in input file.")
+    if len(jobs) > MAX_BATCH_JOBS:
+        _die(f"Too many jobs ({len(jobs)}). Max is {MAX_BATCH_JOBS}.")
+    return jobs
+
+
+# ---------------------------------------------------------------------------
+# CLI parser
+# ---------------------------------------------------------------------------
+
+def _add_shared_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="OpenRouter model ID")
+    parser.add_argument("--prompt", help="Text prompt for generation")
+    parser.add_argument("--prompt-file", help="Read prompt from file")
+    parser.add_argument("--n", type=int, default=1, help="Number of images (1-10)")
+    parser.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS)
+    parser.add_argument("--out", default="output.webp", help="Output path")
+    parser.add_argument("--out-dir", help="Output directory (for batch)")
+    parser.add_argument("--output-format", default=DEFAULT_OUTPUT_FORMAT, choices=["png", "webp", "jpeg"])
+    parser.add_argument("--force", action="store_true", help="Overwrite existing files")
+    parser.add_argument("--dry-run", action="store_true", help="Print payload without calling API")
+    parser.add_argument("--max-attempts", type=int, default=3, help="Retry attempts on transient errors")
+
+    # Prompt augmentation
+    parser.add_argument("--augment", dest="augment", action="store_true")
+    parser.add_argument("--no-augment", dest="augment", action="store_false")
+    parser.set_defaults(augment=True)
+    parser.add_argument("--use-case", help="Augment: use case context")
+    parser.add_argument("--scene", help="Augment: scene/background description")
+    parser.add_argument("--subject", help="Augment: subject description")
+    parser.add_argument("--style", help="Augment: art style/medium")
+    parser.add_argument("--composition", help="Augment: composition/framing")
+    parser.add_argument("--lighting", help="Augment: lighting/mood")
+    parser.add_argument("--palette", help="Augment: color palette")
+    parser.add_argument("--materials", help="Augment: materials/textures")
+    parser.add_argument("--text", help="Augment: text to render (verbatim)")
+    parser.add_argument("--constraints", help="Augment: constraints")
+    parser.add_argument("--negative", help="Augment: things to avoid")
+
+    # Post-processing: downscale
+    parser.add_argument("--downscale-max-dim", type=int, help="Generate additional downscaled copy")
+    parser.add_argument("--downscale-suffix", default=DEFAULT_DOWNSCALE_SUFFIX)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="OpenRouter Image Generation CLI (Codex-grade)"
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # generate
+    gen_parser = subparsers.add_parser("generate", help="Create a new image from text")
+    _add_shared_args(gen_parser)
+    gen_parser.set_defaults(func=_cmd_generate)
+
+    # edit (image-to-image)
+    edit_parser = subparsers.add_parser("edit", help="Edit/transform an existing image")
+    _add_shared_args(edit_parser)
+    edit_parser.add_argument("--image", action="append", required=True, help="Input image path/URL")
+    edit_parser.set_defaults(func=_cmd_edit)
+
+    # generate-batch
+    batch_parser = subparsers.add_parser("generate-batch", help="Batch generate from JSONL")
+    _add_shared_args(batch_parser)
+    batch_parser.add_argument("--input", required=True, help="Path to JSONL file")
+    batch_parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY)
+    batch_parser.add_argument("--fail-fast", action="store_true")
+    batch_parser.set_defaults(func=_cmd_generate_batch)
+
+    args = parser.parse_args()
+
+    # Validation
+    if args.n < 1 or args.n > 10:
+        _die("--n must be between 1 and 10")
+    if args.max_attempts < 1 or args.max_attempts > 10:
+        _die("--max-attempts must be between 1 and 10")
+    if args.command == "generate-batch" and not args.out_dir:
+        _die("generate-batch requires --out-dir")
+
+    args.func(args)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())