From 6cb3dd9ef36e5e799da0c4537ae793b1b65b58d9 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 1 Sep 2026 11:53:10 +0200 Subject: [PATCH] fix(image): validate output sizes --- gemini_nodes.py | 11 ++++++++++- tests/test_gemini_nodes.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/gemini_nodes.py b/gemini_nodes.py index 7b10a04..f26e33e 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -206,6 +206,11 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): "gemini-3.1-flash-image", "gemini-3.1-flash-lite-image", "gemini-3-pro-image-preview", "gemini-3-pro-image", ] + IMAGE_SIZE_BY_SELECTABLE_MODEL = { + "gemini-3.1-flash-image": ("512", "1K", "2K", "4K"), + "gemini-3.1-flash-lite-image": ("1K",), + "gemini-3-pro-image": ("1K", "2K", "4K"), + } MEDIA_RES_MODELS = [ "gemini-3.1-flash-lite", "gemini-3-flash-preview", "gemini-3.1-pro-preview", "gemini-3.5-flash", "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest" @@ -248,7 +253,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): IO.String.Input("retry_pattern", default="", optional=True, multiline=False, tooltip="Regex pattern to match in response text. If matched, retry with new seed. Leave empty to disable."), IO.Int.Input("max_retries", default=3, min=0, max=10, step=1, tooltip="Maximum number of retry attempts when pattern matches. 0 disables retry."), IO.String.Input("timeout_fallback_text", default="", optional=True, multiline=True, tooltip="Text returned when the Gemini request times out. Leave empty to return the standard timeout message."), - IO.Combo.Input("image_size", options=["None", "512", "1K", "2K", "4K"], default="None", tooltip="Generated image resolution. Gemini 3.1 Flash Lite Image supports only 1K; 512 is supported only by Gemini 3.1 Flash Image."), + IO.Combo.Input("image_size", options=["None", "512", "1K", "2K", "4K"], default="None", optional=True, tooltip="Generated image resolution. Gemini 3.1 Flash Lite Image supports only 1K; 512 is supported only by Gemini 3.1 Flash Image."), cls.GeminiVideoConfig.Input("video", optional=True, tooltip="Optional configured Gemini video input with embedded audio and sampling FPS."), IO.Autogrow.Input( "image_inputs", @@ -718,6 +723,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): if aspect_ratio not in (None, "None"): image_config["aspect_ratio"] = aspect_ratio if image_size not in (None, "None"): + supported_image_sizes = cls.IMAGE_SIZE_BY_SELECTABLE_MODEL.get(model) + if supported_image_sizes is not None and image_size not in supported_image_sizes: + allowed = ", ".join(supported_image_sizes) + raise ValueError(f"{model} image_size must be one of: {allowed}.") image_config["image_size"] = image_size if allow_all_people: image_config["person_generation"] = "ALLOW_ALL" diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index 6b607de..68e506b 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -975,6 +975,34 @@ def test_image_config_preserves_requested_image_size(): assert config.image_config.image_size == "2K" +@pytest.mark.parametrize( + ("model", "image_size", "allowed"), + [ + ("gemini-3.1-flash-lite-image", "2K", "1K"), + ("gemini-3-pro-image", "512", "1K, 2K, 4K"), + ], +) +def test_image_config_rejects_unsupported_image_size(model, image_size, allowed): + with pytest.raises(ValueError, match=rf"^{model} image_size must be one of: {allowed}\.$"): + gemini_nodes.SSL_GeminiTextPrompt._build_generate_content_config( + model=model, + temperature=1.0, + top_p=0.95, + top_k=40, + max_output_tokens=8192, + seed=7, + include_images=True, + response_modalities=["IMAGE", "TEXT"], + aspect_ratio="None", + padded_system_instruction="system", + thinking_level="None", + thinking_budget=0, + include_thoughts=False, + media_resolution="unspecified", + image_size=image_size, + ) + + def test_enterprise_image_config_allows_all_people(): config = gemini_nodes.SSL_GeminiTextPrompt._build_generate_content_config( model="gemini-3.1-flash-image",