diff --git a/main.py b/main.py index b575da75c..484b6c5cc 100644 --- a/main.py +++ b/main.py @@ -9237,6 +9237,80 @@ def gemini_model_name(model): value = selected_model(model, "gemini-3-pro-image-preview").strip() return value[len("models/"):] if value.startswith("models/") else value +GEMINI_IMAGE_ASPECT_RATIOS = ( + (1, 1, "1:1"), (1, 4, "1:4"), (1, 8, "1:8"), + (2, 3, "2:3"), (3, 2, "3:2"), (3, 4, "3:4"), + (4, 1, "4:1"), (4, 3, "4:3"), (4, 5, "4:5"), + (5, 4, "5:4"), (8, 1, "8:1"), (9, 16, "9:16"), + (16, 9, "16:9"), (21, 9, "21:9"), +) + +def gemini_supported_aspect_ratio(size, fallback="1:1"): + width, height = parse_size_pair(size) + if not width or not height: + match = re.fullmatch(r"\s*(\d+)\s*:\s*(\d+)\s*", str(size or "")) + if match: + width, height = int(match.group(1)), int(match.group(2)) + if not width or not height: + return fallback + ratio = width / height + # Compare proportionally rather than by raw decimal distance so portrait + # and landscape ratios are treated symmetrically. + return min( + GEMINI_IMAGE_ASPECT_RATIOS, + key=lambda item: abs(math.log(ratio / (item[0] / item[1]))), + )[2] + +def banana_image_request_params(size): + return { + "aspect_ratio": gemini_supported_aspect_ratio(size), + } + +def image_provider_model_ids(provider): + return { + str(item or "").strip().lower() + for item in ((provider or {}).get("image_models") or []) + if str(item or "").strip() + } + +def banana_model_family(model): + value = str(model or "").strip().lower() + if value in {"banana", "nano-banana-pro"}: + return "nano-banana-pro" + match = re.fullmatch(r"(nano-banana-(?:pro|2))(?:-(?:1k|2k|4k))?", value) + return match.group(1) if match else "" + +def looks_like_gemini_image_model(model): + value = str(model or "").strip().lower() + return value.startswith("gemini-") and "image" in value + +def route_openai_image_request(provider, model, size): + requested_model = selected_model(model, IMAGE_MODEL) + family = banana_model_family(requested_model) + params = {} + routed_model = requested_model + if family: + params = banana_image_request_params(size) + available = image_provider_model_ids(provider) + _, resolution = apimart_size_resolution(size) + target_model = family if resolution == "1k" else f"{family}-{resolution}" + if not available or target_model in available: + routed_model = target_model + elif family in available: + routed_model = family + elif requested_model.lower() not in available: + routed_model = target_model + elif looks_like_gemini_image_model(requested_model): + # Some OpenAI-compatible gateways expose Gemini image IDs directly. + # They still require a supported Gemini aspect_ratio. + params = banana_image_request_params(size) + return { + "requested_model": requested_model, + "model": routed_model, + "params": params, + "family": family or ("gemini-image" if params else ""), + } + def gemini_endpoint_url(provider, model): model_name = urllib.parse.quote(gemini_model_name(model), safe="") return provider_endpoint_url(provider, "image_generation_endpoint", f"/v1beta/models/{model_name}:generateContent") @@ -9248,9 +9322,10 @@ def gemini_image_config(size): if raw in {"1K", "2K", "4K"}: return {"aspectRatio": "1:1", "imageSize": raw} if re.fullmatch(r"\d+\s*:\s*\d+", raw): - return {"aspectRatio": raw.replace(" ", ""), "imageSize": "1K"} + return {"aspectRatio": gemini_supported_aspect_ratio(raw), "imageSize": "1K"} return {"aspectRatio": "1:1", "imageSize": "2K"} - aspect_ratio, resolution = apimart_size_resolution(size) + _, resolution = apimart_size_resolution(size) + aspect_ratio = gemini_supported_aspect_ratio(size) return {"aspectRatio": aspect_ratio, "imageSize": resolution.upper()} def gemini_reference_part(ref): @@ -10495,6 +10570,9 @@ async def generate_ai_image(prompt, size, quality, model, reference_images=None, return await generate_gemini_provider_image(prompt, size, model, reference_images, provider) if is_volcengine_provider(provider): return await generate_volcengine_provider_image(prompt, size, model, reference_images, provider) + image_route = route_openai_image_request(provider, model, size) + model = image_route["model"] + routed_image_params = image_route["params"] is_gpt2 = is_gpt_image_2_model(model) is_apimart = is_apimart_provider(provider) # 不对 GPT 尺寸做任何缩小/拦截：用户选什么尺寸就原样发给上游； @@ -10516,6 +10594,12 @@ async def generate_ai_image(prompt, size, quality, model, reference_images=None, response = None async def post_openai_edits(edit_files=None): data = {"model": model, "prompt": prompt, "size": size} + if routed_image_params: + # OpenAI-compatible Banana gateways otherwise infer a ratio + # from arbitrary pixel dimensions (for example 3840x1648 -> + # 240:103), which Gemini rejects. Send the supported preset + # explicitly while retaining size for gateway compatibility. + data.update(routed_image_params) if quality: data["quality"] = quality return await client.post( @@ -10674,6 +10758,8 @@ def post_video_proxy_multipart(): "response_format": "url", "n": 1, "image": image_payload, } + if routed_image_params: + body.update(routed_image_params) if quality: body["quality"] = quality response = await client.post(gen_url, headers=api_headers(provider=provider, model=model), json=body) @@ -10684,6 +10770,8 @@ def post_video_proxy_multipart(): ) else: body = {"model": model, "prompt": prompt, "size": size, "response_format": "url", "n": 1} + if routed_image_params: + body.update(routed_image_params) if quality: body["quality"] = quality response = await client.post( diff --git a/tests/test_banana_outpaint_request.py b/tests/test_banana_outpaint_request.py new file mode 100644 index 000000000..42af5d703 --- /dev/null +++ b/tests/test_banana_outpaint_request.py @@ -0,0 +1,184 @@ +import asyncio +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) +SPEC = importlib.util.spec_from_file_location("infinite_canvas_main", ROOT / "main.py") +APP = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(APP) + + +class FakeResponse: + status_code = 200 + text = "" + reason_phrase = "OK" + + def raise_for_status(self): + return None + + def json(self): + return {"data": [{"url": "https://example.invalid/out.png"}]} + + +class FakeAsyncClient: + last_post = None + + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def post(self, url, **kwargs): + FakeAsyncClient.last_post = {"url": url, **kwargs} + return FakeResponse() + + +class BananaOutpaintRequestTests(unittest.TestCase): + def test_openai_compatible_banana_edit_sends_supported_aspect_ratio(self): + provider = { + "id": "custom-api", + "name": "comfly", + "base_url": "https://ai.comfly.org", + "protocol": "openai", + "image_request_mode": "openai", + "api_key": "test-key", + "model_protocols": {}, + } + reference = [{"url": "/assets/input/source.png", "name": "source.png", "role": "source"}] + with ( + patch.object(APP, "get_api_provider", return_value=provider), + patch.object(APP, "output_file_from_url", return_value=str(__file__)), + patch.object(APP, "api_headers", return_value={}), + patch.object(APP.httpx, "AsyncClient", FakeAsyncClient), + ): + asyncio.run(APP.generate_ai_image( + "Remove white area and fill the scene", + "3840x1648", + "high", + "nano-banana-pro-4k", + reference, + "custom-api", + )) + + data = FakeAsyncClient.last_post["data"] + self.assertEqual(data["model"], "nano-banana-pro-4k") + self.assertEqual(data["size"], "3840x1648") + self.assertEqual(data["aspect_ratio"], "21:9") + self.assertNotIn("image_size", data) + + def test_every_outpaint_preset_maps_to_a_supported_gemini_ratio(self): + allowed = {item[2] for item in APP.GEMINI_IMAGE_ASPECT_RATIOS} + for preset in ("1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9", "9:21"): + with self.subTest(preset=preset): + self.assertIn(APP.gemini_supported_aspect_ratio(preset), allowed) + + def test_route_uses_base_banana_family_and_size_parameters(self): + provider = { + "image_models": [ + "nano-banana-2", "nano-banana-2-2k", "nano-banana-2-4k", + "nano-banana-pro", "nano-banana-pro-2k", "nano-banana-pro-4k", + ] + } + cases = ( + ("nano-banana-pro-4k", "1024x1024", "nano-banana-pro", "1:1"), + ("nano-banana-pro-4k", "2048x2048", "nano-banana-pro-2k", "1:1"), + ("nano-banana-pro-2k", "4096x2304", "nano-banana-pro-4k", "16:9"), + ("nano-banana-2-2k", "3840x1648", "nano-banana-2-4k", "21:9"), + ) + for requested, size, expected_model, expected_ratio in cases: + with self.subTest(requested=requested, size=size): + route = APP.route_openai_image_request(provider, requested, size) + self.assertEqual(route["model"], expected_model) + self.assertEqual(route["params"]["aspect_ratio"], expected_ratio) + self.assertNotIn("image_size", route["params"]) + + def test_route_falls_back_to_base_when_resolution_alias_is_unavailable(self): + provider = {"image_models": ["nano-banana-pro"]} + route = APP.route_openai_image_request(provider, "nano-banana-pro-2k", "4096x2304") + self.assertEqual(route["model"], "nano-banana-pro") + self.assertNotIn("image_size", route["params"]) + + def test_openai_gateway_gemini_image_id_receives_image_parameters(self): + route = APP.route_openai_image_request({}, "gemini-3-pro-image-preview", "1536x1024") + self.assertEqual(route["model"], "gemini-3-pro-image-preview") + self.assertEqual(route["params"], {"aspect_ratio": "3:2"}) + + def test_unrelated_image_model_is_not_routed(self): + route = APP.route_openai_image_request({}, "flux-kontext-pro", "1536x1024") + self.assertEqual(route["model"], "flux-kontext-pro") + self.assertEqual(route["params"], {}) + + def test_text_to_image_request_uses_routed_model_and_parameters(self): + provider = { + "id": "custom-api", + "name": "comfly", + "base_url": "https://ai.comfly.org", + "protocol": "openai", + "image_request_mode": "openai", + "api_key": "test-key", + "model_protocols": {}, + "image_models": ["nano-banana-2", "nano-banana-2-2k", "nano-banana-2-4k"], + } + with ( + patch.object(APP, "get_api_provider", return_value=provider), + patch.object(APP, "api_headers", return_value={}), + patch.object(APP.httpx, "AsyncClient", FakeAsyncClient), + ): + asyncio.run(APP.generate_ai_image( + "A panoramic orchard", + "4096x2304", + "high", + "nano-banana-2-2k", + [], + "custom-api", + )) + + body = FakeAsyncClient.last_post["json"] + self.assertEqual(body["model"], "nano-banana-2-4k") + self.assertEqual(body["aspect_ratio"], "16:9") + self.assertEqual(body["size"], "4096x2304") + self.assertNotIn("image_size", body) + + def test_image_2_request_does_not_receive_banana_only_fields(self): + provider = { + "id": "custom-api", + "name": "comfly", + "base_url": "https://ai.comfly.org", + "protocol": "openai", + "image_request_mode": "openai", + "api_key": "test-key", + "model_protocols": {}, + } + reference = [{"url": "/assets/input/source.png", "name": "source.png", "role": "source"}] + with ( + patch.object(APP, "get_api_provider", return_value=provider), + patch.object(APP, "output_file_from_url", return_value=str(__file__)), + patch.object(APP, "api_headers", return_value={}), + patch.object(APP.httpx, "AsyncClient", FakeAsyncClient), + ): + asyncio.run(APP.generate_ai_image( + "Expand the image", + "1536x1024", + "high", + "gpt-image-2", + reference, + "custom-api", + )) + + data = FakeAsyncClient.last_post["data"] + self.assertNotIn("aspect_ratio", data) + self.assertNotIn("image_size", data) + + +if __name__ == "__main__": + unittest.main()