Support is_input_list and concurrent batching on text prompt nodes

This commit is contained in:
silveroxides
2026-10-16 12:41:24 +02:00
parent 5fe0479c43
commit 725a3d2a83
4 changed files with 457 additions and 29 deletions
+172 -14
View File
@@ -21,6 +21,7 @@ import random
import hashlib
import warnings
from typing import Any, Tuple
from concurrent.futures import ThreadPoolExecutor
class GetKeyAPI(IO.ComfyNode):
@classmethod
@@ -295,6 +296,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
node_id="SSL_GeminiTextPrompt",
display_name="Expanded Gemini Text/Image",
category="API/Gemini",
is_input_list=True,
inputs=[
cls.GemConfig.Input("config"),
IO.String.Input("prompt", multiline=True),
@@ -336,11 +338,11 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
),
],
outputs=[
IO.String.Output("text"),
IO.Image.Output("image"),
IO.Int.Output("final_actual_seed"),
IO.String.Output("structured_output", display_name="Structured output"),
IO.String.Output("thoughts", display_name="Thoughts"),
IO.String.Output("text", is_output_list=True),
IO.Image.Output("image", is_output_list=True),
IO.Int.Output("final_actual_seed", is_output_list=True),
IO.String.Output("structured_output", display_name="Structured output", is_output_list=True),
IO.String.Output("thoughts", display_name="Thoughts", is_output_list=True),
]
)
@@ -418,6 +420,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
):
if image_batch is None:
continue
if isinstance(image_batch, list):
image_batch = image_batch[0] if image_batch else None
if image_batch is None:
continue
if not torch.is_tensor(image_batch) or image_batch.ndim != 4:
raise ValueError(
f"Gemini {socket_name} must be a BHWC image tensor."
@@ -925,15 +931,15 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
raise ValueError(f"Structured response contains invalid JSON constant: {value}.")
@classmethod
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
def _execute_single(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}")
use_cache = bool(config.get("use_cache", False))
@@ -1498,4 +1504,156 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
return IO.NodeOutput(text_output, image_tensor, final_actual_seed, structured_output, thoughts_output)
@classmethod
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
def _is_list_arg(val):
return isinstance(val, list) and not isinstance(val, (str, bytes, dict))
is_prompt_list = _is_list_arg(prompt)
is_video_list = _is_list_arg(video)
is_seed_list = _is_list_arg(seed)
is_batch = is_prompt_list or is_video_list or is_seed_list
if not is_batch:
return cls._execute_single(
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,
)
prompt_list = prompt if is_prompt_list else [prompt]
video_list = video if is_video_list else [video]
seed_list = seed if is_seed_list else [seed]
batch_count = max(
len(prompt_list) if is_prompt_list else 1,
len(video_list) if is_video_list else 1,
len(seed_list) if is_seed_list else 1,
)
def _unwrap_scalar(val, default=None):
while isinstance(val, list) and len(val) > 0 and (isinstance(val[0], list) or not isinstance(val, dict)):
val = val[0]
if isinstance(val, list) and len(val) == 1:
val = val[0]
return val if val is not None else default
def _get_aligned(val, idx, default=None):
if isinstance(val, list) and val:
return val[min(idx, len(val) - 1)]
return val if val is not None else default
config_val = _unwrap_scalar(config, {})
model_val = _unwrap_scalar(model, "gemini-2.5-flash")
temp_val = float(_unwrap_scalar(temperature, 1.0))
top_p_val = float(_unwrap_scalar(top_p, 0.95))
top_k_val = int(_unwrap_scalar(top_k, 40))
max_tokens_val = int(_unwrap_scalar(max_output_tokens, 8192))
inc_imgs_val = bool(_unwrap_scalar(include_images, False))
aspect_val = str(_unwrap_scalar(aspect_ratio, "None"))
bypass_val = str(_unwrap_scalar(bypass_mode, "None"))
budget_val = int(_unwrap_scalar(thinking_budget, 0))
proxy_val = bool(_unwrap_scalar(use_proxy, False))
phost_val = str(_unwrap_scalar(proxy_host, "127.0.0.1"))
pport_val = int(_unwrap_scalar(proxy_port, 7890))
use_seed_val = bool(_unwrap_scalar(use_seed, False))
timeout_val = int(_unwrap_scalar(timeout, 30))
thoughts_val = bool(_unwrap_scalar(include_thoughts, False))
th_level_val = _unwrap_scalar(thinking_level, None)
media_res_val = _unwrap_scalar(media_resolution, None)
retry_pat_val = str(_unwrap_scalar(retry_pattern, ""))
max_retries_val = int(_unwrap_scalar(max_retries, 3))
timeout_fb_val = str(_unwrap_scalar(timeout_fallback_text, ""))
img_size_val = str(_unwrap_scalar(image_size, "None"))
resp_schema_val = _unwrap_scalar(response_schema, None)
def make_item_kwargs(idx):
p = prompt_list[min(idx, len(prompt_list) - 1)] if prompt_list else ""
v = video_list[min(idx, len(video_list) - 1)] if video_list else None
if is_seed_list and len(seed_list) >= batch_count:
s = int(seed_list[idx])
elif is_seed_list and len(seed_list) == 1:
base_s = int(seed_list[0])
s = (base_s + idx) if base_s != 0 else 0
elif not is_seed_list and seed != 0:
s = int(seed) + idx
else:
s = int(_unwrap_scalar(seed, 0))
sys_inst = _get_aligned(system_instruction, idx, "You are a helpful AI assistant.")
return {
"config": config_val,
"prompt": p,
"system_instruction": sys_inst,
"model": model_val,
"temperature": temp_val,
"top_p": top_p_val,
"top_k": top_k_val,
"max_output_tokens": max_tokens_val,
"include_images": inc_imgs_val,
"aspect_ratio": aspect_val,
"bypass_mode": bypass_val,
"thinking_budget": budget_val,
"use_proxy": proxy_val,
"proxy_host": phost_val,
"proxy_port": pport_val,
"use_seed": use_seed_val,
"seed": s,
"timeout": timeout_val,
"include_thoughts": thoughts_val,
"thinking_level": th_level_val,
"media_resolution": media_res_val,
"retry_pattern": retry_pat_val,
"max_retries": max_retries_val,
"timeout_fallback_text": timeout_fb_val,
"image_size": img_size_val,
"video": v,
"image_inputs": image_inputs,
"response_schema": resp_schema_val,
}
if batch_count == 1:
single_res = cls._execute_single(**make_item_kwargs(0))
return IO.NodeOutput(
[single_res[0]],
[single_res[1]],
[single_res[2]],
[single_res[3]],
[single_res[4]],
)
max_workers = min(batch_count, 10)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(cls._execute_single, **make_item_kwargs(idx))
for idx in range(batch_count)
]
batch_results = [f.result() for f in futures]
return IO.NodeOutput(
[r[0] for r in batch_results],
[r[1] for r in batch_results],
[r[2] for r in batch_results],
[r[3] for r in batch_results],
[r[4] for r in batch_results],
)
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes
+180 -15
View File
@@ -23,6 +23,7 @@ import hashlib
import warnings
from typing import Any, Tuple
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from .gemini_nodes import (
GetKeyAPI,
@@ -319,6 +320,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
node_id="SSL_GeminiInteractionsTextPrompt",
display_name="Expanded Gemini Interactions (Text/Image)",
category="API/Gemini/Interactions",
is_input_list=True,
inputs=[
cls.GemConfig.Input("config"),
IO.String.Input("prompt", multiline=True),
@@ -362,12 +364,12 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
),
],
outputs=[
IO.String.Output("text"),
IO.Image.Output("image"),
IO.Int.Output("final_actual_seed"),
IO.String.Output("structured_output", display_name="Structured output"),
IO.String.Output("thoughts", display_name="Thoughts"),
IO.String.Output("interaction_id", display_name="Interaction ID"),
IO.String.Output("text", is_output_list=True),
IO.Image.Output("image", is_output_list=True),
IO.Int.Output("final_actual_seed", is_output_list=True),
IO.String.Output("structured_output", display_name="Structured output", is_output_list=True),
IO.String.Output("thoughts", display_name="Thoughts", is_output_list=True),
IO.String.Output("interaction_id", display_name="Interaction ID", is_output_list=True),
]
)
@@ -443,6 +445,10 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
):
if image_batch is None:
continue
if isinstance(image_batch, list):
image_batch = image_batch[0] if image_batch else None
if image_batch is None:
continue
if not torch.is_tensor(image_batch) or image_batch.ndim != 4:
raise ValueError(
f"Gemini {socket_name} must be a BHWC image tensor."
@@ -739,15 +745,15 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
raise ValueError(f"Structured response contains invalid JSON constant: {value}.")
@classmethod
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None", store=True, previous_interaction_id="",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
def _execute_single(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None", store=True, previous_interaction_id="",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
print(f"[INFO] SSL_GeminiInteractionsTextPrompt execute called, model: {model}")
@@ -1351,6 +1357,165 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
return IO.NodeOutput(text_output, image_tensor, final_actual_seed, structured_output, thoughts_output, interaction_id_output)
@classmethod
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget,
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,
retry_pattern="", max_retries=3, timeout_fallback_text="",
image_size="None", store=True, previous_interaction_id="",
video: GeminiVideoConfig.Type | None = None,
image_inputs: IO.Autogrow.Type | None = None,
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
def _is_list_arg(val):
return isinstance(val, list) and not isinstance(val, (str, bytes, dict))
is_prompt_list = _is_list_arg(prompt)
is_video_list = _is_list_arg(video)
is_seed_list = _is_list_arg(seed)
is_batch = is_prompt_list or is_video_list or is_seed_list
if not is_batch:
return cls._execute_single(
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, store=store,
previous_interaction_id=previous_interaction_id,
video=video, image_inputs=image_inputs, response_schema=response_schema,
)
prompt_list = prompt if is_prompt_list else [prompt]
video_list = video if is_video_list else [video]
seed_list = seed if is_seed_list else [seed]
batch_count = max(
len(prompt_list) if is_prompt_list else 1,
len(video_list) if is_video_list else 1,
len(seed_list) if is_seed_list else 1,
)
def _unwrap_scalar(val, default=None):
while isinstance(val, list) and len(val) > 0 and (isinstance(val[0], list) or not isinstance(val, dict)):
val = val[0]
if isinstance(val, list) and len(val) == 1:
val = val[0]
return val if val is not None else default
def _get_aligned(val, idx, default=None):
if isinstance(val, list) and val:
return val[min(idx, len(val) - 1)]
return val if val is not None else default
config_val = _unwrap_scalar(config, {})
model_val = _unwrap_scalar(model, "gemini-2.5-flash")
temp_val = float(_unwrap_scalar(temperature, 1.0))
top_p_val = float(_unwrap_scalar(top_p, 0.95))
top_k_val = int(_unwrap_scalar(top_k, 40))
max_tokens_val = int(_unwrap_scalar(max_output_tokens, 8192))
inc_imgs_val = bool(_unwrap_scalar(include_images, False))
aspect_val = str(_unwrap_scalar(aspect_ratio, "None"))
bypass_val = str(_unwrap_scalar(bypass_mode, "None"))
budget_val = int(_unwrap_scalar(thinking_budget, 0))
proxy_val = bool(_unwrap_scalar(use_proxy, False))
phost_val = str(_unwrap_scalar(proxy_host, "127.0.0.1"))
pport_val = int(_unwrap_scalar(proxy_port, 7890))
use_seed_val = bool(_unwrap_scalar(use_seed, False))
timeout_val = int(_unwrap_scalar(timeout, 30))
thoughts_val = bool(_unwrap_scalar(include_thoughts, False))
th_level_val = _unwrap_scalar(thinking_level, None)
media_res_val = _unwrap_scalar(media_resolution, None)
retry_pat_val = str(_unwrap_scalar(retry_pattern, ""))
max_retries_val = int(_unwrap_scalar(max_retries, 3))
timeout_fb_val = str(_unwrap_scalar(timeout_fallback_text, ""))
img_size_val = str(_unwrap_scalar(image_size, "None"))
store_val = bool(_unwrap_scalar(store, True))
resp_schema_val = _unwrap_scalar(response_schema, None)
def make_item_kwargs(idx):
p = prompt_list[min(idx, len(prompt_list) - 1)] if prompt_list else ""
v = video_list[min(idx, len(video_list) - 1)] if video_list else None
if is_seed_list and len(seed_list) >= batch_count:
s = int(seed_list[idx])
elif is_seed_list and len(seed_list) == 1:
base_s = int(seed_list[0])
s = (base_s + idx) if base_s != 0 else 0
elif not is_seed_list and seed != 0:
s = int(seed) + idx
else:
s = int(_unwrap_scalar(seed, 0))
sys_inst = _get_aligned(system_instruction, idx, "You are a helpful AI assistant.")
prev_id = _get_aligned(previous_interaction_id, idx, "")
return {
"config": config_val,
"prompt": p,
"system_instruction": sys_inst,
"model": model_val,
"temperature": temp_val,
"top_p": top_p_val,
"top_k": top_k_val,
"max_output_tokens": max_tokens_val,
"include_images": inc_imgs_val,
"aspect_ratio": aspect_val,
"bypass_mode": bypass_val,
"thinking_budget": budget_val,
"use_proxy": proxy_val,
"proxy_host": phost_val,
"proxy_port": pport_val,
"use_seed": use_seed_val,
"seed": s,
"timeout": timeout_val,
"include_thoughts": thoughts_val,
"thinking_level": th_level_val,
"media_resolution": media_res_val,
"retry_pattern": retry_pat_val,
"max_retries": max_retries_val,
"timeout_fallback_text": timeout_fb_val,
"image_size": img_size_val,
"store": store_val,
"previous_interaction_id": prev_id,
"video": v,
"image_inputs": image_inputs,
"response_schema": resp_schema_val,
}
if batch_count == 1:
single_res = cls._execute_single(**make_item_kwargs(0))
return IO.NodeOutput(
[single_res[0]],
[single_res[1]],
[single_res[2]],
[single_res[3]],
[single_res[4]],
[single_res[5]],
)
max_workers = min(batch_count, 10)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(cls._execute_single, **make_item_kwargs(idx))
for idx in range(batch_count)
]
batch_results = [f.result() for f in futures]
return IO.NodeOutput(
[r[0] for r in batch_results],
[r[1] for r in batch_results],
[r[2] for r in batch_results],
[r[3] for r in batch_results],
[r[4] for r in batch_results],
[r[5] for r in batch_results],
)
class SSL_GeminiInteractionsStore(IO.ComfyNode):
@classmethod
+61
View File
@@ -1695,3 +1695,64 @@ def test_gemini_standard_auto_safety_resolution(structured_client):
assert len(sent_safety) == 9
assert all("BLOCK_NONE" in str(s.threshold) for s in sent_safety)
assert all(s.method is None for s in sent_safety)
def test_gemini_batch_prompts(structured_client):
structured_client.responses = [
[("final", "Response 1")],
[("final", "Response 2")],
[("final", "Response 3")],
]
args = _execute_kwargs(_config())
args["prompt"] = ["Prompt one", "Prompt two", "Prompt three"]
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
assert isinstance(output[0], list)
assert output[0] == ["Response 1", "Response 2", "Response 3"]
assert isinstance(output[1], list) and len(output[1]) == 3
assert isinstance(output[2], list) and len(output[2]) == 3
assert len(structured_client.calls) == 3
def test_gemini_batch_seeds(structured_client):
structured_client.responses = [
[("final", "Resp 1")],
[("final", "Resp 2")],
]
args = _execute_kwargs(_config())
args.update(prompt="Single prompt", use_seed=True, seed=[101, 102])
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
assert isinstance(output[0], list) and len(output[0]) == 2
assert output[2] == [101, 102]
assert len(structured_client.calls) == 2
def test_gemini_batch_single_list_element(structured_client):
structured_client.responses = [
[("final", "Only response")],
]
args = _execute_kwargs(_config())
args["prompt"] = ["Only one prompt"]
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
assert isinstance(output[0], list)
assert output[0] == ["Only response"]
assert len(structured_client.calls) == 1
def test_gemini_batch_videos(structured_client):
structured_client.responses = [
[("final", "Vid 1")],
[("final", "Vid 2")],
]
args = _execute_kwargs(_config())
args.update(
prompt="Describe video",
video=[{"video": None, "fps": 1}, {"video": None, "fps": 1}],
)
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
assert isinstance(output[0], list)
assert output[0] == ["Vid 1", "Vid 2"]
assert len(structured_client.calls) == 2
+44
View File
@@ -685,3 +685,47 @@ def test_standard_auto_safety_resolution(mock_interactions_client):
assert all("method" not in s for s in call["safety_settings"])
def test_interactions_batch_prompts(mock_interactions_client):
mock_interactions_client.responses = [
{"text": "Answer 1", "id": "int_b1"},
{"text": "Answer 2", "id": "int_b2"},
]
args = _execute_kwargs(prompt=["Prompt A", "Prompt B"])
output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args)
assert isinstance(output[0], list)
assert output[0] == ["Answer 1", "Answer 2"]
assert isinstance(output[5], list)
assert output[5] == ["int_b1", "int_b2"]
assert len(mock_interactions_client.calls) == 2
def test_interactions_batch_seeds(mock_interactions_client):
mock_interactions_client.responses = [
{"text": "A1", "id": "i1"},
{"text": "A2", "id": "i2"},
]
args = _execute_kwargs(prompt="Single", use_seed=True, seed=[501, 502])
output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args)
assert isinstance(output[0], list) and len(output[0]) == 2
assert output[2] == [501, 502]
assert len(mock_interactions_client.calls) == 2
def test_interactions_batch_videos(mock_interactions_client):
mock_interactions_client.responses = [
{"text": "VidResp 1", "id": "iv1"},
{"text": "VidResp 2", "id": "iv2"},
]
args = _execute_kwargs(
prompt="Describe video",
video=[{"video": None, "fps": 1}, {"video": None, "fps": 1}],
)
output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args)
assert isinstance(output[0], list)
assert output[0] == ["VidResp 1", "VidResp 2"]
assert len(mock_interactions_client.calls) == 2