diff --git a/interactions_nodes.py b/interactions_nodes.py index d23d5f1..ba7e787 100644 --- a/interactions_nodes.py +++ b/interactions_nodes.py @@ -28,6 +28,7 @@ from .gemini_nodes import ( GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiVideoConfig, + SSL_GeminiTextPrompt, ) @@ -657,6 +658,49 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput: print(f"[INFO] SSL_GeminiInteractionsTextPrompt execute called, model: {model}") + + # If model is an image generation model, reroute to SSL_GeminiTextPrompt (GenerateContent) + if model in cls.IMAGE_MODELS: + print(f"[INFO] Model '{model}' is an image generation model; rerouting to GenerateContent engine (SSL_GeminiTextPrompt).") + gen_output = SSL_GeminiTextPrompt.execute( + config=config, + prompt=prompt, + system_instruction=system_instruction, + model=model, + temperature=temperature, + top_p=top_p, + top_k=top_k, + max_output_tokens=max_output_tokens, + include_images=include_images, + aspect_ratio=aspect_ratio, + bypass_mode=bypass_mode, + thinking_budget=thinking_budget, + use_proxy=use_proxy, + proxy_host=proxy_host, + proxy_port=proxy_port, + use_seed=use_seed, + seed=seed, + timeout=timeout, + include_thoughts=include_thoughts, + thinking_level=thinking_level, + media_resolution=media_resolution, + retry_pattern=retry_pattern, + max_retries=max_retries, + timeout_fallback_text=timeout_fallback_text, + image_size=image_size, + video=video, + image_inputs=image_inputs, + response_schema=response_schema, + ) + return IO.NodeOutput( + gen_output[0], # text + gen_output[1], # image + gen_output[2], # final_actual_seed + gen_output[3], # structured_output + gen_output[4], # thoughts + "", # interaction_id + ) + use_cache = bool(config.get("use_cache", False)) cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60)) cache_seed = int(config.get("cache_seed", 0)) diff --git a/tests/test_interactions_nodes.py b/tests/test_interactions_nodes.py index 93a5fc9..d54ef3a 100644 --- a/tests/test_interactions_nodes.py +++ b/tests/test_interactions_nodes.py @@ -289,19 +289,15 @@ def test_structured_output_with_schema(mock_interactions_client): ] -def test_image_generation(mock_interactions_client): - # Create small valid 16x16 PNG image - img = Image.new("RGB", (16, 16), color=(255, 0, 0)) - buf = BytesIO() - img.save(buf, format="PNG") - b64_png = base64.b64encode(buf.getvalue()).decode("utf-8") +def test_image_generation(mock_interactions_client, monkeypatch): + called_with = [] + fake_img = torch.zeros((1, 16, 16, 3), dtype=torch.float32) - mock_interactions_client.responses = [{ - "text": "Generated an image", - "image_data": b64_png, - "image_mime": "image/png", - "id": "int_img_gen_1", - }] + def fake_generate(**kwargs): + called_with.append(kwargs) + return interactions_nodes.IO.NodeOutput("Generated an image", fake_img, 0, "", "") + + monkeypatch.setattr(interactions_nodes.SSL_GeminiTextPrompt, "execute", fake_generate) args = _execute_kwargs( model="gemini-3.1-flash-image", @@ -311,14 +307,16 @@ def test_image_generation(mock_interactions_client): ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + assert output[0] == "Generated an image" image_tensor = output[1] assert torch.is_tensor(image_tensor) assert image_tensor.shape == (1, 16, 16, 3) - - call = mock_interactions_client.calls[0] - assert call["response_format"] == [ - {"type": "image", "aspect_ratio": "16:9", "image_size": "1K"} - ] + assert output[5] == "" # interaction_id is empty for rerouted image models + assert len(called_with) == 1 + assert called_with[0]["model"] == "gemini-3.1-flash-image" + assert called_with[0]["include_images"] is True + assert called_with[0]["aspect_ratio"] == "16:9" + assert called_with[0]["image_size"] == "1K" def test_thinking_and_thoughts_extraction(mock_interactions_client):