diff --git a/gemini_nodes.py b/gemini_nodes.py index bb2e8a0..bbde025 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -76,6 +76,7 @@ class SSL_GeminiTextPrompt(ComfyNodeABC): "max_output_tokens": (IO.INT, {"default": 8192, "min": 1, "max": 8192, "step": 1}), "include_images": (IO.BOOLEAN, {"default": True}), "bypass_mode": (["None", "system_instruction", "prompt", "both"], {"default": "None"}), + "thinking_budget": (IO.INT, {"default": 0, "min": 0, "max": 24576, "step": 1, "tooltip": "0 disables thinking mode, 1 will activate it as default thinking budget and anything above sets specific budget"}), }, "optional": { "input_image": (IO.IMAGE,), @@ -152,12 +153,13 @@ class SSL_GeminiTextPrompt(ComfyNodeABC): return tensor def generate(self, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, bypass_mode, - input_image=None, input_image_2=None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, + thinking_budget, input_image=None, input_image_2=None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=True, seed=0): original_http_proxy = os.environ.get('HTTP_PROXY') original_https_proxy = os.environ.get('HTTPS_PROXY') original_http_proxy_lower = os.environ.get('http_proxy') original_https_proxy_lower = os.environ.get('https_proxy') + thinking_models = ["gemini-2.5-pro", "gemini-2.5-pro-preview-05-06", "gemini-2.5-flash", "gemini-2.5-flash-preview-04-17", "gemini-2.5-flash-lite-preview-06-17", "gemini-beta-3.0-pro"] print(f"[INFO] Starting generation, model: {model}, temperature: {temperature}") @@ -359,6 +361,42 @@ class SSL_GeminiTextPrompt(ComfyNodeABC): ], response_modalities=response_modalities, ) + elif model in thinking_models: + generate_content_config = types.GenerateContentConfig( + thinking_config = types.ThinkingConfig( + thinking_budget=thinking_budget, + ), + temperature=temperature, + top_p=top_p, + top_k=top_k, + max_output_tokens=max_output_tokens, + safety_settings=[ + types.SafetySetting( + category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, + threshold=types.HarmBlockThreshold.BLOCK_NONE, + ), + types.SafetySetting( + category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, + threshold=types.HarmBlockThreshold.BLOCK_NONE, + ), + types.SafetySetting( + category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, + threshold=types.HarmBlockThreshold.BLOCK_NONE, + ), + types.SafetySetting( + category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, + threshold=types.HarmBlockThreshold.BLOCK_NONE, + ), + types.SafetySetting( + category=types.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY, + threshold=types.HarmBlockThreshold.BLOCK_NONE, + ), + ], + response_modalities=response_modalities, + system_instruction=[ + types.Part.from_text(text=padded_system_instruction), + ], + ) else: generate_content_config = types.GenerateContentConfig( temperature=temperature, diff --git a/requirements.txt b/requirements.txt index abe5f0a..31903d2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -google-genai +google-genai>=1.25.0 numpy torch Pillow