update comfy api imports
This commit is contained in:
+56
-57
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user