"""xAI image generation backend. Exposes xAI's ``grok-imagine-image`` models as an :class:`ImageGenProvider`: text-to-image (``/v1/images/generations``, several aspect ratios, 1K/2K) and image editing (``/v1/images/edits``), base64 output saved to cache. Selection precedence (first hit wins): ``model`` kwarg → ``XAI_IMAGE_MODEL`` env → ``image_gen.xai.model`` → :data:`DEFAULT_MODEL`. """ from __future__ import annotations import logging import os from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import requests from agent.image_gen_provider import ( DEFAULT_ASPECT_RATIO, ImageGenProvider, resolve_aspect_ratio, success_response, ) from plugins.image_gen._common import ( catalog_rows, collect_source_images, error_factory, load_image_gen_config, materialize_image, post_json, ) from tools.xai_http import ( build_xai_storage_options, hermes_xai_user_agent, maybe_mark_xai_storage_notice_seen, read_xai_imagine_storage_config, resolve_xai_http_credentials, xai_storage_notice_text, ) logger = logging.getLogger(__name__) _MODELS: Dict[str, Dict[str, Any]] = { "grok-imagine-image": { "display": "Grok Imagine Image", "speed": "~5-10s", "strengths": "Fast, high-quality", }, "grok-imagine-image-2.0": { "display": "Grok Imagine Image 2.0", "speed": "~10-20s", "strengths": "Typography/layout-aware; legible small text; strongest quality.", }, "grok-imagine-image-quality": { "display": "Grok Imagine Image (Quality)", "speed": "~10-20s", "strengths": "Higher fidelity / detail; slower than the standard model.", }, } DEFAULT_MODEL = "grok-imagine-image" # xAI documents the quality model as the edit-capable baseline. _EDIT_FALLBACK_MODEL = "grok-imagine-image-quality" # Live catalog cache: (models_dict, fetched_monotonic). xAI's # ``/image-generation-models`` endpoint is the source of truth so newly # released Imagine models appear in the picker without a code change; the # static ``_MODELS`` table is the offline fallback and supplies curated # speed/strengths text for the models we know about. _LIVE_CACHE: Optional[Tuple[Dict[str, Dict[str, Any]], float]] = None _LIVE_CACHE_TTL = 300.0 _LIVE_TIMEOUT = 10.0 # xAI aspect ratios (more options than FAL/OpenAI) _XAI_ASPECT_RATIOS = { "landscape": "16:9", "square": "1:1", "portrait": "9:16", "4:3": "4:3", "3:4": "3:4", "3:2": "3:2", "2:3": "2:3", } _XAI_RESOLUTIONS = {"1k", "2k"} DEFAULT_RESOLUTION = "1k" _MAX_SOURCE_IMAGES = 3 _REQUEST_TIMEOUT = 120 def _base_url(creds: Dict[str, Any]) -> str: return str(creds.get("base_url") or "https://api.x.ai/v1").strip().rstrip("/") def _fetch_live_models() -> Dict[str, Dict[str, Any]]: """Fetch image models from xAI's ``/image-generation-models`` endpoint. Returns ``{model_id: {"input_modalities": [...], "aliases": [...]}}``. Raises on any failure — callers treat that as "use the static table". """ creds = resolve_xai_http_credentials() api_key = str(creds.get("api_key") or "").strip() if not api_key: raise RuntimeError("no xAI credentials") response = requests.get( f"{_base_url(creds)}/image-generation-models", headers={"Authorization": f"Bearer {api_key}", "User-Agent": hermes_xai_user_agent()}, timeout=_LIVE_TIMEOUT, ) response.raise_for_status() payload = response.json() entries = payload.get("models") or payload.get("data") or [] out: Dict[str, Dict[str, Any]] = {} for entry in entries: if not isinstance(entry, dict): continue model_id = entry.get("id") or entry.get("name") if not isinstance(model_id, str) or not model_id.strip(): continue out[model_id.strip()] = { "input_modalities": entry.get("input_modalities") or [], "aliases": entry.get("aliases") or [], } return out def _live_models() -> Dict[str, Dict[str, Any]]: """Cached live catalog (``{}`` when unreachable).""" global _LIVE_CACHE import time if _LIVE_CACHE is not None and time.monotonic() - _LIVE_CACHE[1] < _LIVE_CACHE_TTL: return _LIVE_CACHE[0] try: live = _fetch_live_models() except Exception as exc: # noqa: BLE001 - offline/unauth → static fallback logger.debug("xAI live image model catalog unavailable: %s", exc) live = {} _LIVE_CACHE = (live, time.monotonic()) return live def _catalog() -> Dict[str, Dict[str, Any]]: """Merged model catalog: live endpoint IDs + curated static metadata. Known models keep their curated display/speed/strengths; models xAI ships later still show up (with generic metadata) so users can pick them the day they launch. Curated entries the live list momentarily omits are kept; static table alone when the API is unreachable. """ live = _live_models() if not live: return dict(_MODELS) merged: Dict[str, Dict[str, Any]] = {} for model_id in live: meta = _MODELS.get(model_id) or { "display": model_id, "speed": "", "strengths": "New xAI Imagine model (from live xAI catalog)", } merged[model_id] = dict(meta) merged[model_id]["input_modalities"] = live[model_id].get("input_modalities") or [] for model_id, meta in _MODELS.items(): merged.setdefault(model_id, dict(meta)) return merged def _load_xai_config() -> Dict[str, Any]: return load_image_gen_config("xai") def _configured_model() -> Optional[str]: value = _load_xai_config().get("model") return value if isinstance(value, str) else None def _resolve_model(caller_model: Optional[str] = None) -> Tuple[str, Dict[str, Any]]: """Return ``(model_id, meta)``: caller kwarg → ``XAI_IMAGE_MODEL`` → config → default. Every candidate is validated against the merged live+static catalog, so a newly released xAI model is selectable the day it appears in the live catalog. """ catalog = _catalog() for candidate in (caller_model, os.environ.get("XAI_IMAGE_MODEL"), _configured_model()): if candidate and candidate in catalog: return candidate, catalog[candidate] return DEFAULT_MODEL, catalog.get(DEFAULT_MODEL, _MODELS[DEFAULT_MODEL]) def _resolve_edit_model(caller_model: Optional[str] = None) -> str: """Model for ``/v1/images/edits``: an explicitly selected model that accepts image input is honored; otherwise the documented quality baseline.""" catalog = _catalog() explicit = caller_model or os.environ.get("XAI_IMAGE_MODEL") or _configured_model() if explicit and explicit in catalog: if "image" in (catalog[explicit].get("input_modalities") or []): return explicit return _EDIT_FALLBACK_MODEL def _resolve_resolution() -> str: res = _load_xai_config().get("resolution") return res if isinstance(res, str) and res in _XAI_RESOLUTIONS else DEFAULT_RESOLUTION def _xai_image_field(source: str) -> Dict[str, str]: """Build the xAI ``image`` field for an edit request. ``/v1/images/edits`` accepts a public HTTPS URL or a base64 data URI; local file paths are read and encoded into a ``data:`` URI. """ source = source.strip() if source.lower().startswith(("http://", "https://", "data:")): return {"url": source, "type": "image_url"} import base64 # Shared credential-read guard before reading local bytes (same boundary # the OpenAI / OpenRouter / Codex image providers apply). from agent.file_safety import raise_if_read_blocked raise_if_read_blocked(source) with open(os.path.expanduser(source), "rb") as fh: # windows-footgun: ok raw = fh.read() ext = (os.path.splitext(source)[1].lstrip(".") or "png").lower() if ext == "jpg": ext = "jpeg" b64 = base64.b64encode(raw).decode("utf-8") return {"url": f"data:image/{ext};base64,{b64}", "type": "image_url"} class XAIImageGenProvider(ImageGenProvider): """xAI ``grok-imagine-image`` backend.""" @property def name(self) -> str: return "xai" @property def display_name(self) -> str: return "xAI (Grok)" def is_available(self) -> bool: return bool(resolve_xai_http_credentials().get("api_key")) def list_models(self) -> List[Dict[str, Any]]: return catalog_rows(_catalog(), ("display", "speed", "strengths")) def get_setup_schema(self) -> Dict[str, Any]: # Auth resolution is delegated to the shared ``xai_grok`` post_setup hook # (``hermes_cli/tools_config.py``), identical to the TTS / video-gen # entries, so every xAI service shows the same OAuth-or-API-key choice. storage_notice = xai_storage_notice_text("image_gen") tag = ( "grok-imagine-image - text-to-image & image editing; uses xAI " "Grok OAuth or XAI_API_KEY" ) if storage_notice: tag += f". {storage_notice}" return { "name": "xAI Grok Imagine (image)", "badge": "paid", "tag": tag, "env_vars": [], "post_setup": "xai_grok", } def capabilities(self) -> Dict[str, Any]: # /v1/images/edits accepts up to 3 total source images. return { "modalities": ["text", "image"], "max_reference_images": 2, "max_source_images": _MAX_SOURCE_IMAGES, } def generate( self, prompt: str, aspect_ratio: str = DEFAULT_ASPECT_RATIO, *, image_url: Optional[str] = None, reference_image_urls: Optional[List[str]] = None, **kwargs: Any, ) -> Dict[str, Any]: """Text-to-image via ``/v1/images/generations``, or image editing via ``/v1/images/edits`` when source images are supplied. Editing uses a JSON body (xAI does not support the OpenAI SDK's multipart ``images.edit()``) and an image-input-capable model. """ creds = resolve_xai_http_credentials() api_key = str(creds.get("api_key") or "").strip() provider_name = str(creds.get("provider") or "xai").strip() or "xai" if not api_key: return error_factory(provider_name, aspect_ratio)( "No xAI credentials found. Configure xAI OAuth in `hermes model` or set XAI_API_KEY.", "missing_api_key", ) model_id, meta = _resolve_model(kwargs.get("model")) aspect = resolve_aspect_ratio(aspect_ratio) xai_res = _resolve_resolution() source_images = collect_source_images(image_url, reference_image_urls) edit_fail = error_factory(provider_name, aspect, model=_EDIT_FALLBACK_MODEL, prompt=prompt) if len(source_images) > _MAX_SOURCE_IMAGES: return edit_fail( f"xAI image editing supports at most {_MAX_SOURCE_IMAGES} source images", "too_many_references", ) for index, source in enumerate(source_images): field = "image_url" if index == 0 and image_url and image_url.strip() == source else "reference_image_urls" if not source.lower().startswith(("http://", "https://", "data:")): if not Path(source).expanduser().is_file(): return edit_fail( f"{field} must be a public HTTPS URL or data URI " "(e.g. the `image`/`public_url` from a prior Imagine result)", "invalid_image_url", ) is_edit = bool(source_images) headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "User-Agent": hermes_xai_user_agent(), } base_url = _base_url(creds) storage_options = build_xai_storage_options( "image_gen", filename_prefix="hermes-xai-image", extension="png", ) storage_notice = maybe_mark_xai_storage_notice_seen("image_gen") storage_cfg = read_xai_imagine_storage_config("image_gen") if is_edit: # An explicit user selection that accepts image input (e.g. # grok-imagine-image-2.0) is honored; otherwise the quality baseline. model_id = _resolve_edit_model(kwargs.get("model")) try: image_fields = [_xai_image_field(source) for source in source_images] except Exception as exc: return edit_fail( f"Could not load source image for editing: {exc}", "io_error", model=model_id, ) payload: Dict[str, Any] = {"model": model_id, "prompt": prompt} if len(image_fields) == 1: payload["image"] = image_fields[0] else: payload["images"] = image_fields endpoint_url = f"{base_url}/images/edits" else: payload = { "model": model_id, "prompt": prompt, "aspect_ratio": _XAI_ASPECT_RATIOS.get(aspect, "1:1"), "resolution": xai_res, } endpoint_url = f"{base_url}/images/generations" if storage_options is not None: payload["storage_options"] = storage_options fail = error_factory(provider_name, aspect, model=model_id, prompt=prompt) result, failure = post_json( endpoint_url, headers=headers, payload=payload, timeout=_REQUEST_TIMEOUT, label="xAI", ) if failure: if failure.kind == "http": logger.error("xAI image gen failed (%d): %s", failure.status, failure.message) return fail(failure.error, failure.error_type) # xAI returns data[0].b64_json, data[0].url, and optionally # data[0].file_output when storage_options were requested. data = result.get("data", []) if not data: return fail("xAI returned no image data", "empty_response") first = data[0] file_output = first.get("file_output") if isinstance(first, dict) else None file_output = file_output if isinstance(file_output, dict) else {} public_url = file_output.get("public_url") if isinstance(file_output.get("public_url"), str) else None if public_url: image_ref = public_url else: # grok-imagine-image URLs (``imgen.x.ai/xai-tmp-*``) 404 within # minutes; materialise locally so downstream consumers (Telegram # send_photo …) get a stable path. image_ref, err = materialize_image( first.get("b64_json"), first.get("url"), prefix=f"xai_{model_id}", label="xAI", provider="xai", model=model_id, prompt=prompt, aspect=aspect, log=logger, ) if err: return err extra: Dict[str, Any] = {"storage_enabled": bool(storage_cfg["enabled"])} if not is_edit: extra["resolution"] = xai_res if storage_notice: extra["storage_notice"] = storage_notice if public_url: extra["public_url"] = public_url for key in ( "filename", "expires_at", "public_url_expires_at", "public_url_error", "storage_error", ): if key in file_output: extra[key] = file_output[key] if result.get("usage"): extra["usage"] = result["usage"] return success_response( image=image_ref, model=model_id, prompt=prompt, aspect_ratio=aspect, provider="xai", modality="image" if is_edit else "text", extra=extra, ) def register(ctx: Any) -> None: """Register this provider with the image gen registry.""" ctx.register_image_gen_provider(XAIImageGenProvider())