#!/usr/bin/env python3 """CLI for explicit image generation or editing through Novamailio Image API. Novamailio exposes an OpenAI-compatible /v1 Image API. This script keeps the same request shape as the imagegen skill's GPT Image CLI while defaulting the provider object, API key, and base URL to Novamailio. Defaults to gpt-image-2 and a structured prompt augmentation workflow. """ from __future__ import annotations import argparse import asyncio import base64 import json import os from pathlib import Path import re import sys import time from typing import Any, Dict, Iterable, List, Optional, Tuple from io import BytesIO DEFAULT_MODEL = "gpt-image-2" DEFAULT_SIZE = "auto" DEFAULT_QUALITY = "medium" DEFAULT_OUTPUT_FORMAT = "png" DEFAULT_CONCURRENCY = 5 DEFAULT_DOWNSCALE_SUFFIX = "-web" DEFAULT_OUTPUT_PATH = "output/novamailio-imagegen/output.png" DEFAULT_BASE_URL = "https://api.novamailio.com/v1" DEFAULT_MAX_ATTEMPTS = 3 DEFAULT_RETRY_DELAY = 30.0 DEFAULT_TIMEOUT = 600.0 GPT_IMAGE_MODEL_PREFIX = "gpt-image-" ALLOWED_LEGACY_SIZES = {"1024x1024", "1536x1024", "1024x1536", "auto"} ALLOWED_QUALITIES = {"low", "medium", "high", "auto"} ALLOWED_BACKGROUNDS = {"transparent", "opaque", "auto", None} GPT_IMAGE_2_MODEL = "gpt-image-2" GPT_IMAGE_2_MIN_PIXELS = 655_360 GPT_IMAGE_2_MAX_PIXELS = 8_294_400 GPT_IMAGE_2_MAX_EDGE = 3840 GPT_IMAGE_2_MAX_RATIO = 3.0 MAX_IMAGE_BYTES = 50 * 1024 * 1024 MAX_BATCH_JOBS = 500 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 _dependency_hint(package: str, *, upgrade: bool = False) -> str: command = f"uv pip install {'-U ' if upgrade else ''}{package}" return ( "Activate the repo-selected environment first, then install it with " f"`{command}`. If this repo uses a local virtualenv, start with " "`source .venv/bin/activate`; otherwise use this repo's configured shared fallback " "environment. If your project declares dependencies, prefer that project's normal " "`uv sync` flow." ) def _normalize_base_url(value: Optional[str]) -> str: url = (value or DEFAULT_BASE_URL).strip().rstrip("/") if not url: return DEFAULT_BASE_URL if url == "https://api.novamailio.com": return DEFAULT_BASE_URL return url def _resolve_api_key(*, api_key: Optional[str], dry_run: bool) -> Optional[str]: if api_key: print("Using Novamailio API key from --api-key.", file=sys.stderr) return api_key env_key = os.getenv("NOVAMAILIO_API_KEY") if env_key: print("NOVAMAILIO_API_KEY is set.", file=sys.stderr) return env_key if dry_run: _warn("NOVAMAILIO_API_KEY is not set; dry-run only.") return None _die("NOVAMAILIO_API_KEY is not set. Configure it in the backend server environment.") return None # unreachable def _is_auth_error(exc: Exception) -> bool: status_code = getattr(exc, "status_code", None) or getattr(exc, "code", None) if status_code in {401, 403}: return True name = exc.__class__.__name__.lower() msg = str(exc).lower() return ( "authentication" in name or "permission" in name or "invalid_api_key" in msg or "invalid api key" in msg or "unauthorized" in msg or "forbidden" in msg or "401" in msg or "403" in msg ) def _refresh_api_key_after_auth_error(args: argparse.Namespace, exc: Exception) -> bool: _warn(f"Novamailio API key appears invalid: {exc}") return False 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: path = Path(prompt_file) if not path.exists(): _die(f"Prompt file not found: {path}") return path.read_text(encoding="utf-8").strip() if prompt: return prompt.strip() _die("Missing prompt. Use --prompt or --prompt-file.") return "" # unreachable def _check_image_paths(paths: Iterable[str]) -> List[Path]: resolved: List[Path] = [] for raw in paths: path = Path(raw) if not path.exists(): _die(f"Image file not found: {path}") if path.stat().st_size > MAX_IMAGE_BYTES: _warn(f"Image exceeds 50MB limit: {path}") resolved.append(path) return resolved def _normalize_output_format(fmt: Optional[str]) -> str: if not fmt: return DEFAULT_OUTPUT_FORMAT fmt = fmt.lower() if fmt not in {"png", "jpeg", "jpg", "webp"}: _die("output-format must be png, jpeg, jpg, or webp.") return "jpeg" if fmt == "jpg" else fmt def _parse_size(size: str) -> Optional[Tuple[int, int]]: match = re.fullmatch(r"([1-9][0-9]*)x([1-9][0-9]*)", size) if not match: return None return int(match.group(1)), int(match.group(2)) def _validate_gpt_image_2_size(size: str) -> None: if size == "auto": return parsed = _parse_size(size) if parsed is None: _die("size must be auto or WIDTHxHEIGHT, for example 1024x1024.") width, height = parsed max_edge = max(width, height) min_edge = min(width, height) total_pixels = width * height if max_edge > GPT_IMAGE_2_MAX_EDGE: _die("gpt-image-2 size maximum edge length must be less than or equal to 3840px.") if width % 16 != 0 or height % 16 != 0: _die("gpt-image-2 size width and height must be multiples of 16px.") if max_edge / min_edge > GPT_IMAGE_2_MAX_RATIO: _die("gpt-image-2 size long edge to short edge ratio must not exceed 3:1.") if total_pixels < GPT_IMAGE_2_MIN_PIXELS or total_pixels > GPT_IMAGE_2_MAX_PIXELS: _die( "gpt-image-2 size total pixels must be at least 655,360 and no more than 8,294,400." ) def _validate_size(size: str, model: str) -> None: if model == GPT_IMAGE_2_MODEL: _validate_gpt_image_2_size(size) return if size not in ALLOWED_LEGACY_SIZES: _die( "size must be one of 1024x1024, 1536x1024, 1024x1536, or auto for this GPT Image model." ) def _validate_quality(quality: str) -> None: if quality not in ALLOWED_QUALITIES: _die("quality must be one of low, medium, high, or auto.") def _validate_background(background: Optional[str]) -> None: if background not in ALLOWED_BACKGROUNDS: _die("background must be one of transparent, opaque, or auto.") def _validate_model(model: str) -> None: if not model.startswith(GPT_IMAGE_MODEL_PREFIX): _die( "model must be a GPT Image model (for example gpt-image-1.5, gpt-image-1, or gpt-image-1-mini)." ) def _validate_transparency(background: Optional[str], output_format: str) -> None: if background == "transparent" and output_format not in {"png", "webp"}: _die("transparent background requires output-format png or webp.") def _validate_model_specific_options( *, model: str, background: Optional[str], ) -> None: if model != GPT_IMAGE_2_MODEL: return def _validate_generate_payload(payload: Dict[str, Any]) -> None: model = str(payload.get("model", DEFAULT_MODEL)) _validate_model(model) n = int(payload.get("n", 1)) if n < 1 or n > 10: _die("n must be between 1 and 10") size = str(payload.get("size", DEFAULT_SIZE)) quality = str(payload.get("quality", DEFAULT_QUALITY)) background = payload.get("background") _validate_size(size, model) _validate_quality(quality) _validate_background(background) _validate_model_specific_options(model=model, background=background) oc = payload.get("output_compression") if oc is not None and not (0 <= int(oc) <= 100): _die("output_compression must be between 0 and 100") def _build_output_paths( out: str, output_format: str, count: int, out_dir: Optional[str], ) -> List[Path]: ext = "." + output_format if out_dir: out_base = Path(out_dir) out_base.mkdir(parents=True, exist_ok=True) return [out_base / f"image_{i}{ext}" for i in range(1, count + 1)] out_path = Path(out) if out_path.exists() and out_path.is_dir(): out_path.mkdir(parents=True, exist_ok=True) return [out_path / f"image_{i}{ext}" for i in range(1, count + 1)] if out_path.suffix == "": out_path = out_path.with_suffix(ext) elif output_format and out_path.suffix.lstrip(".").lower() != output_format: _warn( f"Output extension {out_path.suffix} does not match output-format {output_format}." ) 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 _augment_prompt(args: argparse.Namespace, prompt: str) -> str: fields = _fields_from_args(args) return _augment_prompt_fields(args.augment, prompt, fields) def _augment_prompt_fields(augment: bool, prompt: str, fields: Dict[str, Optional[str]]) -> str: if not augment: return prompt sections: List[str] = [] if fields.get("use_case"): sections.append(f"Use case: {fields['use_case']}") sections.append(f"Primary request: {prompt}") if fields.get("scene"): sections.append(f"Scene/background: {fields['scene']}") if fields.get("subject"): sections.append(f"Subject: {fields['subject']}") if fields.get("style"): sections.append(f"Style/medium: {fields['style']}") if fields.get("composition"): sections.append(f"Composition/framing: {fields['composition']}") if fields.get("lighting"): sections.append(f"Lighting/mood: {fields['lighting']}") if fields.get("palette"): sections.append(f"Color palette: {fields['palette']}") if fields.get("materials"): sections.append(f"Materials/textures: {fields['materials']}") if fields.get("text"): sections.append(f"Text (verbatim): \"{fields['text']}\"") if fields.get("constraints"): sections.append(f"Constraints: {fields['constraints']}") if fields.get("negative"): sections.append(f"Avoid: {fields['negative']}") return "\n".join(sections) 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 _print_request(payload: dict) -> None: print(json.dumps(payload, indent=2, sort_keys=True)) def _decode_and_write(images: List[str], outputs: List[Path], force: bool) -> None: for idx, image_b64 in enumerate(images): if idx >= len(outputs): break 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(base64.b64decode(image_b64)) print(f"Wrote {out_path}") def _derive_downscale_path(path: Path, suffix: str) -> Path: if suffix and not suffix.startswith("-") 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 Exception: _die(f"Downscaling requires Pillow. {_dependency_hint('pillow')}") if max_dim < 1: _die("--downscale-max-dim must be >= 1") 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 == "jpg": fmt = "jpeg" if fmt == "jpeg": if resized.mode in ("RGBA", "LA") or ("transparency" in getattr(resized, "info", {})): bg = Image.new("RGB", resized.size, (255, 255, 255)) bg.paste(resized.convert("RGBA"), mask=resized.convert("RGBA").split()[-1]) resized = bg else: resized = resized.convert("RGB") out = BytesIO() resized.save(out, format=fmt.upper()) return out.getvalue() def _decode_write_and_downscale( images: List[str], outputs: List[Path], *, force: bool, downscale_max_dim: Optional[int], downscale_suffix: str, output_format: str, ) -> None: for idx, image_b64 in enumerate(images): if idx >= len(outputs): break 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) raw = base64.b64decode(image_b64) out_path.write_bytes(raw) print(f"Wrote {out_path}") if downscale_max_dim is None: continue derived = _derive_downscale_path(out_path, downscale_suffix) if derived.exists() and not force: _die(f"Output already exists: {derived} (use --force to overwrite)") derived.parent.mkdir(parents=True, exist_ok=True) resized = _downscale_image_bytes(raw, max_dim=downscale_max_dim, output_format=output_format) derived.write_bytes(resized) print(f"Wrote {derived}") def _create_client(*, api_key: str, base_url: str, timeout: float): try: from openai import OpenAI except ImportError: _die(f"openai SDK not installed in the active environment. {_dependency_hint('openai')}") return OpenAI(api_key=api_key, base_url=base_url, timeout=timeout) def _create_async_client(*, api_key: str, base_url: str, timeout: float): try: from openai import AsyncOpenAI except ImportError: try: import openai as _openai # noqa: F401 except ImportError: _die( f"openai SDK not installed in the active environment. {_dependency_hint('openai')}" ) _die( "AsyncOpenAI not available in this openai SDK version. " f"{_dependency_hint('openai', upgrade=True)}" ) return AsyncOpenAI(api_key=api_key, base_url=base_url, timeout=timeout) 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 "job" def _normalize_job(job: Any, idx: int) -> Dict[str, Any]: if isinstance(job, str): prompt = job.strip() if not prompt: _die(f"Empty prompt at job {idx}") return {"prompt": prompt} if isinstance(job, dict): if "prompt" not in job or not str(job["prompt"]).strip(): _die(f"Missing prompt for job {idx}") return job _die(f"Invalid job at index {idx}: expected string or object.") return {} # unreachable 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: item: Any if line.startswith("{"): item = json.loads(line) else: item = line jobs.append(_normalize_job(item, idx=line_no)) except json.JSONDecodeError as exc: _die(f"Invalid JSON on line {line_no}: {exc}") 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 def _merge_non_null(dst: Dict[str, Any], src: Dict[str, Any]) -> Dict[str, Any]: merged = dict(dst) for k, v in src.items(): if v is not None: merged[k] = v return merged def _job_output_paths( *, out_dir: Path, output_format: str, idx: int, prompt: str, n: int, explicit_out: Optional[str], ) -> List[Path]: out_dir.mkdir(parents=True, exist_ok=True) ext = "." + output_format if explicit_out: base = Path(explicit_out) if base.suffix == "": base = base.with_suffix(ext) elif base.suffix.lstrip(".").lower() != output_format: _warn( f"Job {idx}: output extension {base.suffix} does not match output-format {output_format}." ) base = out_dir / base.name else: slug = _slugify(prompt[:80]) base = out_dir / f"{idx:03d}-{slug}{ext}" if n == 1: return [base] return [ base.with_name(f"{base.stem}-{i}{base.suffix}") for i in range(1, n + 1) ] def _extract_retry_after_seconds(exc: Exception) -> Optional[float]: # Best-effort: openai SDK errors vary by version. Prefer a conservative fallback. for attr in ("retry_after", "retry_after_seconds"): val = getattr(exc, attr, None) if isinstance(val, (int, float)) and val >= 0: return float(val) msg = str(exc) m = re.search(r"retry[- ]after[:= ]+([0-9]+(?:\\.[0-9]+)?)", msg, re.IGNORECASE) if m: try: return float(m.group(1)) except Exception: return None return None def _is_rate_limit_error(exc: Exception) -> bool: name = exc.__class__.__name__.lower() if "ratelimit" in name or "rate_limit" in name: return True msg = str(exc).lower() return "429" in msg or "rate limit" in msg or "too many requests" in msg def _is_transient_error(exc: Exception) -> bool: if _is_rate_limit_error(exc): return True name = exc.__class__.__name__.lower() if any( marker in name for marker in ("timeout", "timedout", "tempor", "connection", "protocol") ): return True msg = str(exc).lower() return ( "timeout" in msg or "timed out" in msg or "connection error" in msg or "connection reset" in msg or "server disconnected" in msg or "remote protocol error" in msg or "503" in msg or "service unavailable" in msg or "no available compatible accounts" in msg ) def _format_transient_error(exc: Exception, *, attempts: int) -> str: msg = str(exc) if "no available compatible accounts" in msg.lower(): return ( "Novamailio returned 503: no available compatible image-generation accounts. " f"Retried {attempts} time(s). Try again later, lower concurrency, or switch to a Novamailio account/model with available image capacity." ) return f"Novamailio request failed after {attempts} attempt(s): {msg}" def _call_with_retries( call, *, attempts: int, retry_delay: float, label: str, ) -> Any: last_exc: Optional[Exception] = None for attempt in range(1, attempts + 1): try: return call() except Exception as exc: last_exc = exc if not _is_transient_error(exc): raise if attempt == attempts: _die(_format_transient_error(exc, attempts=attempts)) sleep_s = _extract_retry_after_seconds(exc) if sleep_s is None: sleep_s = retry_delay print( f"{label} attempt {attempt}/{attempts} failed ({exc}); retrying in {sleep_s:.1f}s", file=sys.stderr, ) time.sleep(sleep_s) raise last_exc or RuntimeError("unknown error") async def _generate_one_with_retries( client: Any, payload: Dict[str, Any], *, attempts: int, job_label: str, ) -> Any: last_exc: Optional[Exception] = None for attempt in range(1, attempts + 1): try: return await client.images.generate(**payload) except Exception as exc: last_exc = exc if not _is_transient_error(exc): raise if attempt == attempts: raise sleep_s = _extract_retry_after_seconds(exc) if sleep_s is None: sleep_s = min(60.0, 2.0**attempt) print( f"{job_label} attempt {attempt}/{attempts} failed ({exc.__class__.__name__}); retrying in {sleep_s:.1f}s", file=sys.stderr, ) await asyncio.sleep(sleep_s) raise last_exc or RuntimeError("unknown error") async def _run_generate_batch(args: argparse.Namespace) -> int: jobs = _read_jobs_jsonl(args.input) out_dir = Path(args.out_dir) base_fields = _fields_from_args(args) base_payload = { "model": args.model, "n": args.n, "size": args.size, "quality": args.quality, "background": args.background, "output_format": args.output_format, "output_compression": args.output_compression, "moderation": args.moderation, } if args.dry_run: for i, job in enumerate(jobs, start=1): prompt = str(job["prompt"]).strip() fields = _merge_non_null(base_fields, job.get("fields", {})) # Allow flat job keys as well (use_case, scene, etc.) fields = _merge_non_null(fields, {k: job.get(k) for k in base_fields.keys()}) augmented = _augment_prompt_fields(args.augment, prompt, fields) job_payload = dict(base_payload) job_payload["prompt"] = augmented job_payload = _merge_non_null(job_payload, {k: job.get(k) for k in base_payload.keys()}) job_payload = {k: v for k, v in job_payload.items() if v is not None} _validate_generate_payload(job_payload) effective_output_format = _normalize_output_format(job_payload.get("output_format")) _validate_transparency(job_payload.get("background"), effective_output_format) job_payload["output_format"] = effective_output_format n = int(job_payload.get("n", 1)) outputs = _job_output_paths( out_dir=out_dir, output_format=effective_output_format, idx=i, prompt=prompt, n=n, explicit_out=job.get("out"), ) downscaled = None if args.downscale_max_dim is not None: downscaled = [ str(_derive_downscale_path(p, args.downscale_suffix)) for p in outputs ] _print_request( { "provider": "novamailio", "base_url": args.base_url, "endpoint": f"{args.base_url}/images/generations", "job": i, "outputs": [str(p) for p in outputs], "outputs_downscaled": downscaled, **job_payload, } ) return 0 client = _create_async_client( api_key=args.api_key, base_url=args.base_url, timeout=args.timeout, ) sem = asyncio.Semaphore(args.concurrency) any_failed = False async def run_job(i: int, job: Dict[str, Any]) -> Tuple[int, Optional[str]]: nonlocal any_failed prompt = str(job["prompt"]).strip() job_label = f"[job {i}/{len(jobs)}]" fields = _merge_non_null(base_fields, job.get("fields", {})) fields = _merge_non_null(fields, {k: job.get(k) for k in base_fields.keys()}) augmented = _augment_prompt_fields(args.augment, prompt, fields) payload = dict(base_payload) payload["prompt"] = augmented payload = _merge_non_null(payload, {k: job.get(k) for k in base_payload.keys()}) payload = {k: v for k, v in payload.items() if v is not None} n = int(payload.get("n", 1)) _validate_generate_payload(payload) effective_output_format = _normalize_output_format(payload.get("output_format")) _validate_transparency(payload.get("background"), effective_output_format) payload["output_format"] = effective_output_format outputs = _job_output_paths( out_dir=out_dir, output_format=effective_output_format, idx=i, prompt=prompt, n=n, explicit_out=job.get("out"), ) try: async with sem: print(f"{job_label} starting", file=sys.stderr) started = time.time() result = await _generate_one_with_retries( client, payload, attempts=args.max_attempts, job_label=job_label, ) elapsed = time.time() - started print(f"{job_label} completed in {elapsed:.1f}s", file=sys.stderr) images = [item.b64_json for item in result.data] _decode_write_and_downscale( images, outputs, force=args.force, downscale_max_dim=args.downscale_max_dim, downscale_suffix=args.downscale_suffix, output_format=effective_output_format, ) return i, None except Exception as exc: any_failed = True print(f"{job_label} failed: {exc}", file=sys.stderr) if args.fail_fast: raise return i, str(exc) tasks = [asyncio.create_task(run_job(i, job)) for i, job in enumerate(jobs, start=1)] try: await asyncio.gather(*tasks) except Exception: for t in tasks: if not t.done(): t.cancel() raise return 1 if any_failed else 0 def _generate_batch(args: argparse.Namespace) -> None: exit_code = asyncio.run(_run_generate_batch(args)) if exit_code: raise SystemExit(exit_code) def _generate(args: argparse.Namespace) -> None: prompt = _read_prompt(args.prompt, args.prompt_file) prompt = _augment_prompt(args, prompt) payload = { "model": args.model, "prompt": prompt, "n": args.n, "size": args.size, "quality": args.quality, "background": args.background, "output_format": args.output_format, "output_compression": args.output_compression, "moderation": args.moderation, } payload = {k: v for k, v in payload.items() if v is not None} output_format = _normalize_output_format(args.output_format) _validate_transparency(args.background, output_format) payload["output_format"] = output_format output_paths = _build_output_paths(args.out, output_format, args.n, args.out_dir) downscaled = None if args.downscale_max_dim is not None: downscaled = [str(_derive_downscale_path(p, args.downscale_suffix)) for p in output_paths] if args.dry_run: _print_request( { "provider": "novamailio", "base_url": args.base_url, "endpoint": f"{args.base_url}/images/generations", "outputs": [str(p) for p in output_paths], "outputs_downscaled": downscaled, **payload, } ) return print( f"Calling Novamailio Image API (generation) at {args.base_url}. This can take up to {args.timeout / 60:.0f} minutes.", file=sys.stderr, ) started = time.time() client = _create_client(api_key=args.api_key, base_url=args.base_url, timeout=args.timeout) try: result = _call_with_retries( lambda: client.images.generate(**payload), attempts=args.max_attempts, retry_delay=args.retry_delay, label="Generation", ) except Exception as exc: if not _is_auth_error(exc) or not _refresh_api_key_after_auth_error(args, exc): raise client = _create_client(api_key=args.api_key, base_url=args.base_url, timeout=args.timeout) result = _call_with_retries( lambda: client.images.generate(**payload), attempts=args.max_attempts, retry_delay=args.retry_delay, label="Generation after key refresh", ) elapsed = time.time() - started print(f"Generation completed in {elapsed:.1f}s.", file=sys.stderr) images = [item.b64_json for item in result.data] _decode_write_and_downscale( images, output_paths, force=args.force, downscale_max_dim=args.downscale_max_dim, downscale_suffix=args.downscale_suffix, output_format=output_format, ) def _edit(args: argparse.Namespace) -> None: prompt = _read_prompt(args.prompt, args.prompt_file) prompt = _augment_prompt(args, prompt) image_paths = _check_image_paths(args.image) mask_path = Path(args.mask) if args.mask else None if mask_path: if not mask_path.exists(): _die(f"Mask file not found: {mask_path}") if mask_path.suffix.lower() != ".png": _warn(f"Mask should be a PNG with an alpha channel: {mask_path}") if mask_path.stat().st_size > MAX_IMAGE_BYTES: _warn(f"Mask exceeds 50MB limit: {mask_path}") payload = { "model": args.model, "prompt": prompt, "n": args.n, "size": args.size, "quality": args.quality, "background": args.background, "output_format": args.output_format, "output_compression": args.output_compression, "moderation": args.moderation, } payload = {k: v for k, v in payload.items() if v is not None} output_format = _normalize_output_format(args.output_format) _validate_transparency(args.background, output_format) payload["output_format"] = output_format output_paths = _build_output_paths(args.out, output_format, args.n, args.out_dir) downscaled = None if args.downscale_max_dim is not None: downscaled = [str(_derive_downscale_path(p, args.downscale_suffix)) for p in output_paths] if args.dry_run: payload_preview = dict(payload) payload_preview["image"] = [str(p) for p in image_paths] if mask_path: payload_preview["mask"] = str(mask_path) _print_request( { "provider": "novamailio", "base_url": args.base_url, "endpoint": f"{args.base_url}/images/edits", "outputs": [str(p) for p in output_paths], "outputs_downscaled": downscaled, **payload_preview, } ) return print( f"Calling Novamailio Image API (edit) at {args.base_url} with {len(image_paths)} image(s).", file=sys.stderr, ) started = time.time() client = _create_client(api_key=args.api_key, base_url=args.base_url, timeout=args.timeout) try: def _edit_once(): with _open_files(image_paths) as image_files, _open_mask(mask_path) as mask_file: request = dict(payload) request["image"] = image_files if len(image_files) > 1 else image_files[0] if mask_file is not None: request["mask"] = mask_file return client.images.edit(**request) result = _call_with_retries( _edit_once, attempts=args.max_attempts, retry_delay=args.retry_delay, label="Edit", ) except Exception as exc: if not _is_auth_error(exc) or not _refresh_api_key_after_auth_error(args, exc): raise client = _create_client(api_key=args.api_key, base_url=args.base_url, timeout=args.timeout) def _edit_after_key_refresh(): with _open_files(image_paths) as image_files, _open_mask(mask_path) as mask_file: request = dict(payload) request["image"] = image_files if len(image_files) > 1 else image_files[0] if mask_file is not None: request["mask"] = mask_file return client.images.edit(**request) result = _call_with_retries( _edit_after_key_refresh, attempts=args.max_attempts, retry_delay=args.retry_delay, label="Edit after key refresh", ) elapsed = time.time() - started print(f"Edit completed in {elapsed:.1f}s.", file=sys.stderr) images = [item.b64_json for item in result.data] _decode_write_and_downscale( images, output_paths, force=args.force, downscale_max_dim=args.downscale_max_dim, downscale_suffix=args.downscale_suffix, output_format=output_format, ) def _open_files(paths: List[Path]): return _FileBundle(paths) def _open_mask(mask_path: Optional[Path]): if mask_path is None: return _NullContext() return _SingleFile(mask_path) class _NullContext: def __enter__(self): return None def __exit__(self, exc_type, exc, tb): return False class _SingleFile: def __init__(self, path: Path): self._path = path self._handle = None def __enter__(self): self._handle = self._path.open("rb") return self._handle def __exit__(self, exc_type, exc, tb): if self._handle: try: self._handle.close() except Exception: pass return False class _FileBundle: def __init__(self, paths: List[Path]): self._paths = paths self._handles: List[object] = [] def __enter__(self): self._handles = [p.open("rb") for p in self._paths] return self._handles def __exit__(self, exc_type, exc, tb): for handle in self._handles: try: handle.close() except Exception: pass return False def _add_shared_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--base-url", default=os.getenv("NOVAMAILIO_BASE_URL", DEFAULT_BASE_URL), help="Novamailio OpenAI-compatible API root. Defaults to https://api.novamailio.com/v1.", ) parser.add_argument( "--api-key", help="Novamailio API key. On servers, prefer the NOVAMAILIO_API_KEY environment variable.", ) parser.add_argument("--model", default=DEFAULT_MODEL) parser.add_argument("--prompt") parser.add_argument("--prompt-file") parser.add_argument("--n", type=int, default=1) parser.add_argument("--size", default=DEFAULT_SIZE) parser.add_argument("--quality", default=DEFAULT_QUALITY) parser.add_argument("--background") parser.add_argument("--output-format") parser.add_argument("--output-compression", type=int) parser.add_argument("--moderation") parser.add_argument("--out", default=DEFAULT_OUTPUT_PATH) parser.add_argument("--out-dir") parser.add_argument("--force", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--augment", dest="augment", action="store_true") parser.add_argument("--no-augment", dest="augment", action="store_false") parser.set_defaults(augment=True) # Prompt augmentation hints parser.add_argument("--use-case") parser.add_argument("--scene") parser.add_argument("--subject") parser.add_argument("--style") parser.add_argument("--composition") parser.add_argument("--lighting") parser.add_argument("--palette") parser.add_argument("--materials") parser.add_argument("--text") parser.add_argument("--constraints") parser.add_argument("--negative") # Post-processing (optional): generate an additional downscaled copy for fast web loading. parser.add_argument("--downscale-max-dim", type=int) parser.add_argument("--downscale-suffix", default=DEFAULT_DOWNSCALE_SUFFIX) parser.add_argument( "--max-attempts", type=int, default=DEFAULT_MAX_ATTEMPTS, help="Maximum attempts for a single generate/edit request when Novamailio returns transient errors.", ) parser.add_argument( "--retry-delay", type=float, default=DEFAULT_RETRY_DELAY, help="Seconds to wait between transient generate/edit retries when the server does not provide Retry-After.", ) parser.add_argument( "--timeout", type=float, default=DEFAULT_TIMEOUT, help="Maximum seconds to wait for each Novamailio API request. Defaults to 600 seconds (10 minutes).", ) def main() -> int: parser = argparse.ArgumentParser( description="Generate or edit images through Novamailio's OpenAI-compatible Image API" ) subparsers = parser.add_subparsers(dest="command", required=True) gen_parser = subparsers.add_parser("generate", help="Create a new image") _add_shared_args(gen_parser) gen_parser.set_defaults(func=_generate) batch_parser = subparsers.add_parser( "generate-batch", help="Generate multiple prompts concurrently (JSONL input)", ) _add_shared_args(batch_parser) batch_parser.add_argument("--input", required=True, help="Path to JSONL file (one job per line)") batch_parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY) batch_parser.add_argument("--fail-fast", action="store_true") batch_parser.set_defaults(func=_generate_batch) edit_parser = subparsers.add_parser("edit", help="Edit an existing image") _add_shared_args(edit_parser) edit_parser.add_argument("--image", action="append", required=True) edit_parser.add_argument("--mask") edit_parser.set_defaults(func=_edit) args = parser.parse_args() if args.n < 1 or args.n > 10: _die("--n must be between 1 and 10") if getattr(args, "concurrency", 1) < 1 or getattr(args, "concurrency", 1) > 25: _die("--concurrency must be between 1 and 25") if getattr(args, "max_attempts", 3) < 1 or getattr(args, "max_attempts", 3) > 10: _die("--max-attempts must be between 1 and 10") if getattr(args, "retry_delay", DEFAULT_RETRY_DELAY) < 0: _die("--retry-delay must be >= 0") if getattr(args, "timeout", DEFAULT_TIMEOUT) <= 0: _die("--timeout must be > 0") if args.output_compression is not None and not (0 <= args.output_compression <= 100): _die("--output-compression must be between 0 and 100") if args.command == "generate-batch" and not args.out_dir: _die("generate-batch requires --out-dir") if getattr(args, "downscale_max_dim", None) is not None and args.downscale_max_dim < 1: _die("--downscale-max-dim must be >= 1") _validate_model(args.model) _validate_size(args.size, args.model) _validate_quality(args.quality) _validate_background(args.background) _validate_model_specific_options( model=args.model, background=args.background, ) args.base_url = _normalize_base_url(args.base_url) args.api_key = _resolve_api_key( api_key=args.api_key, dry_run=args.dry_run, ) args.func(args) return 0 if __name__ == "__main__": raise SystemExit(main())