update comfy api imports

This commit is contained in:
silveroxides
2025-12-29 19:35:51 +01:00
parent dbd2be83c8
commit cc154b82ae
+56 -57
View File
@@ -8,8 +8,7 @@ import cv2
from PIL import Image
import io
import folder_paths # type: ignore[reportMissingImports]
from comfy_api.latest import ComfyExtension, ui # type: ignore[reportMissingImports]
from comfy_api.latest import io as cio
from comfy_api.latest import ComfyExtension, UI, IO # type: ignore[reportMissingImports]
from google import genai
from google.genai import types
import time
@@ -45,26 +44,26 @@ try:
except Exception as e:
print(f"[WARNING] Error checking dependencies: {str(e)}")
class GetKeyAPI(cio.ComfyNode):
class GetKeyAPI(IO.ComfyNode):
@classmethod
def define_schema(cls) -> cio.Schema:
return cio.Schema(
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="GetKeyAPI",
display_name="Get API Key from JSON",
category="utils/api_keys",
inputs=[
cio.String.Input("json_path", default="./input/apikeys.json", multiline=False, tooltip="Path to a .json file with simple top level structure with name as key and api-key as value. See example in custom node folder."),
cio.Combo.Input("key_id_method", options=["custom", "random_rotate", "increment_rotate"], default="custom", tooltip="custom sets api-key to the api-key with the name set in the key_id widget. random_rotate randomly switches between keys if multiple in the .json and increment_rotate does it in order from first to last, then repeats."),
cio.Int.Input("rotation_interval", default=0, min=0, tooltip="how many steps to jump when doing rotate."),
cio.String.Input("key_id", default="placeholder", multiline=False, optional=True, tooltip="Put name of key in the .json here if using custom in key_id_method."),
IO.String.Input("json_path", default="./input/apikeys.json", multiline=False, tooltip="Path to a .json file with simple top level structure with name as key and api-key as value. See example in custom node folder."),
IO.Combo.Input("key_id_method", options=["custom", "random_rotate", "increment_rotate"], default="custom", tooltip="custom sets api-key to the api-key with the name set in the key_id widget. random_rotate randomly switches between keys if multiple in the .json and increment_rotate does it in order from first to last, then repeats."),
IO.Int.Input("rotation_interval", default=0, min=0, tooltip="how many steps to jump when doing rotate."),
IO.String.Input("key_id", default="placeholder", multiline=False, optional=True, tooltip="Put name of key in the .json here if using custom in key_id_method."),
],
outputs=[
cio.String.Output("API_KEY")
IO.String.Output("API_KEY")
]
)
@classmethod
def execute(cls, json_path: str, key_id_method: str, rotation_interval: int, key_id: str | None = "placeholder") -> cio.NodeOutput:
def execute(cls, json_path: str, key_id_method: str, rotation_interval: int, key_id: str | None = "placeholder") -> IO.NodeOutput:
api_keys_data = None
absolute_json_path = os.path.abspath(json_path)
@@ -117,25 +116,25 @@ class GetKeyAPI(cio.ComfyNode):
raise ValueError(f"RotateKeyAPI Error: Retrieved value for selected key is not a valid string. Value: {selected_key_value}")
print(f"RotateKeyAPI: Successfully retrieved API key using method '{key_id_method}'.")
return cio.NodeOutput(selected_key_value)
return IO.NodeOutput(selected_key_value)
class SSL_GeminiAPIKeyConfig(cio.ComfyNode):
GemConfig = cio.Custom("GEMINI_CONFIG")
class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
GemConfig = IO.Custom("GEMINI_CONFIG")
@classmethod
def define_schema(cls) -> cio.Schema:
return cio.Schema(
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="SSL_GeminiAPIKeyConfig",
display_name="Configure Gemini API Key",
category="API/Gemini",
inputs=[
cio.String.Input("api_key", multiline=False),
cio.Combo.Input("api_version", options=["v1", "v1alpha", "v1beta", "v2beta"], default="v1alpha"),
cio.Boolean.Input("vertexai", default=False),
cio.String.Input("vertexai_project", default="placeholder", optional=True),
cio.String.Input("vertexai_location", default="placeholder", optional=True),
IO.String.Input("api_key", multiline=False),
IO.Combo.Input("api_version", options=["v1", "v1alpha", "v1beta", "v2beta"], default="v1alpha"),
IO.Boolean.Input("vertexai", default=False),
IO.String.Input("vertexai_project", default="placeholder", optional=True),
IO.String.Input("vertexai_location", default="placeholder", optional=True),
],
outputs=[
cls.GemConfig.Output("config")
@@ -143,14 +142,14 @@ class SSL_GeminiAPIKeyConfig(cio.ComfyNode):
)
@classmethod
def execute(cls, api_key: str, api_version: str, vertexai: bool, vertexai_project: str | None = "placeholder", vertexai_location: str | None = "placeholder") -> cio.NodeOutput:
def execute(cls, api_key: str, api_version: str, vertexai: bool, vertexai_project: str | None = "placeholder", vertexai_location: str | None = "placeholder") -> IO.NodeOutput:
config = {"api_key": api_key, "api_version": api_version, "vertexai": vertexai, "vertexai_project": vertexai_project, "vertexai_location": vertexai_location}
return cio.NodeOutput(config)
return IO.NodeOutput(config)
class SSL_GeminiTextPrompt(cio.ComfyNode):
GemConfig = cio.Custom("GEMINI_CONFIG")
class SSL_GeminiTextPrompt(IO.ComfyNode):
GemConfig = IO.Custom("GEMINI_CONFIG")
_cache: dict = {}
# Define model lists centrally to ensure consistency between cache logic and execution logic
@@ -171,40 +170,40 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
]
@classmethod
def define_schema(cls) -> cio.Schema:
return cio.Schema(
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="SSL_GeminiTextPrompt",
display_name="Expanded Gemini Text/Image",
category="API/Gemini",
inputs=[
cls.GemConfig.Input("config"),
cio.String.Input("prompt", multiline=True),
cio.String.Input("system_instruction", default="You are a helpful AI assistant.", multiline=True),
cio.Combo.Input("model", options=["learnlm-2.0-flash-experimental", "gemini-exp-1206", "gemini-2.0-flash", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-flash-thinking-exp-1219", "gemini-2.5-pro", "gemini-2.5-pro-preview-05-06", "gemini-2.5-flash", "gemini-2.5-flash-preview-09-2025", "gemini-2.5-flash-lite-preview-09-2025", "gemini-2.5-flash-lite", "gemini-3-pro-preview", "gemini-2.5-flash-image-preview", "nano-banana-pro-preview", "gemini-3-flash-preview"], default="gemini-2.0-flash"),
cio.Float.Input("temperature", default=1.0, min=0.0, max=1.0, step=0.01),
cio.Float.Input("top_p", default=0.95, min=0.0, max=1.0, step=0.01),
cio.Int.Input("top_k", default=40, min=1, max=100, step=1),
cio.Int.Input("max_output_tokens", default=8192, min=1, max=65536, step=1),
cio.Boolean.Input("include_images", default=False),
cio.Combo.Input("aspect_ratio", options=["None", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"], default="None"),
cio.Combo.Input("bypass_mode", options=["None", "system_instruction", "prompt", "both"], default="None"),
cio.Int.Input("thinking_budget", default=0, min=-1, max=24576, step=1, tooltip="0 disables thinking mode, -1 will activate it as default dynamic thinking and anything above 0 sets specific budget"),
cio.Image.Input("input_image", optional=True),
cio.Image.Input("input_image_2", optional=True),
cio.Boolean.Input("use_proxy", default=False),
cio.String.Input("proxy_host", default="127.0.0.1"),
cio.Int.Input("proxy_port", default=7890, min=1, max=65535),
cio.Boolean.Input("use_seed", default=True),
cio.Int.Input("seed", default=0, min=0, max=2147483647),
cio.Int.Input("timeout", default=30, min=15, max=300, step=15),
cio.Boolean.Input("include_thoughts", default=False),
cio.Combo.Input("thinking_level", options=["None", "low", "medium", "high"], default="None", tooltip="Does not work at the same time as 'thinking_budget'. if this is set, then thinking budget is ignored."),
cio.Combo.Input("media_resolution", options=["unspecified", "low", "medium", "high"], default="unspecified", tooltip="Set input media resolution for image, video and pdf. This changes tokens consumed."),
IO.String.Input("prompt", multiline=True),
IO.String.Input("system_instruction", default="You are a helpful AI assistant.", multiline=True),
IO.Combo.Input("model", options=["learnlm-2.0-flash-experimental", "gemini-exp-1206", "gemini-2.0-flash", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-flash-thinking-exp-1219", "gemini-2.5-pro", "gemini-2.5-pro-preview-05-06", "gemini-2.5-flash", "gemini-2.5-flash-preview-09-2025", "gemini-2.5-flash-lite-preview-09-2025", "gemini-2.5-flash-lite", "gemini-3-pro-preview", "gemini-2.5-flash-image-preview", "nano-banana-pro-preview", "gemini-3-flash-preview"], default="gemini-2.0-flash"),
IO.Float.Input("temperature", default=1.0, min=0.0, max=1.0, step=0.01),
IO.Float.Input("top_p", default=0.95, min=0.0, max=1.0, step=0.01),
IO.Int.Input("top_k", default=40, min=1, max=100, step=1),
IO.Int.Input("max_output_tokens", default=8192, min=1, max=65536, step=1),
IO.Boolean.Input("include_images", default=False),
IO.Combo.Input("aspect_ratio", options=["None", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"], default="None"),
IO.Combo.Input("bypass_mode", options=["None", "system_instruction", "prompt", "both"], default="None"),
IO.Int.Input("thinking_budget", default=0, min=-1, max=24576, step=1, tooltip="0 disables thinking mode, -1 will activate it as default dynamic thinking and anything above 0 sets specific budget"),
IO.Image.Input("input_image", optional=True),
IO.Image.Input("input_image_2", optional=True),
IO.Boolean.Input("use_proxy", default=False),
IO.String.Input("proxy_host", default="127.0.0.1"),
IO.Int.Input("proxy_port", default=7890, min=1, max=65535),
IO.Boolean.Input("use_seed", default=True),
IO.Int.Input("seed", default=0, min=0, max=2147483647),
IO.Int.Input("timeout", default=30, min=15, max=300, step=15),
IO.Boolean.Input("include_thoughts", default=False),
IO.Combo.Input("thinking_level", options=["None", "low", "medium", "high"], default="None", tooltip="Does not work at the same time as 'thinking_budget'. if this is set, then thinking budget is ignored."),
IO.Combo.Input("media_resolution", options=["unspecified", "low", "medium", "high"], default="unspecified", tooltip="Set input media resolution for image, video and pdf. This changes tokens consumed."),
],
outputs=[
cio.String.Output("text"),
cio.Image.Output("image"),
cio.Int.Output("final_actual_seed")
IO.String.Output("text"),
IO.Image.Output("image"),
IO.Int.Output("final_actual_seed")
]
)
@@ -455,7 +454,7 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget, input_image=None, input_image_2=None,
use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=False, seed=0, timeout=30,
include_thoughts=False, thinking_level=None, media_resolution=None) -> cio.NodeOutput:
include_thoughts=False, thinking_level=None, media_resolution=None) -> IO.NodeOutput:
fingerprint, cached = cls._compute_fingerprint_and_check_cache(
config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
@@ -468,7 +467,7 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
if cached is not None:
cached_text, cached_image, cached_seed = cached
print(f"[INFO] Returning cached result for fingerprint {fingerprint}")
return cio.NodeOutput(cached_text, cached_image, cached_seed)
return IO.NodeOutput(cached_text, cached_image, cached_seed)
# --- keep most of the original implementation but converted to classmethod usage ---
original_http_proxy = os.environ.get('HTTP_PROXY')
@@ -545,7 +544,7 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
client = genai.Client(api_key=config.get("api_key"), http_options=types.HttpOptions(api_version=config.get("api_version")), **client_options) # type: ignore
except Exception as e:
print(f"[ERROR] Gemini client initialization failed: {str(e)}")
return cio.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
return IO.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
# Network test (best-effort)
try:
@@ -592,7 +591,7 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
contents = img_parts + [{"text": padded_prompt}]
except Exception as e:
print(f"[ERROR] Error processing input image: {str(e)}")
return cio.NodeOutput(f"Error processing input image: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
return IO.NodeOutput(f"Error processing input image: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
else:
contents = padded_prompt
@@ -721,7 +720,7 @@ class SSL_GeminiTextPrompt(cio.ComfyNode):
pass
try:
return cio.NodeOutput(text_output, image_tensor, final_actual_seed)
return IO.NodeOutput(text_output, image_tensor, final_actual_seed)
finally:
if original_http_proxy:
os.environ['HTTP_PROXY'] = original_http_proxy