diff --git a/plugins/image_gen/krea/__init__.py b/plugins/image_gen/krea/__init__.py index a897302175b..02565d237c1 100644 --- a/plugins/image_gen/krea/__init__.py +++ b/plugins/image_gen/krea/__init__.py @@ -25,6 +25,7 @@ from __future__ import annotations import logging import os import time +import uuid from typing import Any, Dict, List, Optional, Tuple import requests @@ -63,6 +64,13 @@ _MODELS: Dict[str, Dict[str, Any]] = { "price": "$0.060 (text) / $0.065 (style refs) / $0.070 (moodboards)", "path": "large", }, + "krea-2-medium-turbo": { + "display": "Krea 2 Medium Turbo", + "speed": "~8-15s", + "strengths": "Fastest Krea 2 — medium quality at lower latency / cost.", + "price": "$0.015 (text) / $0.0175 (style refs)", + "path": "medium-turbo", + }, } DEFAULT_MODEL = "krea-2-medium" @@ -116,8 +124,16 @@ def _load_krea_config() -> Dict[str, Any]: return {} -def _resolve_model() -> Tuple[str, Dict[str, Any]]: - """Decide which model to use and return ``(model_id, meta)``.""" +def _resolve_model(explicit: Optional[str] = None) -> Tuple[str, Dict[str, Any]]: + """Decide which model to use and return ``(model_id, meta)``. + + Precedence: explicit caller override (e.g. managed-mode routing or a direct + ``model`` kwarg) → ``KREA_IMAGE_MODEL`` env → ``image_gen.krea.model`` → + ``image_gen.model`` → :data:`DEFAULT_MODEL`. + """ + if isinstance(explicit, str) and explicit.strip() in _MODELS: + return explicit.strip(), _MODELS[explicit.strip()] + env_override = os.environ.get("KREA_IMAGE_MODEL") if env_override and env_override in _MODELS: return env_override, _MODELS[env_override] @@ -140,6 +156,44 @@ def _resolve_model() -> Tuple[str, Dict[str, Any]]: return DEFAULT_MODEL, _MODELS[DEFAULT_MODEL] +def _resolve_managed_krea_gateway(): + """Return managed Krea gateway config when the user is on the managed path. + + Mirrors ``_resolve_managed_fal_gateway`` in ``tools/image_generation_tool.py``: + the Nous-hosted Krea gateway wins when it is resolvable AND either no direct + ``KREA_API_KEY`` is configured or the user explicitly opted into the gateway + for ``image_gen``. Returns ``None`` (direct/BYO path) otherwise, and never + raises — plugin discovery and availability scans must stay robust. + """ + try: + from tools.managed_tool_gateway import resolve_managed_tool_gateway + from tools.tool_backend_helpers import prefers_gateway + except Exception as exc: # noqa: BLE001 + logger.debug("Managed Krea gateway resolution unavailable: %s", exc) + return None + + if os.environ.get("KREA_API_KEY") and not prefers_gateway("image_gen"): + return None + + try: + return resolve_managed_tool_gateway("krea") + except Exception as exc: # noqa: BLE001 + logger.debug("Managed Krea gateway resolution failed: %s", exc) + return None + + +def _managed_krea_gateway_ready() -> bool: + """Cheap, offline-friendly probe for managed Krea availability.""" + try: + from tools.managed_tool_gateway import is_managed_tool_gateway_ready + except Exception: # noqa: BLE001 + return False + try: + return bool(is_managed_tool_gateway_ready("krea")) + except Exception: # noqa: BLE001 + return False + + def _resolve_creativity(value: Optional[str]) -> str: """Coerce ``creativity`` kwarg to a valid Krea value (default ``medium``).""" if isinstance(value, str): @@ -171,7 +225,10 @@ class KreaImageGenProvider(ImageGenProvider): return "Krea" def is_available(self) -> bool: - return bool(os.environ.get("KREA_API_KEY")) + # Available with a direct Krea key OR via the managed Nous gateway + # (Nous Subscription), so portal users with no Krea key can still + # reach Krea 2 through the gateway. + return bool(os.environ.get("KREA_API_KEY")) or _managed_krea_gateway_ready() def list_models(self) -> List[Dict[str, Any]]: return [ @@ -192,7 +249,7 @@ class KreaImageGenProvider(ImageGenProvider): return { "name": "Krea", "badge": "paid", - "tag": "Krea 2 foundation model — Medium ($0.03) + Large ($0.06). Style transfer, moodboards, reference-guided generation.", + "tag": "Krea 2 foundation model — Medium ($0.03), Large ($0.06), Medium Turbo ($0.015). Style transfer, moodboards, reference-guided generation. Direct key or managed Nous Subscription gateway.", "env_vars": [ { "key": "KREA_API_KEY", @@ -265,22 +322,67 @@ class KreaImageGenProvider(ImageGenProvider): aspect_ratio=aspect, ) - api_key = os.environ.get("KREA_API_KEY") - if not api_key: - return error_response( - error=( - "KREA_API_KEY not set. Run `hermes tools` → Image " - "Generation → Krea to configure, or get a key at " - "https://www.krea.ai/settings/api-tokens." - ), - error_type="auth_required", - provider="krea", - aspect_ratio=aspect, - ) + # Route through the managed Nous gateway (Nous Subscription) when the + # user is on the managed path; otherwise use the direct Krea API with a + # BYO ``KREA_API_KEY``. The gateway owns the shared Krea credential and + # meters/bills per generation, so the caller token is the Nous access + # token, not a Krea key. + managed = _resolve_managed_krea_gateway() + if managed is not None: + base_url = managed.gateway_origin.rstrip("/") + auth_token = managed.nous_user_token + else: + base_url = BASE_URL + auth_token = os.environ.get("KREA_API_KEY") + if not auth_token: + return error_response( + error=( + "KREA_API_KEY not set. Run `hermes tools` → Image " + "Generation → Krea to configure, get a key at " + "https://www.krea.ai/settings/api-tokens, or sign in to " + "a Nous account with the managed Krea gateway enabled " + "(`hermes setup`)." + ), + error_type="auth_required", + provider="krea", + aspect_ratio=aspect, + ) - model_id, meta = _resolve_model() + model_id, meta = _resolve_model(kwargs.get("model")) creativity = _resolve_creativity(kwargs.get("creativity")) + # The managed gateway only prices base text-to-image and URL + # ``image_style_references`` tiers. Trained styles (LoRAs) and + # moodboards have no managed price and are rejected at the gateway, so + # fail fast here with actionable guidance instead of a raw 400. + if managed is not None: + if isinstance(kwargs.get("styles"), list) and kwargs.get("styles"): + return error_response( + error=( + "Managed Krea (Nous Subscription) does not support " + "trained styles (LoRAs). Set KREA_API_KEY to use Krea " + "directly, or omit `styles`." + ), + error_type="unsupported_argument", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + if isinstance(kwargs.get("moodboards"), list) and kwargs.get("moodboards"): + return error_response( + error=( + "Managed Krea (Nous Subscription) does not support " + "moodboards. Set KREA_API_KEY to use Krea directly, or " + "omit `moodboards`." + ), + error_type="unsupported_argument", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) + payload: Dict[str, Any] = { "prompt": prompt, "aspect_ratio": krea_ar, @@ -309,13 +411,19 @@ class KreaImageGenProvider(ImageGenProvider): payload["moodboards"] = moodboards[:1] headers = { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {auth_token}", "Content-Type": "application/json", "User-Agent": "Hermes-Agent/1.0 (krea-image-gen)", } + if managed is not None: + # The gateway derives the per-generation billing idempotency + # boundary from this header (else it falls back to a body + # fingerprint). A fresh key per submit keeps each generation a + # distinct billable execution. + headers["x-idempotency-key"] = str(uuid.uuid4()) # 1. Submit job. - submit_url = f"{BASE_URL}/generate/image/krea/krea-2/{meta['path']}" + submit_url = f"{base_url}/generate/image/krea/krea-2/{meta['path']}" try: response = requests.post( submit_url, @@ -337,6 +445,32 @@ class KreaImageGenProvider(ImageGenProvider): except Exception: # noqa: BLE001 err_msg = resp.text[:300] if resp is not None else str(exc) logger.error("Krea submit failed (%d): %s", status, err_msg) + # On a managed 4xx, surface actionable remediation mirroring the + # FAL managed gateway path: the model may not be enabled/priced on + # the Nous Portal, or the gateway's shared Krea key hit its + # concurrency cap (429). + if managed is not None and 400 <= status < 500: + hint = ( + "Krea's shared-key concurrency cap was hit — retry shortly." + if status == 429 + else ( + f"Model '{model_id}' may not be enabled/priced on the " + "Nous Portal's Krea gateway. Set KREA_API_KEY to use " + "Krea directly, or pick a different model via " + "`hermes tools` → Image Generation." + ) + ) + return error_response( + error=( + f"Nous Subscription Krea gateway rejected '{model_id}' " + f"(HTTP {status}): {err_msg}. {hint}" + ), + error_type="api_error", + provider="krea", + model=model_id, + prompt=prompt, + aspect_ratio=aspect, + ) return error_response( error=f"Krea image generation failed ({status}): {err_msg}", error_type="api_error", @@ -387,10 +521,12 @@ class KreaImageGenProvider(ImageGenProvider): aspect_ratio=aspect, ) - # 2. Poll for completion. - job_url = f"{BASE_URL}/jobs/{job_id}" + # 2. Poll for completion. Status/result polling is bound to the same + # principal at the gateway, so the managed path polls the gateway's + # ``/jobs/{id}`` with the Nous token (404 on cross-user/unknown jobs). + job_url = f"{base_url}/jobs/{job_id}" poll_headers = { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {auth_token}", "User-Agent": "Hermes-Agent/1.0 (krea-image-gen)", } interval = _POLL_INITIAL_INTERVAL diff --git a/plugins/image_gen/krea/plugin.yaml b/plugins/image_gen/krea/plugin.yaml index bc650dc5222..9f42979ab75 100644 --- a/plugins/image_gen/krea/plugin.yaml +++ b/plugins/image_gen/krea/plugin.yaml @@ -1,6 +1,6 @@ name: krea -version: 1.0.0 -description: "Krea image generation backend (Krea 2 Large + Krea 2 Medium foundation models)." +version: 1.1.0 +description: "Krea image generation backend (Krea 2 Large + Medium + Medium Turbo foundation models). Direct KREA_API_KEY or managed Nous Subscription gateway." author: NousResearch kind: backend requires_env: diff --git a/tests/plugins/image_gen/test_krea_provider.py b/tests/plugins/image_gen/test_krea_provider.py index cc9dcd5a6b0..4f7b7919e51 100644 --- a/tests/plugins/image_gen/test_krea_provider.py +++ b/tests/plugins/image_gen/test_krea_provider.py @@ -76,10 +76,22 @@ class TestKreaImageGenProvider: def test_is_available_without_key(self, monkeypatch): monkeypatch.delenv("KREA_API_KEY", raising=False) + import plugins.image_gen.krea as krea_mod from plugins.image_gen.krea import KreaImageGenProvider + # No direct key AND no managed gateway → unavailable. + monkeypatch.setattr(krea_mod, "_managed_krea_gateway_ready", lambda: False) assert KreaImageGenProvider().is_available() is False + def test_is_available_via_managed_gateway_without_key(self, monkeypatch): + monkeypatch.delenv("KREA_API_KEY", raising=False) + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + # No direct key but the managed Nous gateway is ready → available. + monkeypatch.setattr(krea_mod, "_managed_krea_gateway_ready", lambda: True) + assert KreaImageGenProvider().is_available() is True + def test_list_models(self): from plugins.image_gen.krea import KreaImageGenProvider @@ -608,6 +620,188 @@ class TestPollRetryPolicy: assert mock_get.call_count == 2 +# --------------------------------------------------------------------------- +# Managed Nous gateway path +# --------------------------------------------------------------------------- + + +def _managed_cfg( + origin: str = "https://krea-gateway.example.com", + token: str = "nous-tok-abc", +): + from types import SimpleNamespace + + return SimpleNamespace( + vendor="krea", + gateway_origin=origin, + nous_user_token=token, + managed_mode=True, + ) + + +class TestManagedGateway: + def test_managed_submit_uses_gateway_origin_and_nous_token(self, monkeypatch): + """Managed mode submits to the gateway origin with the Nous token.""" + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + # Even with a direct key present, an active managed gateway wins. + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + submit = _submit_response() + poll = _poll_response(_completed_job()) + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll) as mock_get, \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="A managed lamp") + + assert result["success"] is True + post_url = mock_post.call_args[0][0] + assert post_url == ( + "https://krea-gateway.example.com/generate/image/krea/krea-2/medium" + ) + headers = mock_post.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer nous-tok-abc" + # Idempotency key drives the gateway's per-generation billing boundary. + assert headers["x-idempotency-key"] + # Poll is bound to the same gateway + Nous token. + poll_url = mock_get.call_args[0][0] + assert poll_url.startswith("https://krea-gateway.example.com/jobs/") + poll_headers = mock_get.call_args.kwargs["headers"] + assert poll_headers["Authorization"] == "Bearer nous-tok-abc" + + def test_managed_available_without_direct_key(self, monkeypatch): + """No KREA_API_KEY but an active gateway → generate proceeds (no auth_required).""" + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + monkeypatch.delenv("KREA_API_KEY", raising=False) + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + submit = _submit_response() + poll = _poll_response(_completed_job()) + with patch("plugins.image_gen.krea.requests.post", return_value=submit), \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is True + + def test_managed_4xx_returns_actionable_remediation(self, monkeypatch): + import requests as req_lib + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + resp = req_lib.Response() + resp.status_code = 402 + resp._content = b'{"error": {"message": "out of credits"}}' + resp.headers["Content-Type"] = "application/json" + resp.raise_for_status = MagicMock(side_effect=req_lib.HTTPError(response=resp)) + + with patch("plugins.image_gen.krea.requests.post", return_value=resp): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert result["error_type"] == "api_error" + assert "402" in result["error"] + assert "Nous Subscription Krea gateway" in result["error"] + assert "KREA_API_KEY" in result["error"] + + def test_managed_429_concurrency_hint(self, monkeypatch): + import requests as req_lib + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + resp = req_lib.Response() + resp.status_code = 429 + resp._content = b'{"error": {"message": "maximum number of concurrent jobs"}}' + resp.headers["Content-Type"] = "application/json" + resp.raise_for_status = MagicMock(side_effect=req_lib.HTTPError(response=resp)) + + with patch("plugins.image_gen.krea.requests.post", return_value=resp): + result = KreaImageGenProvider().generate(prompt="test") + + assert result["success"] is False + assert "429" in result["error"] + assert "concurrency" in result["error"].lower() + + def test_managed_blocks_styles(self, monkeypatch): + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + with patch("plugins.image_gen.krea.requests.post") as mock_post: + result = KreaImageGenProvider().generate( + prompt="test", + styles=[{"id": "lora-1"}], + ) + + assert result["success"] is False + assert result["error_type"] == "unsupported_argument" + assert "LoRA" in result["error"] or "styles" in result["error"] + # Never hit the network with an unsupported tier. + mock_post.assert_not_called() + + def test_managed_blocks_moodboards(self, monkeypatch): + import plugins.image_gen.krea as krea_mod + from plugins.image_gen.krea import KreaImageGenProvider + + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: _managed_cfg()) + + with patch("plugins.image_gen.krea.requests.post") as mock_post: + result = KreaImageGenProvider().generate( + prompt="test", + moodboards=[{"url": "https://x.com/m.png"}], + ) + + assert result["success"] is False + assert result["error_type"] == "unsupported_argument" + assert "moodboard" in result["error"].lower() + mock_post.assert_not_called() + + +class TestExplicitModelOverride: + def test_model_kwarg_overrides_config(self, monkeypatch): + """An explicit ``model`` kwarg (managed routing) wins over config/default.""" + from plugins.image_gen.krea import _resolve_model + + model_id, meta = _resolve_model("krea-2-large") + assert model_id == "krea-2-large" + assert meta["path"] == "large" + + def test_turbo_routes_to_medium_turbo_endpoint(self): + from plugins.image_gen.krea import KreaImageGenProvider + + submit = _submit_response() + poll = _poll_response(_completed_job()) + with patch("plugins.image_gen.krea.requests.post", return_value=submit) as mock_post, \ + patch("plugins.image_gen.krea.requests.get", return_value=poll), \ + patch( + "plugins.image_gen.krea.save_url_image", + return_value=Path("/tmp/x.png"), + ), \ + patch("plugins.image_gen.krea.time.sleep"): + result = KreaImageGenProvider().generate(prompt="test", model="krea-2-medium-turbo") + + assert result["success"] is True + assert result["model"] == "krea-2-medium-turbo" + post_url = mock_post.call_args[0][0] + assert post_url.endswith("/generate/image/krea/krea-2/medium-turbo") + + # --------------------------------------------------------------------------- # Registration # --------------------------------------------------------------------------- diff --git a/tests/tools/test_image_generation.py b/tests/tools/test_image_generation.py index df7d3a34abb..a548b6e0bdc 100644 --- a/tests/tools/test_image_generation.py +++ b/tests/tools/test_image_generation.py @@ -501,3 +501,134 @@ class TestManagedGatewayErrorTranslation: with pytest.raises(ConnectionError): image_tool._submit_fal_request("fal-ai/flux-2-pro", {"prompt": "x"}) + + +class TestKreaModelNormalization: + """Native ``krea-2-*`` detection for managed Krea routing.""" + + def test_native_models_detected(self, image_tool): + for mid in ("krea-2-medium", "krea-2-large", "krea-2-medium-turbo"): + assert image_tool.is_krea_model(mid) is True + assert image_tool._normalize_krea_model(mid) == mid + + def test_fal_krea_models_are_not_native_krea(self, image_tool): + # fal-ai/krea/v2/* stays on the FAL path — not the Krea plugin. + for mid in ( + "fal-ai/krea/v2/medium/text-to-image", + "fal-ai/krea/v2/large/text-to-image", + "fal-ai/krea/v2/medium", + "fal-ai/krea/v2/large/edit", + ): + assert image_tool.is_krea_model(mid) is False + assert image_tool._normalize_krea_model(mid) is None + + def test_non_krea_models_are_not_krea(self, image_tool): + for mid in ("fal-ai/flux-2/klein/9b", "fal-ai/nano-banana-pro", None, "", 123): + assert image_tool.is_krea_model(mid) is False + assert image_tool._normalize_krea_model(mid) is None + + +class TestManagedKreaRouting: + """`_maybe_route_managed_krea` only fires for Krea models in managed mode.""" + + def test_no_route_when_model_not_krea(self, image_tool, monkeypatch): + monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: None) + monkeypatch.setattr( + image_tool, "_read_configured_image_model", lambda: "fal-ai/flux-2/klein/9b" + ) + assert image_tool._maybe_route_managed_krea("p", "square") is None + + def test_no_route_when_provider_is_krea_plugin(self, image_tool, monkeypatch): + # provider == "krea" is handled by the normal plugin dispatch instead. + monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: "krea") + monkeypatch.setattr( + image_tool, "_read_configured_image_model", lambda: "krea-2-medium" + ) + assert image_tool._maybe_route_managed_krea("p", "square") is None + + def test_no_route_for_fal_krea_model_in_managed_mode(self, image_tool, monkeypatch): + # fal-ai/krea/v2/* stays on FAL even when the Krea gateway is available. + monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: None) + monkeypatch.setattr( + image_tool, + "_read_configured_image_model", + lambda: "fal-ai/krea/v2/medium/text-to-image", + ) + import plugins.image_gen.krea as krea_mod + from types import SimpleNamespace + + monkeypatch.setattr( + krea_mod, + "_resolve_managed_krea_gateway", + lambda: SimpleNamespace( + vendor="krea", + gateway_origin="https://krea-gateway.example.com", + nous_user_token="tok", + managed_mode=True, + ), + ) + assert image_tool._maybe_route_managed_krea("p", "square") is None + + def test_no_route_for_krea_model_in_direct_mode(self, image_tool, monkeypatch): + # Native krea-2-* selected, but no managed gateway (BYO/direct) → fall through. + monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: None) + monkeypatch.setattr( + image_tool, + "_read_configured_image_model", + lambda: "krea-2-medium", + ) + import plugins.image_gen.krea as krea_mod + + monkeypatch.setattr(krea_mod, "_resolve_managed_krea_gateway", lambda: None) + assert image_tool._maybe_route_managed_krea("p", "square") is None + + def test_routes_native_krea_model_to_krea_plugin_in_managed_mode( + self, image_tool, monkeypatch + ): + from types import SimpleNamespace + from unittest.mock import MagicMock + import json as _json + + monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: None) + monkeypatch.setattr( + image_tool, + "_read_configured_image_model", + lambda: "krea-2-large", + ) + import plugins.image_gen.krea as krea_mod + + monkeypatch.setattr( + krea_mod, + "_resolve_managed_krea_gateway", + lambda: SimpleNamespace( + vendor="krea", + gateway_origin="https://krea-gateway.example.com", + nous_user_token="tok", + managed_mode=True, + ), + ) + + fake_provider = MagicMock() + fake_provider.generate.return_value = {"success": True, "image": "/tmp/x.png"} + monkeypatch.setattr( + "agent.image_gen_registry.get_provider", lambda name: fake_provider + ) + monkeypatch.setattr( + "hermes_cli.plugins._ensure_plugins_discovered", lambda *a, **k: None + ) + + out = image_tool._maybe_route_managed_krea("a cat", "portrait") + assert out is not None + assert _json.loads(out)["success"] is True + kwargs = fake_provider.generate.call_args.kwargs + assert kwargs["model"] == "krea-2-large" + assert kwargs["prompt"] == "a cat" + assert kwargs["aspect_ratio"] == "portrait" + + +class TestFalKreaCatalog: + """Krea 2 on FAL remains in the FAL picker for FAL-billed users.""" + + def test_fal_krea_models_in_fal_catalog(self, image_tool): + assert "fal-ai/krea/v2/medium/text-to-image" in image_tool.FAL_MODELS + assert "fal-ai/krea/v2/large/text-to-image" in image_tool.FAL_MODELS diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index 81c6491f9d9..7806db57efb 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -374,20 +374,15 @@ FAL_MODELS: Dict[str, Dict[str, Any]] = { }, "max_reference_images": 3, }, - # Krea 2 — Krea's first foundation image model, day-0 partner launch on - # fal (2026-05-27). Same model family as our direct ``plugins/image_gen/krea`` - # backend, exposed here for users who prefer to bill through their - # existing FAL key / Nous Portal subscription rather than register - # directly with Krea. Both variants share the same parameter schema — - # only model id, price, and recommended use case differ. + # Krea 2 on FAL — same model family as ``plugins/image_gen/krea``, but billed + # through FAL / the FAL managed gateway. Native ``krea-2-*`` ids route to the + # dedicated Krea plugin instead. "fal-ai/krea/v2/medium/text-to-image": { "display": "Krea 2 Medium", "speed": "~15-25s", "strengths": "Illustration, anime, painting, expressive/artistic styles", "price": "$0.030 (text) / $0.035 (style refs)", "size_style": "aspect_ratio", - # Krea natively accepts 1:1, 4:3, 3:2, 16:9, 2.35:1, 4:5, 2:3, 9:16 — - # we map our 3 abstract ratios to the closest match. "sizes": { "landscape": "16:9", "square": "1:1", @@ -1406,6 +1401,115 @@ def _dispatch_to_plugin_provider( return json.dumps(result) +# --------------------------------------------------------------------------- +# Managed-mode Krea routing +# --------------------------------------------------------------------------- +# +# Native ``krea-2-*`` plugin model ids are served by the dedicated Krea managed +# gateway. ``fal-ai/krea/v2/*`` FAL catalog ids stay on the FAL path (BYO key +# or FAL managed gateway). Routing only fires in managed mode; direct/BYO users +# keep their unchanged pipeline. + +_KREA_NATIVE_MODELS = {"krea-2-medium", "krea-2-large", "krea-2-medium-turbo"} + + +def _normalize_krea_model(model_id: Optional[str]) -> Optional[str]: + """Return the native Krea plugin model id when ``model_id`` is ``krea-2-*``.""" + if not isinstance(model_id, str): + return None + candidate = model_id.strip() + if candidate in _KREA_NATIVE_MODELS: + return candidate + return None + + +def is_krea_model(model_id: Optional[str]) -> bool: + """True when ``model_id`` is a native Krea plugin id (``krea-2-*``).""" + return _normalize_krea_model(model_id) is not None + + +def _maybe_route_managed_krea( + prompt: str, + aspect_ratio: str, + image_url: Optional[str] = None, + reference_image_urls: Optional[list] = None, +) -> Optional[str]: + """Route a native ``krea-2-*`` model to the managed Krea gateway, in managed mode. + + Returns a JSON result string when handled by the Krea managed gateway, or + ``None`` to fall through to the normal plugin/FAL pipeline. Fires only when + all hold: + - the configured image model is a native ``krea-2-*`` id, AND + - the user isn't already routed to the Krea plugin via + ``image_gen.provider`` (that path dispatches normally), AND + - the managed Krea gateway is resolvable (portal/managed mode). + + Direct/BYO users (no managed gateway) fall through untouched. + """ + # ``provider == "krea"`` is already handled by the standard plugin dispatch. + if _read_configured_image_provider() == "krea": + return None + + normalized = _normalize_krea_model(_read_configured_image_model()) + if normalized is None: + return None + + # Only intercept on the managed path; BYO/direct users keep their pipeline. + try: + from plugins.image_gen.krea import _resolve_managed_krea_gateway + + if _resolve_managed_krea_gateway() is None: + return None + except Exception as exc: # noqa: BLE001 + logger.debug("Managed Krea routing probe failed: %s", exc) + return None + + try: + from agent.image_gen_registry import get_provider + from hermes_cli.plugins import _ensure_plugins_discovered + + _ensure_plugins_discovered() + provider = get_provider("krea") + except Exception as exc: # noqa: BLE001 + logger.debug("Managed Krea routing: provider unavailable: %s", exc) + return None + if provider is None: + return None + + kwargs: Dict[str, Any] = { + "prompt": prompt, + "aspect_ratio": aspect_ratio, + "model": normalized, + } + try: + if isinstance(image_url, str) and image_url.strip(): + kwargs["image_url"] = image_url.strip() + norm_refs = None + if reference_image_urls is not None: + from agent.image_gen_provider import normalize_reference_images + + norm_refs = normalize_reference_images(reference_image_urls) + if norm_refs: + kwargs["reference_image_urls"] = norm_refs + result = provider.generate(**kwargs) + except Exception as exc: # noqa: BLE001 + logger.warning("Managed Krea routing failed: %s", exc) + return json.dumps({ + "success": False, + "image": None, + "error": f"Managed Krea generation error: {exc}", + "error_type": "provider_exception", + }) + if not isinstance(result, dict): + return json.dumps({ + "success": False, + "image": None, + "error": "Krea provider returned a non-dict result", + "error_type": "provider_contract", + }) + return json.dumps(result) + + def _handle_image_generate(args, **kw): prompt = args.get("prompt", "") if not prompt: @@ -1416,7 +1520,8 @@ def _handle_image_generate(args, **kw): task_id = kw.get("task_id") # Route to a plugin-registered provider if one is active (and it's - # not the in-tree FAL path). + # not the in-tree FAL path). When ``image_gen.provider == "krea"`` this + # already reaches the Krea plugin's managed gateway path. dispatched = _dispatch_to_plugin_provider( prompt, aspect_ratio, image_url=image_url, @@ -1425,6 +1530,19 @@ def _handle_image_generate(args, **kw): if dispatched is not None: return _postprocess_image_generate_result(dispatched, task_id=task_id) + # Managed-mode Krea routing: when no explicit plugin provider is configured + # but the selected model is a native ``krea-2-*`` id, a portal user routes to + # the dedicated Krea managed gateway. ``fal-ai/krea/v2/*`` models stay on the + # FAL path below. Runs after plugin dispatch (which returns None when no + # provider is set) so the BYO/direct FAL path stays untouched. + krea_routed = _maybe_route_managed_krea( + prompt, aspect_ratio, + image_url=image_url, + reference_image_urls=reference_image_urls, + ) + if krea_routed is not None: + return _postprocess_image_generate_result(krea_routed, task_id=task_id) + raw = image_generate_tool( prompt=prompt, aspect_ratio=aspect_ratio,