Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

plugins/image_gen/deepinfra/__init__.py

@@ -0,0 +1,336 @@
+"""DeepInfra image generation backend.
+
+Exposes DeepInfra's image-gen catalog (FLUX, Qwen-Image-Edit, …) through
+the OpenAI-compatible ``/v1/openai/images/generations`` endpoint as an
+:class:`ImageGenProvider` implementation.
+
+**Fully dynamic model discovery.** Unlike the other image-gen plugins in
+this tree (which ship a hardcoded ``_MODELS`` dict), DeepInfra publishes
+a single tagged catalog at
+``https://api.deepinfra.com/v1/openai/models?filter=true&sort_by=hermes``
+where each entry's ``metadata.tags`` declares its surface (``image-gen``
+here). ``list_models()`` filters that catalog via
+:func:`hermes_cli.models._fetch_deepinfra_models_by_tag` so newly added
+models show up in ``hermes tools`` automatically. No model ids are
+hardcoded in this file — if a model is retired upstream, it disappears
+from hermes the next time the catalog is fetched, no patch required.
+
+Model selection (first hit wins):
+
+1. ``DEEPINFRA_IMAGE_MODEL`` env var
+2. ``image_gen.deepinfra.model`` in ``config.yaml``
+3. First model from the live catalog
+
+When all three are absent (catalog unreachable, nothing configured),
+``generate()`` returns an :func:`error_response` rather than guessing.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+
+from agent.secret_scope import get_secret
+from agent.image_gen_provider import (
+    DEFAULT_ASPECT_RATIO,
+    ImageGenProvider,
+    error_response,
+    resolve_aspect_ratio,
+    save_b64_image,
+    save_url_image,
+    success_response,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# DeepInfra accepts standard OpenAI ``size`` strings. Mirrors the
+# OpenAI plugin's mapping so aspect_ratio semantics stay consistent
+# across the agent's image_generate tool surface.
+_SIZES = {
+    "landscape": "1536x1024",
+    "square": "1024x1024",
+    "portrait": "1024x1536",
+}
+
+
+def _load_deepinfra_image_config() -> Dict[str, Any]:
+    """Read ``image_gen.deepinfra`` from config.yaml."""
+    try:
+        from hermes_cli.config import load_config
+
+        cfg = load_config()
+        section = cfg.get("image_gen") if isinstance(cfg, dict) else None
+        di_section = section.get("deepinfra") if isinstance(section, dict) else None
+        return di_section if isinstance(di_section, dict) else {}
+    except Exception as exc:
+        logger.debug("Could not load image_gen.deepinfra config: %s", exc)
+        return {}
+
+
+def _live_models() -> Optional[List[Dict[str, Any]]]:
+    """Fetch ``image-gen``-tagged models from the DeepInfra catalog."""
+    try:
+        from hermes_cli.models import _fetch_deepinfra_models_by_tag
+    except Exception as exc:
+        logger.debug("Cannot import _fetch_deepinfra_models_by_tag: %s", exc)
+        return None
+    return _fetch_deepinfra_models_by_tag("image-gen")
+
+
+def _format_catalog_row(item: Dict[str, Any]) -> Dict[str, Any]:
+    """Format a catalog item into the picker row shape."""
+    mid = item.get("id", "")
+    metadata = item.get("metadata") or {}
+    pricing = metadata.get("pricing") if isinstance(metadata, dict) else None
+    price = ""
+    if isinstance(pricing, dict) and pricing.get("per_image_unit") is not None:
+        try:
+            price = f"${float(pricing['per_image_unit']):.4f}/image"
+        except (TypeError, ValueError):
+            price = ""
+    row: Dict[str, Any] = {
+        "id": mid,
+        "display": mid.split("/", 1)[-1] if "/" in mid else mid,
+        "strengths": metadata.get("description", "") if isinstance(metadata, dict) else "",
+    }
+    if price:
+        row["price"] = price
+    if isinstance(metadata, dict):
+        for key in ("default_width", "default_height", "default_iterations"):
+            if metadata.get(key) is not None:
+                row[key] = metadata[key]
+    return row
+
+
+def _resolve_model(catalog: List[Dict[str, Any]], cfg: Dict[str, Any]) -> Optional[str]:
+    """Pick the model id (env > config > first live result, else None).
+
+    Takes the already-loaded ``image_gen.deepinfra`` config so ``generate()``
+    reads config once instead of via a second ``load_config`` deepcopy.
+    """
+    env_override = os.environ.get("DEEPINFRA_IMAGE_MODEL", "").strip()
+    if env_override:
+        return env_override
+    cfg_model = cfg.get("model") if isinstance(cfg, dict) else None
+    if isinstance(cfg_model, str) and cfg_model.strip():
+        return cfg_model.strip()
+    if catalog:
+        first = catalog[0].get("id")
+        if isinstance(first, str) and first:
+            return first
+    return None
+
+
+class DeepInfraImageGenProvider(ImageGenProvider):
+    """DeepInfra ``images.generations`` backend.
+
+    Catalog is discovered live from the DeepInfra ``/models`` endpoint
+    filtered by the ``image-gen`` surface tag.
+    """
+
+    @property
+    def name(self) -> str:
+        return "deepinfra"
+
+    @property
+    def display_name(self) -> str:
+        return "DeepInfra"
+
+    def is_available(self) -> bool:
+        return bool((get_secret("DEEPINFRA_API_KEY", "") or "").strip())
+
+    def list_models(self) -> List[Dict[str, Any]]:
+        live = _live_models()
+        if not live:
+            return []
+        return [_format_catalog_row(item) for item in live]
+
+    def default_model(self) -> Optional[str]:
+        rows = self.list_models()
+        if rows:
+            return rows[0].get("id")
+        return None
+
+    def capabilities(self) -> Dict[str, Any]:
+        """DeepInfra's OpenAI-compatible generation surface is text-only."""
+        return {"modalities": ["text"], "max_reference_images": 0}
+
+    def get_setup_schema(self) -> Dict[str, Any]:
+        return {
+            "name": "DeepInfra",
+            "badge": "paid",
+            "tag": "FLUX, Qwen-Image, … — live catalog from api.deepinfra.com",
+            "env_vars": [
+                {
+                    "key": "DEEPINFRA_API_KEY",
+                    "prompt": "DeepInfra API key",
+                    "url": "https://deepinfra.com/dash/api_keys",
+                },
+            ],
+        }
+
+    def generate(
+        self,
+        prompt: str,
+        aspect_ratio: str = DEFAULT_ASPECT_RATIO,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        prompt = (prompt or "").strip()
+        aspect = resolve_aspect_ratio(aspect_ratio)
+
+        if kwargs.get("image_url") or kwargs.get("reference_image_urls"):
+            return error_response(
+                error=(
+                    "DeepInfra image generation is text-to-image only in this "
+                    "backend; image_url and reference_image_urls are unsupported."
+                ),
+                error_type="modality_unsupported",
+                provider="deepinfra",
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+
+        if not prompt:
+            return error_response(
+                error="Prompt is required and must be a non-empty string",
+                error_type="invalid_argument",
+                provider="deepinfra",
+                aspect_ratio=aspect,
+            )
+
+        api_key = (get_secret("DEEPINFRA_API_KEY", "") or "").strip()
+        if not api_key:
+            return error_response(
+                error=(
+                    "DEEPINFRA_API_KEY not set. Run `hermes tools` → Image "
+                    "Generation → DeepInfra to configure, or `hermes setup` "
+                    "to add the key."
+                ),
+                error_type="auth_required",
+                provider="deepinfra",
+                aspect_ratio=aspect,
+            )
+
+        di_cfg = _load_deepinfra_image_config()
+        catalog = _live_models() or []
+        model_id = _resolve_model(catalog, di_cfg)
+        if not model_id:
+            return error_response(
+                error=(
+                    "No DeepInfra image-gen model available. Pin one in "
+                    "config.yaml under image_gen.deepinfra.model, set "
+                    "DEEPINFRA_IMAGE_MODEL, or check connectivity to "
+                    "api.deepinfra.com so the live catalog can be fetched."
+                ),
+                error_type="no_model_available",
+                provider="deepinfra",
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+        size = _SIZES.get(aspect, _SIZES["square"])
+        from hermes_cli.models import deepinfra_base_url
+        base_url = deepinfra_base_url(di_cfg)
+
+        # DeepInfra's /images/generations is OpenAI-compatible — use the
+        # openai SDK so we inherit its retry, timeout, and error mapping
+        # (mirrors the existing OpenAI image-gen plugin).
+        try:
+            import openai
+        except ImportError:
+            return error_response(
+                error="openai Python package not installed (pip install openai)",
+                error_type="missing_dependency",
+                provider="deepinfra",
+                aspect_ratio=aspect,
+            )
+
+        client = openai.OpenAI(api_key=api_key, base_url=base_url)
+        try:
+            response = client.images.generate(
+                model=model_id,
+                prompt=prompt,
+                size=size,
+                n=1,
+            )
+        except Exception as exc:
+            logger.debug("DeepInfra image generation failed", exc_info=True)
+            return error_response(
+                error=f"DeepInfra image generation failed: {exc}",
+                error_type="api_error",
+                provider="deepinfra",
+                model=model_id,
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+        finally:
+            close = getattr(client, "close", None)
+            if callable(close):
+                close()
+
+        data = getattr(response, "data", None) or []
+        if not data:
+            return error_response(
+                error="DeepInfra returned no image data",
+                error_type="empty_response",
+                provider="deepinfra",
+                model=model_id,
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+
+        first = data[0]
+        b64 = getattr(first, "b64_json", None)
+        url = getattr(first, "url", None)
+
+        # Drop the ``vendor/`` prefix and any colons so the saved filename
+        # stays a single path component on every OS.
+        short = model_id.split("/", 1)[-1].replace(":", "_")
+
+        if b64:
+            try:
+                saved_path = save_b64_image(b64, prefix=f"deepinfra_{short}")
+            except Exception as exc:
+                return error_response(
+                    error=f"Could not save image to cache: {exc}",
+                    error_type="io_error",
+                    provider="deepinfra",
+                    model=model_id,
+                    prompt=prompt,
+                    aspect_ratio=aspect,
+                )
+            image_ref = str(saved_path)
+        elif url:
+            # Materialise the (often short-lived) delivery URL locally so a
+            # downstream consumer (Telegram send_photo, browser fetch) doesn't
+            # get a dead link — mirrors the openai/xai/krea image plugins.
+            # Best-effort: fall back to the bare URL if the download fails.
+            try:
+                image_ref = str(save_url_image(url, prefix=f"deepinfra_{short}"))
+            except Exception as exc:
+                logger.debug("DeepInfra: caching delivery URL failed (%s); returning URL", exc)
+                image_ref = url
+        else:
+            return error_response(
+                error="DeepInfra response contained neither b64_json nor URL",
+                error_type="empty_response",
+                provider="deepinfra",
+                model=model_id,
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+
+        return success_response(
+            image=image_ref,
+            model=model_id,
+            prompt=prompt,
+            aspect_ratio=aspect,
+            provider="deepinfra",
+            extra={"size": size},
+        )
+
+
+def register(ctx) -> None:
+    """Plugin entry point — wire ``DeepInfraImageGenProvider`` into the registry."""
+    ctx.register_image_gen_provider(DeepInfraImageGenProvider())