diff --git a/agent/image_gen_provider.py b/agent/image_gen_provider.py index a3eeb1e4c8..b76bc8c01a 100644 --- a/agent/image_gen_provider.py +++ b/agent/image_gen_provider.py @@ -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``. """ diff --git a/agent/video_gen_provider.py b/agent/video_gen_provider.py index 8630f8f204..c444979f7e 100644 --- a/agent/video_gen_provider.py +++ b/agent/video_gen_provider.py @@ -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``. """ diff --git a/plugins/image_gen/fal/__init__.py b/plugins/image_gen/fal/__init__.py index f59ceebb49..face534286 100644 --- a/plugins/image_gen/fal/__init__.py +++ b/plugins/image_gen/fal/__init__.py @@ -143,6 +143,7 @@ class FalImageGenProvider(ImageGenProvider): "num_images", "output_format", "seed", + "upscale", ) if key in kwargs and kwargs[key] is not None } diff --git a/plugins/image_gen/krea/__init__.py b/plugins/image_gen/krea/__init__.py index 64c53c447f..86bae4a986 100644 --- a/plugins/image_gen/krea/__init__.py +++ b/plugins/image_gen/krea/__init__.py @@ -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"] diff --git a/plugins/video_gen/fal/__init__.py b/plugins/video_gen/fal/__init__.py index f2a76dca1e..e4aaa7431d 100644 --- a/plugins/video_gen/fal/__init__.py +++ b/plugins/video_gen/fal/__init__.py @@ -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"] diff --git a/tests/plugins/image_gen/test_krea_provider.py b/tests/plugins/image_gen/test_krea_provider.py index a2b2cdf35d..cdb9ebd6cc 100644 --- a/tests/plugins/image_gen/test_krea_provider.py +++ b/tests/plugins/image_gen/test_krea_provider.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/plugins/video_gen/test_fal_plugin.py b/tests/plugins/video_gen/test_fal_plugin.py index f26364a923..45771cc1a1 100644 --- a/tests/plugins/video_gen/test_fal_plugin.py +++ b/tests/plugins/video_gen/test_fal_plugin.py @@ -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 diff --git a/tests/tools/test_image_generation.py b/tests/tools/test_image_generation.py index a919116c20..6cea8ff9a3 100644 --- a/tests/tools/test_image_generation.py +++ b/tests/tools/test_image_generation.py @@ -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 diff --git a/tests/tools/test_video_generation_dispatch.py b/tests/tools/test_video_generation_dispatch.py index 4ee4f71dc8..c7101bd4ba 100644 --- a/tests/tools/test_video_generation_dispatch.py +++ b/tests/tools/test_video_generation_dispatch.py @@ -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 diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index 664a5c5f7e..238a290e20 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -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) diff --git a/tools/video_generation_tool.py b/tools/video_generation_tool.py index d35138441c..36e1b2dd85 100644 --- a/tools/video_generation_tool.py +++ b/tools/video_generation_tool.py @@ -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} diff --git a/website/docs/user-guide/features/image-generation.md b/website/docs/user-guide/features/image-generation.md index 9c27e7eebe..e33544c716 100644 --- a/website/docs/user-guide/features/image-generation.md +++ b/website/docs/user-guide/features/image-generation.md @@ -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:` tag that platform adapters convert to native media. ## Debugging