From 892a215f130d5a7c9e303dd8bce2e2b102d0c740 Mon Sep 17 00:00:00 2001
From: silveroxides <ishimarukaito@gmail.com>
Date: Tue, 24 Jun 2025 16:58:53 +0200
Subject: [PATCH] Make image generation work with 2.0-flash-exp

---
 gemini_nodes.py | 83 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 55 insertions(+), 28 deletions(-)

diff --git a/gemini_nodes.py b/gemini_nodes.py
index 771d9cd..63511c5 100644
--- a/gemini_nodes.py
+++ b/gemini_nodes.py
@@ -329,34 +329,61 @@ class SSL_GeminiTextPrompt(ComfyNodeABC):
 
             print(padded_prompt)
             print(padded_system_instruction)
-            generate_content_config = types.GenerateContentConfig(
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                max_output_tokens=max_output_tokens,
-                safety_settings=[
-                    types.SafetySetting(
-                        category="HARM_CATEGORY_HARASSMENT",
-                        threshold="BLOCK_NONE",
-                    ),
-                    types.SafetySetting(
-                        category="HARM_CATEGORY_HATE_SPEECH",
-                        threshold="BLOCK_NONE",
-                    ),
-                    types.SafetySetting(
-                        category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                        threshold="BLOCK_NONE",
-                    ),
-                    types.SafetySetting(
-                        category="HARM_CATEGORY_DANGEROUS_CONTENT",
-                        threshold="BLOCK_NONE",
-                    ),
-                ],
-                response_modalities=response_modalities,
-                system_instruction=[
-                    types.Part.from_text(text=padded_system_instruction),
-                ],
-            )
+            if include_images == True:
+                generate_content_config = types.GenerateContentConfig(
+                    temperature=temperature,
+                    top_p=top_p,
+                    top_k=top_k,
+                    max_output_tokens=max_output_tokens,
+                    safety_settings=[
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_HARASSMENT",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_HATE_SPEECH",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_DANGEROUS_CONTENT",
+                            threshold="BLOCK_NONE",
+                        ),
+                    ],
+                    response_modalities=response_modalities,
+                )
+            else:
+                generate_content_config = types.GenerateContentConfig(
+                    temperature=temperature,
+                    top_p=top_p,
+                    top_k=top_k,
+                    max_output_tokens=max_output_tokens,
+                    safety_settings=[
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_HARASSMENT",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_HATE_SPEECH",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                            threshold="BLOCK_NONE",
+                        ),
+                        types.SafetySetting(
+                            category="HARM_CATEGORY_DANGEROUS_CONTENT",
+                            threshold="BLOCK_NONE",
+                        ),
+                    ],
+                    response_modalities=response_modalities,
+                    system_instruction=[
+                        types.Part.from_text(text=padded_system_instruction),
+                    ],
+                )
 
             if use_seed == True and actual_seed is not None:
                 try: