feat(media): opt-in upscale pass for image_generate and video_generate across FAL and Krea
The generated-media surface previously had almost no upscaler coverage: only fal-ai/flux-2-pro chained Clarity Upscaler (hardcoded catalog default), every other image model returned ~1MP output with no high-res path, and video had no upscaler at all. Krea's API treats the enhancer as a standard second pass; this brings the same shape to Hermes. - image_generate: new optional 'upscale' boolean in the tool schema. Explicit true chains the backend upscaler on ANY model (including edits); explicit false disables flux-2-pro's automatic default; omitted keeps per-model catalog behavior. Response now reports 'upscaled' so the agent knows which resolution it got. - FAL image path: explicit flag overrides the catalog 'upscale' default (Clarity Upscaler, 2x). Failure falls back to the native image. - Krea plugin: upscale=true chains Krea Enhance (/generate/enhance/krea/enhance, 2x, prompt-guided) through the same BYO/managed base URL + auth as generation, with a best-effort poll loop that never fails a successful generation. - video_generate: new optional 'upscale' boolean; FAL video plugin chains ByteDance SeedVR2 (fal-ai/seedvr/upscale/video, 2x factor mode). Providers without upscalers ignore the kwarg per the ABC contract (documented in both ABCs). Validation: targeted suites green (123 tests across 6 files, including new coverage for override-wins/default-kept/failure-fallback on all three paths); live E2E on direct FAL verified both chains end-to-end (klein 9b + Clarity upscaled image; pixverse-v6 1s 360p + SeedVR2 upscaled video).
This commit is contained in:
Teknium
@@ -185,6 +185,12 @@ class ImageGenProvider(abc.ABC):
|
||||
or :func:`error_response`. ``kwargs`` may contain forward-compat
|
||||
parameters future versions of the schema will expose —
|
||||
implementations MUST ignore unknown keys (no TypeError).
|
||||
|
||||
Known optional kwarg: ``upscale`` (bool) — when true, the caller
|
||||
requests a post-generation high-resolution pass through the
|
||||
backend's upscaler/enhancer. Providers without an upscaler simply
|
||||
ignore it; providers that honor it should report ``upscaled: True``
|
||||
in the response ``extra``.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -193,6 +193,12 @@ class VideoGenProvider(abc.ABC):
|
||||
or :func:`error_response`. ``kwargs`` may contain forward-compat
|
||||
parameters future versions of the schema will expose —
|
||||
implementations MUST ignore unknown keys (no TypeError).
|
||||
|
||||
Known optional kwarg: ``upscale`` (bool) — when true, the caller
|
||||
requests a post-generation high-resolution pass through the
|
||||
backend's video upscaler. Providers without an upscaler simply
|
||||
ignore it; providers that honor it should report ``upscaled: True``
|
||||
in the response ``extra``.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -143,6 +143,7 @@ class FalImageGenProvider(ImageGenProvider):
|
||||
"num_images",
|
||||
"output_format",
|
||||
"seed",
|
||||
"upscale",
|
||||
)
|
||||
if key in kwargs and kwargs[key] is not None
|
||||
}
|
||||
|
||||
@@ -111,6 +111,12 @@ _RETRYABLE_POLL_STATUSES = frozenset({408, 409, 425, 429, 500, 502, 503, 504})
|
||||
|
||||
_TERMINAL_STATES = {"completed", "failed", "cancelled"}
|
||||
|
||||
# Krea Enhance — the upscale/enhancer endpoint used for the optional
|
||||
# ``upscale`` pass after generation ("1.5K native, 4K via Enhancer" is
|
||||
# Krea's own pipeline shape). Cheap creative enhancer, max 8K.
|
||||
_ENHANCE_PATH = "/generate/enhance/krea/enhance"
|
||||
_ENHANCE_SCALE_FACTOR = 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
@@ -214,6 +220,121 @@ def _resolve_creativity(value: Optional[str]) -> str:
|
||||
return "medium"
|
||||
|
||||
|
||||
def _poll_krea_job(
|
||||
base_url: str,
|
||||
auth_token: str,
|
||||
job_id: str,
|
||||
*,
|
||||
timeout_seconds: float = _POLL_TIMEOUT_SECONDS,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Poll ``/jobs/{job_id}`` until terminal; return the job dict or None.
|
||||
|
||||
Best-effort variant of the main generate() poll loop used for secondary
|
||||
jobs (the Enhance upscale pass): any failure returns ``None`` so the
|
||||
caller can fall back instead of failing the whole generation.
|
||||
"""
|
||||
job_url = f"{base_url}/jobs/{job_id}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {auth_token}",
|
||||
"User-Agent": "Hermes-Agent/1.0 (krea-image-gen)",
|
||||
}
|
||||
interval = _POLL_INITIAL_INTERVAL
|
||||
deadline = time.monotonic() + timeout_seconds
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
interval = min(interval * _POLL_BACKOFF, _POLL_MAX_INTERVAL)
|
||||
try:
|
||||
resp = requests.get(job_url, headers=headers, timeout=30)
|
||||
resp.raise_for_status()
|
||||
job = resp.json()
|
||||
except requests.HTTPError as exc:
|
||||
status = exc.response.status_code if exc.response is not None else 0
|
||||
if status not in _RETRYABLE_POLL_STATUSES or time.monotonic() >= deadline:
|
||||
logger.warning("Krea enhance poll failed (%d) for job %s", status, job_id)
|
||||
return None
|
||||
continue
|
||||
except Exception as exc: # noqa: BLE001 — timeout/connection/JSON
|
||||
if time.monotonic() >= deadline:
|
||||
logger.warning("Krea enhance poll gave up for job %s: %s", job_id, exc)
|
||||
return None
|
||||
continue
|
||||
|
||||
if isinstance(job, dict):
|
||||
status_str = job.get("status")
|
||||
if status_str in _TERMINAL_STATES or job.get("completed_at"):
|
||||
return job
|
||||
if time.monotonic() >= deadline:
|
||||
logger.warning("Krea enhance job %s did not finish in %ds", job_id, int(timeout_seconds))
|
||||
return None
|
||||
|
||||
|
||||
def _extract_result_url(job: Optional[Dict[str, Any]]) -> Optional[str]:
|
||||
"""Pull the first result URL out of a terminal Krea job dict."""
|
||||
if not isinstance(job, dict):
|
||||
return None
|
||||
result = job.get("result")
|
||||
if not isinstance(result, dict):
|
||||
return None
|
||||
urls = result.get("urls")
|
||||
if isinstance(urls, list):
|
||||
for candidate in urls:
|
||||
if isinstance(candidate, str) and candidate.strip():
|
||||
return candidate.strip()
|
||||
single = result.get("url")
|
||||
if isinstance(single, str) and single.strip():
|
||||
return single.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _enhance_image(
|
||||
base_url: str,
|
||||
auth_token: str,
|
||||
image_url: str,
|
||||
prompt: str,
|
||||
*,
|
||||
managed: bool,
|
||||
) -> Optional[str]:
|
||||
"""Run Krea Enhance on ``image_url``; return the enhanced URL or None.
|
||||
|
||||
Best-effort: any submit/poll/result failure logs and returns ``None`` so
|
||||
the caller falls back to the original (un-upscaled) image — an upscale
|
||||
failure must never destroy an already-successful generation.
|
||||
"""
|
||||
headers = {
|
||||
"Authorization": f"Bearer {auth_token}",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "Hermes-Agent/1.0 (krea-image-gen)",
|
||||
}
|
||||
if managed:
|
||||
headers["x-idempotency-key"] = str(uuid.uuid4())
|
||||
payload: Dict[str, Any] = {
|
||||
"image_url": image_url,
|
||||
"image_scaling_factor": _ENHANCE_SCALE_FACTOR,
|
||||
# Keep the enhancer faithful to the generated composition: the
|
||||
# original prompt guides detail, and default ai_strength stays
|
||||
# conservative (Krea default 0.4 adds detail without redrawing).
|
||||
"prompt": prompt,
|
||||
}
|
||||
try:
|
||||
resp = requests.post(
|
||||
f"{base_url}{_ENHANCE_PATH}", headers=headers, json=payload, timeout=30,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
job_id = (resp.json() or {}).get("job_id")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Krea Enhance submit failed: %s", exc)
|
||||
return None
|
||||
if not isinstance(job_id, str) or not job_id:
|
||||
logger.warning("Krea Enhance submit response missing job_id")
|
||||
return None
|
||||
|
||||
job = _poll_krea_job(base_url, auth_token, job_id)
|
||||
if not isinstance(job, dict) or job.get("status") in {"failed", "cancelled"}:
|
||||
logger.warning("Krea Enhance job %s did not complete successfully", job_id)
|
||||
return None
|
||||
return _extract_result_url(job)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -700,6 +821,33 @@ class KreaImageGenProvider(ImageGenProvider):
|
||||
aspect_ratio=aspect,
|
||||
)
|
||||
|
||||
# Optional high-resolution pass (Krea Enhance). Explicit agent/user
|
||||
# opt-in via the ``upscale`` kwarg; config default via
|
||||
# ``image_gen.krea.upscale``. Best-effort: failure falls back to the
|
||||
# original image rather than failing the generation.
|
||||
upscaled = False
|
||||
upscale_requested = kwargs.get("upscale")
|
||||
if not isinstance(upscale_requested, bool):
|
||||
cfg_krea = _load_krea_config().get("krea")
|
||||
upscale_requested = bool(
|
||||
isinstance(cfg_krea, dict) and cfg_krea.get("upscale") is True
|
||||
)
|
||||
if upscale_requested:
|
||||
enhanced_url = _enhance_image(
|
||||
base_url,
|
||||
auth_token,
|
||||
result_image_url,
|
||||
prompt,
|
||||
managed=managed is not None,
|
||||
)
|
||||
if enhanced_url:
|
||||
result_image_url = enhanced_url
|
||||
upscaled = True
|
||||
else:
|
||||
logger.warning(
|
||||
"Krea Enhance pass failed — returning native-resolution image"
|
||||
)
|
||||
|
||||
# Materialise locally — Krea result URLs may expire, mirroring
|
||||
# what we do for xAI / OpenAI URL responses (#26942).
|
||||
try:
|
||||
@@ -719,7 +867,10 @@ class KreaImageGenProvider(ImageGenProvider):
|
||||
"resolution": DEFAULT_RESOLUTION,
|
||||
"creativity": creativity,
|
||||
"job_id": job_id,
|
||||
"upscaled": upscaled,
|
||||
}
|
||||
if upscaled:
|
||||
extra["upscale_factor"] = _ENHANCE_SCALE_FACTOR
|
||||
if isinstance(job.get("completed_at"), str):
|
||||
extra["completed_at"] = job["completed_at"]
|
||||
|
||||
|
||||
@@ -531,6 +531,45 @@ def _check_fal_video_available() -> bool:
|
||||
return _resolve_managed_fal_video_gateway() is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Upscaler (SeedVR2 — video upscale pass)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# ByteDance SeedVR2 on FAL: $0.001/megapixel of output video. A 5s 720p→1440p
|
||||
# 2x pass is roughly $0.44. Faithful restoration-style upscaler (the same
|
||||
# model family Krea exposes as its "SeedVR2" video enhancer).
|
||||
UPSCALER_ENDPOINT = "fal-ai/seedvr/upscale/video"
|
||||
UPSCALER_FACTOR = 2
|
||||
|
||||
|
||||
def _upscale_video(video_url: str) -> Optional[str]:
|
||||
"""Upscale a generated video via SeedVR2; return the new URL or None.
|
||||
|
||||
Best-effort: any failure logs and returns ``None`` so the caller falls
|
||||
back to the native-resolution video — an upscale failure must never
|
||||
destroy an already-successful generation.
|
||||
"""
|
||||
try:
|
||||
logger.info("Upscaling video with SeedVR2 (%dx)...", UPSCALER_FACTOR)
|
||||
handle = _submit_fal_video_request(UPSCALER_ENDPOINT, {
|
||||
"video_url": video_url,
|
||||
"upscale_mode": "factor",
|
||||
"upscale_factor": UPSCALER_FACTOR,
|
||||
})
|
||||
result = handle.get()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Video upscale failed: %s", exc)
|
||||
return None
|
||||
|
||||
video = (result or {}).get("video") if isinstance(result, dict) else None
|
||||
if isinstance(video, dict) and video.get("url"):
|
||||
return video["url"]
|
||||
if isinstance(video, str) and video:
|
||||
return video
|
||||
logger.warning("Video upscaler returned no URL")
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -618,6 +657,7 @@ class FALVideoGenProvider(VideoGenProvider):
|
||||
negative_prompt: Optional[str] = None,
|
||||
audio: Optional[bool] = None,
|
||||
seed: Optional[int] = None,
|
||||
upscale: Optional[bool] = None,
|
||||
**kwargs: Any,
|
||||
) -> Dict[str, Any]:
|
||||
if not _check_fal_video_available():
|
||||
@@ -722,9 +762,25 @@ class FALVideoGenProvider(VideoGenProvider):
|
||||
provider="fal", model=family_id, prompt=prompt,
|
||||
)
|
||||
|
||||
extra: Dict[str, Any] = {"endpoint": endpoint}
|
||||
# Optional high-resolution pass (SeedVR2). Explicit agent/user opt-in
|
||||
# only; best-effort — failure falls back to the native-resolution
|
||||
# video rather than failing the generation.
|
||||
upscaled = False
|
||||
if upscale:
|
||||
upscaled_url = _upscale_video(url)
|
||||
if upscaled_url:
|
||||
url = upscaled_url
|
||||
upscaled = True
|
||||
else:
|
||||
logger.warning(
|
||||
"Video upscale pass failed — returning native-resolution video"
|
||||
)
|
||||
|
||||
extra: Dict[str, Any] = {"endpoint": endpoint, "upscaled": upscaled}
|
||||
if upscaled:
|
||||
extra["upscale_factor"] = UPSCALER_FACTOR
|
||||
if isinstance(video, dict):
|
||||
if video.get("file_size"):
|
||||
if video.get("file_size") and not upscaled:
|
||||
extra["file_size"] = video["file_size"]
|
||||
if video.get("content_type"):
|
||||
extra["content_type"] = video["content_type"]
|
||||
|
||||
@@ -581,6 +581,85 @@ class TestExplicitModelOverride:
|
||||
assert post_url.endswith("/generate/image/krea/krea-2/medium-turbo")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Upscale pass (Krea Enhance)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestUpscalePass:
|
||||
def _run_generate(self, *, upscale, enhance_job):
|
||||
"""Drive generate() with sequenced post/get mocks.
|
||||
|
||||
Sequence: generation submit POST → generation poll GET; then (when
|
||||
upscale fires) enhance submit POST → enhance poll GET.
|
||||
"""
|
||||
from plugins.image_gen.krea import KreaImageGenProvider
|
||||
|
||||
gen_submit = _submit_response()
|
||||
gen_poll = _poll_response(_completed_job("https://krea.cdn/native.png"))
|
||||
enh_submit = _submit_response("00000000-0000-0000-0000-00000000e0e0")
|
||||
enh_poll = _poll_response(enhance_job) if enhance_job else None
|
||||
|
||||
posts = [gen_submit, enh_submit]
|
||||
gets = [gen_poll] + ([enh_poll] if enh_poll else [])
|
||||
|
||||
with patch("plugins.image_gen.krea.requests.post", side_effect=posts) as mock_post, \
|
||||
patch("plugins.image_gen.krea.requests.get", side_effect=gets) as mock_get, \
|
||||
patch(
|
||||
"plugins.image_gen.krea.save_url_image",
|
||||
side_effect=lambda url, prefix: Path(f"/tmp/{url.rsplit('/', 1)[-1]}"),
|
||||
), \
|
||||
patch("plugins.image_gen.krea.time.sleep"):
|
||||
result = KreaImageGenProvider().generate(prompt="a lamp", upscale=upscale)
|
||||
return result, mock_post, mock_get
|
||||
|
||||
def test_upscale_routes_through_enhance_endpoint(self):
|
||||
enhance_job = {
|
||||
"job_id": "00000000-0000-0000-0000-00000000e0e0",
|
||||
"status": "completed",
|
||||
"created_at": "2026-05-27T00:00:00Z",
|
||||
"completed_at": "2026-05-27T00:01:00Z",
|
||||
"result": {"urls": ["https://krea.cdn/enhanced.png"]},
|
||||
}
|
||||
result, mock_post, _ = self._run_generate(upscale=True, enhance_job=enhance_job)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["upscaled"] is True
|
||||
assert result["upscale_factor"] == 2
|
||||
assert result["image"].endswith("enhanced.png")
|
||||
# Second POST hit the Enhance endpoint with the native image + factor.
|
||||
assert mock_post.call_count == 2
|
||||
enh_url = mock_post.call_args_list[1][0][0]
|
||||
assert enh_url.endswith("/generate/enhance/krea/enhance")
|
||||
enh_payload = mock_post.call_args_list[1].kwargs["json"]
|
||||
assert enh_payload["image_url"] == "https://krea.cdn/native.png"
|
||||
assert enh_payload["image_scaling_factor"] == 2
|
||||
assert enh_payload["prompt"] == "a lamp"
|
||||
|
||||
def test_upscale_failure_falls_back_to_native(self):
|
||||
failed_job = {
|
||||
"job_id": "00000000-0000-0000-0000-00000000e0e0",
|
||||
"status": "failed",
|
||||
"created_at": "2026-05-27T00:00:00Z",
|
||||
"completed_at": "2026-05-27T00:01:00Z",
|
||||
"result": None,
|
||||
}
|
||||
result, mock_post, _ = self._run_generate(upscale=True, enhance_job=failed_job)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["upscaled"] is False
|
||||
assert result["image"].endswith("native.png")
|
||||
assert mock_post.call_count == 2 # enhance attempted, fell back
|
||||
|
||||
def test_no_upscale_by_default(self):
|
||||
result, mock_post, _ = self._run_generate(upscale=None, enhance_job=None)
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["upscaled"] is False
|
||||
assert result["image"].endswith("native.png")
|
||||
assert mock_post.call_count == 1 # only the generation submit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -341,3 +341,78 @@ class TestPayloadBuilder:
|
||||
)
|
||||
# Only prompt — no payload bloat for fields we can't verify
|
||||
assert p == {"prompt": "a horse galloping"}
|
||||
|
||||
|
||||
class TestUpscalePass:
|
||||
"""Opt-in SeedVR2 upscale chain after generation."""
|
||||
|
||||
@pytest.fixture
|
||||
def with_fake_fal(self, monkeypatch):
|
||||
"""Stub fal_client.submit, capturing every endpoint hit in order."""
|
||||
import sys
|
||||
import types
|
||||
|
||||
captured = {"calls": []}
|
||||
|
||||
class FakeHandle:
|
||||
def __init__(self, endpoint):
|
||||
self._endpoint = endpoint
|
||||
|
||||
def get(self):
|
||||
if self._endpoint.endswith("upscale/video"):
|
||||
return {"video": {"url": "https://fake/upscaled.mp4"}}
|
||||
return {"video": {"url": "https://fake/native.mp4"}}
|
||||
|
||||
fake = types.ModuleType("fal_client")
|
||||
def _submit(endpoint, arguments=None, headers=None):
|
||||
captured["calls"].append((endpoint, arguments))
|
||||
return FakeHandle(endpoint)
|
||||
fake.submit = _submit # type: ignore
|
||||
monkeypatch.setitem(sys.modules, "fal_client", fake)
|
||||
|
||||
from plugins.video_gen import fal as fal_plugin
|
||||
fal_plugin._fal_client = None
|
||||
fal_plugin._managed_fal_video_client = None
|
||||
fal_plugin._managed_fal_video_client_config = None
|
||||
|
||||
monkeypatch.setenv("FAL_KEY", "test")
|
||||
monkeypatch.setattr(fal_plugin, "_resolve_managed_fal_video_gateway", lambda: None)
|
||||
return captured
|
||||
|
||||
def test_upscale_chains_seedvr(self, with_fake_fal):
|
||||
from plugins.video_gen.fal import FALVideoGenProvider, UPSCALER_ENDPOINT
|
||||
|
||||
result = FALVideoGenProvider().generate(
|
||||
"a dog", model="pixverse-v6", upscale=True,
|
||||
)
|
||||
assert result["success"] is True
|
||||
assert result["video"] == "https://fake/upscaled.mp4"
|
||||
assert result["upscaled"] is True
|
||||
assert result["upscale_factor"] == 2
|
||||
endpoints = [c[0] for c in with_fake_fal["calls"]]
|
||||
assert endpoints == ["fal-ai/pixverse/v6/text-to-video", UPSCALER_ENDPOINT]
|
||||
# Upscale request carries the native URL + factor mode.
|
||||
upscale_args = with_fake_fal["calls"][1][1]
|
||||
assert upscale_args["video_url"] == "https://fake/native.mp4"
|
||||
assert upscale_args["upscale_mode"] == "factor"
|
||||
|
||||
def test_no_upscale_by_default(self, with_fake_fal):
|
||||
from plugins.video_gen.fal import FALVideoGenProvider
|
||||
|
||||
result = FALVideoGenProvider().generate("a dog", model="pixverse-v6")
|
||||
assert result["success"] is True
|
||||
assert result["video"] == "https://fake/native.mp4"
|
||||
assert result["upscaled"] is False
|
||||
assert len(with_fake_fal["calls"]) == 1
|
||||
|
||||
def test_upscale_failure_falls_back_to_native(self, with_fake_fal, monkeypatch):
|
||||
from plugins.video_gen import fal as fal_plugin
|
||||
from plugins.video_gen.fal import FALVideoGenProvider
|
||||
|
||||
monkeypatch.setattr(fal_plugin, "_upscale_video", lambda url: None)
|
||||
result = FALVideoGenProvider().generate(
|
||||
"a dog", model="pixverse-v6", upscale=True,
|
||||
)
|
||||
assert result["success"] is True
|
||||
assert result["video"] == "https://fake/native.mp4"
|
||||
assert result["upscaled"] is False
|
||||
|
||||
@@ -274,12 +274,13 @@ class TestRegistryIntegration:
|
||||
|
||||
def test_schema_exposes_expected_agent_params(self, image_tool):
|
||||
"""The agent-facing schema exposes the unified text+image surface:
|
||||
prompt (required), aspect_ratio, and the image-to-image inputs
|
||||
image_url + reference_image_urls. Model selection stays a user-level
|
||||
config choice, never an agent-level arg."""
|
||||
prompt (required), aspect_ratio, the image-to-image inputs
|
||||
image_url + reference_image_urls, and the opt-in upscale pass. Model
|
||||
selection stays a user-level config choice, never an agent-level arg."""
|
||||
props = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]
|
||||
assert set(props.keys()) == {
|
||||
"prompt", "aspect_ratio", "image_url", "reference_image_urls",
|
||||
"upscale",
|
||||
}
|
||||
assert image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["required"] == ["prompt"]
|
||||
|
||||
@@ -447,3 +448,124 @@ class TestFalKreaCatalog:
|
||||
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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Opt-in upscale pass
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _FakeHandle:
|
||||
def __init__(self, result):
|
||||
self._result = result
|
||||
|
||||
def get(self):
|
||||
return self._result
|
||||
|
||||
|
||||
class TestUpscaleOptIn:
|
||||
"""Explicit ``upscale`` overrides the per-model catalog default."""
|
||||
|
||||
def _run(self, image_tool, monkeypatch, *, model, upscale, upscaler_called):
|
||||
monkeypatch.setenv("FAL_IMAGE_MODEL", model)
|
||||
monkeypatch.setattr(image_tool, "fal_key_is_configured", lambda: True)
|
||||
monkeypatch.setattr(image_tool, "_resolve_managed_fal_gateway", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
image_tool, "_submit_fal_request",
|
||||
lambda endpoint, arguments=None: _FakeHandle(
|
||||
{"images": [{"url": "https://fal/native.png", "width": 1024, "height": 768}]}
|
||||
),
|
||||
)
|
||||
calls = []
|
||||
|
||||
def _fake_upscale(url, prompt):
|
||||
calls.append(url)
|
||||
return {
|
||||
"url": "https://fal/upscaled.png", "width": 2048, "height": 1536,
|
||||
"upscaled": True, "upscale_factor": 2,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(image_tool, "_upscale_image", _fake_upscale)
|
||||
|
||||
import json as _json
|
||||
out = _json.loads(image_tool.image_generate_tool("a cat", upscale=upscale))
|
||||
assert out["success"] is True
|
||||
assert bool(calls) is upscaler_called
|
||||
assert out["upscaled"] is upscaler_called
|
||||
expected_url = "https://fal/upscaled.png" if upscaler_called else "https://fal/native.png"
|
||||
assert out["image"] == expected_url
|
||||
|
||||
def test_explicit_true_upscales_non_default_model(self, image_tool, monkeypatch):
|
||||
"""Klein has upscale=False in the catalog — explicit True wins."""
|
||||
self._run(image_tool, monkeypatch,
|
||||
model="fal-ai/flux-2/klein/9b", upscale=True, upscaler_called=True)
|
||||
|
||||
def test_explicit_false_disables_flux2_pro_default(self, image_tool, monkeypatch):
|
||||
"""flux-2-pro defaults to upscale=True — explicit False wins."""
|
||||
self._run(image_tool, monkeypatch,
|
||||
model="fal-ai/flux-2-pro", upscale=False, upscaler_called=False)
|
||||
|
||||
def test_omitted_keeps_catalog_default_off(self, image_tool, monkeypatch):
|
||||
self._run(image_tool, monkeypatch,
|
||||
model="fal-ai/flux-2/klein/9b", upscale=None, upscaler_called=False)
|
||||
|
||||
def test_omitted_keeps_catalog_default_on(self, image_tool, monkeypatch):
|
||||
self._run(image_tool, monkeypatch,
|
||||
model="fal-ai/flux-2-pro", upscale=None, upscaler_called=True)
|
||||
|
||||
def test_upscale_failure_falls_back_to_native(self, image_tool, monkeypatch):
|
||||
monkeypatch.setenv("FAL_IMAGE_MODEL", "fal-ai/flux-2/klein/9b")
|
||||
monkeypatch.setattr(image_tool, "fal_key_is_configured", lambda: True)
|
||||
monkeypatch.setattr(image_tool, "_resolve_managed_fal_gateway", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
image_tool, "_submit_fal_request",
|
||||
lambda endpoint, arguments=None: _FakeHandle(
|
||||
{"images": [{"url": "https://fal/native.png"}]}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(image_tool, "_upscale_image", lambda url, prompt: None)
|
||||
|
||||
import json as _json
|
||||
out = _json.loads(image_tool.image_generate_tool("a cat", upscale=True))
|
||||
assert out["success"] is True
|
||||
assert out["image"] == "https://fal/native.png"
|
||||
assert out["upscaled"] is False
|
||||
|
||||
|
||||
class TestUpscaleDispatchForwarding:
|
||||
"""The tool handler forwards explicit upscale to plugin providers."""
|
||||
|
||||
def test_dispatch_forwards_upscale(self, image_tool, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
import json as _json
|
||||
|
||||
monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: "krea")
|
||||
monkeypatch.setattr(image_tool, "_read_configured_image_model", lambda: None)
|
||||
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._dispatch_to_plugin_provider("a cat", "square", upscale=True)
|
||||
assert _json.loads(out)["success"] is True
|
||||
assert fake_provider.generate.call_args.kwargs["upscale"] is True
|
||||
|
||||
def test_dispatch_omits_upscale_when_unset(self, image_tool, monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
monkeypatch.setattr(image_tool, "_read_configured_image_provider", lambda: "krea")
|
||||
monkeypatch.setattr(image_tool, "_read_configured_image_model", lambda: None)
|
||||
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
|
||||
)
|
||||
|
||||
image_tool._dispatch_to_plugin_provider("a cat", "square")
|
||||
assert "upscale" not in fake_provider.generate.call_args.kwargs
|
||||
|
||||
@@ -94,3 +94,20 @@ class TestUnifiedDispatch:
|
||||
props = VIDEO_GENERATE_SCHEMA["parameters"]["properties"]
|
||||
assert "operation" not in props
|
||||
assert "video_url" not in props
|
||||
|
||||
|
||||
def test_upscale_in_schema_and_forwarded(self):
|
||||
"""`upscale` is an agent-facing param, forwarded to providers when
|
||||
set and omitted (not None) when unset."""
|
||||
from tools.video_generation_tool import VIDEO_GENERATE_SCHEMA
|
||||
props = VIDEO_GENERATE_SCHEMA["parameters"]["properties"]
|
||||
assert props["upscale"]["type"] == "boolean"
|
||||
|
||||
provider = _RecordingProvider()
|
||||
video_gen_registry.register_provider(provider)
|
||||
result = self._run({"prompt": "a dog", "upscale": True}, configured="fake")
|
||||
assert result["success"] is True
|
||||
assert provider.last_kwargs["upscale"] is True
|
||||
|
||||
self._run({"prompt": "a dog"}, configured="fake")
|
||||
assert "upscale" not in provider.last_kwargs
|
||||
|
||||
@@ -1079,6 +1079,7 @@ def image_generate_tool(
|
||||
seed: Optional[int] = None,
|
||||
image_url: Optional[str] = None,
|
||||
reference_image_urls: Optional[list] = None,
|
||||
upscale: Optional[bool] = None,
|
||||
) -> str:
|
||||
"""Generate an image from a text prompt, or edit a source image, via FAL.
|
||||
|
||||
@@ -1203,9 +1204,15 @@ def image_generate_tool(
|
||||
if not images:
|
||||
raise ValueError("No images were generated")
|
||||
|
||||
# Edit endpoints already return the final composition; the Clarity
|
||||
# upscaler is a text-to-image quality pass, so skip it for edits.
|
||||
should_upscale = bool(meta.get("upscale", False)) and not use_edit
|
||||
# Explicit ``upscale`` (agent/user opt-in via the tool schema) wins
|
||||
# over the per-model catalog default — including for edits, where an
|
||||
# explicit request is intentional. The catalog default keeps skipping
|
||||
# edits (Clarity is a text-to-image quality pass and must not alter
|
||||
# edit compositions silently).
|
||||
if upscale is not None:
|
||||
should_upscale = bool(upscale)
|
||||
else:
|
||||
should_upscale = bool(meta.get("upscale", False)) and not use_edit
|
||||
|
||||
formatted_images = []
|
||||
for img in images:
|
||||
@@ -1241,6 +1248,7 @@ def image_generate_tool(
|
||||
"success": True,
|
||||
"image": formatted_images[0]["url"] if formatted_images else None,
|
||||
"modality": modality,
|
||||
"upscaled": bool(formatted_images and formatted_images[0].get("upscaled")),
|
||||
}
|
||||
|
||||
debug_call_data["success"] = True
|
||||
@@ -1448,6 +1456,16 @@ IMAGE_GENERATE_SCHEMA = {
|
||||
"capped per-model; the description above indicates the max."
|
||||
),
|
||||
},
|
||||
"upscale": {
|
||||
"type": "boolean",
|
||||
"description": (
|
||||
"Optional high-resolution pass: when true, the generated "
|
||||
"image is run through the active backend's upscaler/"
|
||||
"enhancer (extra cost and latency, roughly 2x resolution). "
|
||||
"Use when the user asks for high-res / print / wallpaper "
|
||||
"quality output. Omit for the model's native resolution."
|
||||
),
|
||||
},
|
||||
},
|
||||
"required": ["prompt"],
|
||||
},
|
||||
@@ -1498,6 +1516,7 @@ def _dispatch_to_plugin_provider(
|
||||
aspect_ratio: str,
|
||||
image_url: Optional[str] = None,
|
||||
reference_image_urls: Optional[list] = None,
|
||||
upscale: Optional[bool] = None,
|
||||
):
|
||||
"""Route the call to a plugin-registered provider when one is selected.
|
||||
|
||||
@@ -1512,7 +1531,9 @@ def _dispatch_to_plugin_provider(
|
||||
|
||||
``image_url`` / ``reference_image_urls`` enable image-to-image / editing:
|
||||
they are forwarded to the provider's ``generate()`` so the backend can
|
||||
route to its edit endpoint.
|
||||
route to its edit endpoint. ``upscale`` (when explicitly set) requests a
|
||||
post-generation high-resolution pass; providers without upscale support
|
||||
ignore it via their ``**kwargs`` (the ABC contract).
|
||||
"""
|
||||
configured = _read_configured_image_provider()
|
||||
if not configured or configured == "fal":
|
||||
@@ -1568,6 +1589,8 @@ def _dispatch_to_plugin_provider(
|
||||
norm_refs = normalize_reference_images(reference_image_urls)
|
||||
if norm_refs:
|
||||
kwargs["reference_image_urls"] = norm_refs
|
||||
if upscale is not None:
|
||||
kwargs["upscale"] = bool(upscale)
|
||||
result = provider.generate(**kwargs)
|
||||
except TypeError as exc:
|
||||
# A provider whose generate() signature predates image_url support
|
||||
@@ -1655,6 +1678,7 @@ def _maybe_route_managed_krea(
|
||||
aspect_ratio: str,
|
||||
image_url: Optional[str] = None,
|
||||
reference_image_urls: Optional[list] = None,
|
||||
upscale: Optional[bool] = None,
|
||||
) -> Optional[str]:
|
||||
"""Route a native ``krea-2-*`` model to the managed Krea gateway, in managed mode.
|
||||
|
||||
@@ -1713,6 +1737,8 @@ def _maybe_route_managed_krea(
|
||||
norm_refs = normalize_reference_images(reference_image_urls)
|
||||
if norm_refs:
|
||||
kwargs["reference_image_urls"] = norm_refs
|
||||
if upscale is not None:
|
||||
kwargs["upscale"] = bool(upscale)
|
||||
result = provider.generate(**kwargs)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Managed Krea routing failed: %s", exc)
|
||||
@@ -1781,6 +1807,9 @@ def _handle_image_generate(args, **kw):
|
||||
aspect_ratio = args.get("aspect_ratio", DEFAULT_ASPECT_RATIO)
|
||||
image_url = args.get("image_url")
|
||||
reference_image_urls = args.get("reference_image_urls")
|
||||
upscale = args.get("upscale")
|
||||
if not isinstance(upscale, bool):
|
||||
upscale = None
|
||||
task_id = kw.get("task_id")
|
||||
|
||||
# Terminal-backend confinement chokepoint: convert path-like sources to
|
||||
@@ -1799,6 +1828,7 @@ def _handle_image_generate(args, **kw):
|
||||
prompt, aspect_ratio,
|
||||
image_url=image_url,
|
||||
reference_image_urls=reference_image_urls,
|
||||
upscale=upscale,
|
||||
)
|
||||
if dispatched is not None:
|
||||
return _postprocess_image_generate_result(dispatched, task_id=task_id)
|
||||
@@ -1812,6 +1842,7 @@ def _handle_image_generate(args, **kw):
|
||||
prompt, aspect_ratio,
|
||||
image_url=image_url,
|
||||
reference_image_urls=reference_image_urls,
|
||||
upscale=upscale,
|
||||
)
|
||||
if krea_routed is not None:
|
||||
return _postprocess_image_generate_result(krea_routed, task_id=task_id)
|
||||
@@ -1821,6 +1852,7 @@ def _handle_image_generate(args, **kw):
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_url=image_url,
|
||||
reference_image_urls=reference_image_urls,
|
||||
upscale=upscale,
|
||||
)
|
||||
return _postprocess_image_generate_result(raw, task_id=task_id)
|
||||
|
||||
|
||||
@@ -145,6 +145,17 @@ VIDEO_GENERATE_SCHEMA: Dict[str, Any] = {
|
||||
"dependent)."
|
||||
),
|
||||
},
|
||||
"upscale": {
|
||||
"type": "boolean",
|
||||
"description": (
|
||||
"Optional high-resolution pass: when true, the generated "
|
||||
"video is run through the active backend's video upscaler "
|
||||
"(extra cost and latency, roughly 2x resolution). Use when "
|
||||
"the user asks for high-res / 4K output. Omit for the "
|
||||
"model's native resolution. Ignored by backends without "
|
||||
"an upscaler."
|
||||
),
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
@@ -328,6 +339,7 @@ def _handle_video_generate(args: Dict[str, Any], **_kw: Any) -> str:
|
||||
negative_prompt = (args.get("negative_prompt") or "").strip() or None
|
||||
audio = _coerce_bool(args.get("audio"))
|
||||
seed = _coerce_int(args.get("seed"))
|
||||
upscale = _coerce_bool(args.get("upscale"))
|
||||
model_override = (args.get("model") or "").strip() or None
|
||||
|
||||
# Soft validation — providers do their own. Prompt is required by the
|
||||
@@ -361,6 +373,7 @@ def _handle_video_generate(args: Dict[str, Any], **_kw: Any) -> str:
|
||||
"negative_prompt": negative_prompt,
|
||||
"audio": audio,
|
||||
"seed": seed,
|
||||
"upscale": upscale,
|
||||
}
|
||||
# Drop None entries so providers see clean defaults.
|
||||
kwargs = {k: v for k, v in kwargs.items() if v is not None}
|
||||
|
||||
@@ -158,9 +158,32 @@ GPT Image 2 maps to 4:3 presets rather than 16:9 because its minimum pixel count
|
||||
|
||||
This translation happens in `_build_fal_payload()` — agent code never has to know about per-model schema differences.
|
||||
|
||||
## Automatic Upscaling
|
||||
## Upscaling
|
||||
|
||||
Upscaling via FAL's **Clarity Upscaler** is gated per-model:
|
||||
### On-demand (any model)
|
||||
|
||||
The agent-facing `upscale` parameter requests a high-resolution pass after
|
||||
generation on **any** model — ask for "high-res", "print quality", or
|
||||
"wallpaper" output and the agent sets `upscale: true`:
|
||||
|
||||
| Backend | Upscaler | Result |
|
||||
|---|---|---|
|
||||
| **FAL.ai** (all models) | Clarity Upscaler | ~2× resolution, +$0.03/MP |
|
||||
| **Krea** (Krea 2 family) | Krea Enhance | 2× resolution (up to 8K ceiling) |
|
||||
| Other backends | — | parameter is ignored (native resolution returned) |
|
||||
|
||||
An explicit `upscale: false` also *disables* the automatic pass on models
|
||||
that default to it (currently `flux-2-pro`). Passing `upscale: true` with an
|
||||
image edit runs the pass on the edited output too.
|
||||
|
||||
`video_generate` accepts the same `upscale` parameter on the FAL backend,
|
||||
chaining ByteDance's **SeedVR2** video upscaler (2×, $0.001/MP of output
|
||||
video) after generation.
|
||||
|
||||
### Automatic (per-model default)
|
||||
|
||||
Upscaling via FAL's **Clarity Upscaler** also runs automatically for models
|
||||
whose catalog entry sets `upscale: True`:
|
||||
|
||||
| Model | Upscale? | Why |
|
||||
|---|---|---|
|
||||
@@ -177,14 +200,14 @@ When upscaling runs, it uses these settings:
|
||||
| Guidance scale | 4 |
|
||||
| Inference steps | 18 |
|
||||
|
||||
If upscaling fails (network issue, rate limit), the original image is returned automatically.
|
||||
If upscaling fails (network issue, rate limit), the original image is returned automatically. The response reports `upscaled: true/false` so the agent knows which resolution it got.
|
||||
|
||||
## How It Works Internally
|
||||
|
||||
1. **Model resolution** — `_resolve_fal_model()` reads `image_gen.model` from `config.yaml`, falls back to the `FAL_IMAGE_MODEL` env var, then to `fal-ai/flux-2/klein/9b`.
|
||||
2. **Payload building** — `_build_fal_payload()` translates your `aspect_ratio` into the model's native format (preset enum, aspect-ratio enum, or GPT literal), merges the model's default params, applies any caller overrides, then filters to the model's `supports` whitelist so unsupported keys are never sent.
|
||||
3. **Submission** — `_submit_fal_request()` routes via direct FAL credentials or the managed Nous gateway.
|
||||
4. **Upscaling** — runs only if the model's metadata has `upscale: True`.
|
||||
4. **Upscaling** — runs when the agent passed `upscale: true`, or when the model's metadata has `upscale: True` (explicit `upscale: false` wins over the metadata default).
|
||||
5. **Delivery** — final image URL returned to the agent, which emits a `MEDIA:<url>` tag that platform adapters convert to native media.
|
||||
|
||||
## Debugging
|
||||
|
||||
Reference in New Issue
Block a user