Fix seed list flattening with context cache and seed-only batching

This commit is contained in:
silveroxides
2026-10-24 01:54:22 +02:00
parent d16be4fef5
commit 1d433c9931
4 changed files with 130 additions and 25 deletions
+46 -8
View File
@@ -1607,11 +1607,30 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
def _is_list_arg(val):
return isinstance(val, list) and not isinstance(val, (str, bytes, dict))
def _extract_seed_list(val):
if not isinstance(val, (list, tuple)):
return None
flat = []
for item in val:
if isinstance(item, (list, tuple)):
sub = _extract_seed_list(item)
if sub is not None:
flat.extend(sub)
elif torch.is_tensor(item):
flat.extend(int(x) for x in item.detach().cpu().reshape(-1).tolist())
elif isinstance(item, np.ndarray):
flat.extend(int(x) for x in item.reshape(-1).tolist())
elif item is not None:
flat.append(int(item))
return flat
is_prompt_list = _is_list_arg(prompt)
is_video_list = _is_list_arg(video)
is_seed_list = _is_list_arg(seed)
extracted_seeds = _extract_seed_list(seed)
is_seed_list = extracted_seeds is not None
is_config_list = _is_list_arg(config)
is_batch = is_prompt_list or is_video_list or is_seed_list
is_batch = is_prompt_list or is_video_list or is_seed_list or is_config_list
if not is_batch:
return cls._execute_single(
@@ -1628,12 +1647,14 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
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]
seed_list = extracted_seeds if is_seed_list else [seed]
config_list = config if is_config_list else [config]
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,
len(config_list) if is_config_list else 1,
)
def _unwrap_scalar(val, default=None):
@@ -1671,25 +1692,42 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
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)
has_multi_seed_input = is_seed_list and len(seed_list) > 1
def make_item_kwargs(idx, prepared_imgs=None, precomp_vid=None):
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])
if has_multi_seed_input:
if idx < len(seed_list):
s = int(seed_list[idx])
else:
base_s = int(seed_list[-1])
s = (base_s + (idx - len(seed_list) + 1)) if base_s != 0 else 0
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
s = (base_s + idx) if (base_s != 0 and batch_count > 1) else base_s
elif not is_seed_list and seed != 0:
s = int(seed) + idx
s = (int(seed) + idx) if batch_count > 1 else int(seed)
else:
s = int(_unwrap_scalar(seed, 0))
raw_cfg = _get_aligned(config_list, idx, config_val)
item_cfg = dict(raw_cfg) if isinstance(raw_cfg, dict) else dict(config_val)
if item_cfg.get("use_cache"):
if has_multi_seed_input:
item_cfg["cache_seed"] = s
elif batch_count > 1 and not (is_config_list and len(config_list) > 1):
base_cs = int(item_cfg.get("cache_seed", 0))
if base_cs != 0:
item_cfg["cache_seed"] = base_cs + idx
elif s != 0:
item_cfg["cache_seed"] = s
sys_inst = _get_aligned(system_instruction, idx, "You are a helpful AI assistant.")
return {
"config": config_val,
"config": item_cfg,
"prompt": p,
"system_instruction": sys_inst,
"model": model_val,
+72 -14
View File
@@ -623,7 +623,13 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
def _normalize_store_mode(cls, store: Any) -> str:
if store is None:
return "new"
if isinstance(store, bool):
return "new" if store else "disabled"
mode = str(store).strip().lower()
if mode == "true":
return "new"
if mode == "false":
return "disabled"
if mode not in cls.STORE_MODES:
raise ValueError(f"store must be one of: {', '.join(cls.STORE_MODES)}.")
return mode
@@ -1467,11 +1473,30 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
def _is_list_arg(val):
return isinstance(val, list) and not isinstance(val, (str, bytes, dict))
def _extract_seed_list(val):
if not isinstance(val, (list, tuple)):
return None
flat = []
for item in val:
if isinstance(item, (list, tuple)):
sub = _extract_seed_list(item)
if sub is not None:
flat.extend(sub)
elif torch.is_tensor(item):
flat.extend(int(x) for x in item.detach().cpu().reshape(-1).tolist())
elif isinstance(item, np.ndarray):
flat.extend(int(x) for x in item.reshape(-1).tolist())
elif item is not None:
flat.append(int(item))
return flat
is_prompt_list = _is_list_arg(prompt)
is_video_list = _is_list_arg(video)
is_seed_list = _is_list_arg(seed)
extracted_seeds = _extract_seed_list(seed)
is_seed_list = extracted_seeds is not None
is_config_list = _is_list_arg(config)
is_batch = is_prompt_list or is_video_list or is_seed_list
is_batch = is_prompt_list or is_video_list or is_seed_list or is_config_list
if not is_batch:
return cls._execute_single(
@@ -1489,12 +1514,14 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
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]
seed_list = extracted_seeds if is_seed_list else [seed]
config_list = config if is_config_list else [config]
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,
len(config_list) if is_config_list else 1,
)
def _unwrap_scalar(val, default=None):
@@ -1534,24 +1561,43 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
store_val = cls._normalize_store_mode(_unwrap_scalar(store, "new"))
resp_schema_val = _unwrap_scalar(response_schema, None)
is_prev_id_list = _is_list_arg(previous_interaction_id) and len(previous_interaction_id) > 1
has_multi_seed_input = is_seed_list and len(seed_list) > 1
has_multi_turn_list = (is_prompt_list and len(prompt_list) > 1) or (is_video_list and len(video_list) > 1)
def make_item_kwargs(idx, chained_prev_id="", prepared_imgs=None, precomp_vid=None):
def make_item_kwargs(idx, chained_prev_id="", prepared_imgs=None, precomp_vid=None, store_override=None):
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])
if has_multi_seed_input:
if idx < len(seed_list):
s = int(seed_list[idx])
else:
base_s = int(seed_list[-1])
s = (base_s + (idx - len(seed_list) + 1)) if base_s != 0 else 0
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
s = (base_s + idx) if (base_s != 0 and batch_count > 1) else base_s
elif not is_seed_list and seed != 0:
s = int(seed) + idx
s = (int(seed) + idx) if batch_count > 1 else int(seed)
else:
s = int(_unwrap_scalar(seed, 0))
raw_cfg = _get_aligned(config_list, idx, config_val)
item_cfg = dict(raw_cfg) if isinstance(raw_cfg, dict) else dict(config_val)
if item_cfg.get("use_cache"):
if has_multi_seed_input:
item_cfg["cache_seed"] = s
elif batch_count > 1 and not (is_config_list and len(config_list) > 1):
base_cs = int(item_cfg.get("cache_seed", 0))
if base_cs != 0:
item_cfg["cache_seed"] = base_cs + idx
elif s != 0:
item_cfg["cache_seed"] = s
sys_inst = _get_aligned(system_instruction, idx, "You are a helpful AI assistant.")
if store_val == "manual":
if is_prev_id_list or idx == 0:
eff_store = store_override if store_override is not None else store_val
if eff_store == "manual":
if is_prev_id_list or idx == 0 or not has_multi_turn_list:
prev_id = str(_get_aligned(previous_interaction_id, idx, "") or "").strip()
else:
prev_id = ""
@@ -1559,7 +1605,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
prev_id = ""
return {
"config": config_val,
"config": item_cfg,
"prompt": p,
"system_instruction": sys_inst,
"model": model_val,
@@ -1584,7 +1630,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
"max_retries": max_retries_val,
"timeout_fallback_text": timeout_fb_val,
"image_size": img_size_val,
"store": store_val,
"store": eff_store,
"previous_interaction_id": prev_id,
"video": v,
"image_inputs": image_inputs,
@@ -1629,18 +1675,30 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
batch_results = []
chained_interaction_id = ""
initial_previous_id = ""
if not has_multi_turn_list and store_val == "previous":
initial_previous_id = GeminiInteractionsStoreManager.get_last_interaction_id()
for idx in range(batch_count):
cls._check_interrupt()
if has_multi_turn_list:
item_chained_id = chained_interaction_id
item_store_override = None
else:
# Seed/config variations on the same prompt & video: do not chain variations into each other
item_chained_id = initial_previous_id
item_store_override = "manual" if (store_val == "previous" and idx > 0) else None
item_res = cls._execute_single(
**make_item_kwargs(
idx,
chained_prev_id=chained_interaction_id,
chained_prev_id=item_chained_id,
prepared_imgs=prepared_images,
precomp_vid=precomputed_video,
store_override=item_store_override,
)
)
batch_results.append(item_res)
if store_val != "disabled" and item_res[5]:
if has_multi_turn_list and store_val != "disabled" and item_res[5]:
chained_interaction_id = str(item_res[5]).strip()
return IO.NodeOutput(
+5 -2
View File
@@ -1761,13 +1761,16 @@ def test_gemini_batch_seeds(structured_client):
[("final", "Resp 1")],
[("final", "Resp 2")],
]
args = _execute_kwargs(_config())
args.update(prompt="Single prompt", use_seed=True, seed=[101, 102])
args = _execute_kwargs(_config(use_cache=True, cache_seed=615665156))
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[0] == ["Resp 1", "Resp 2"]
assert output[2] == [101, 102]
assert len(structured_client.calls) == 2
assert structured_client.calls[0]["config"].seed == 101
assert structured_client.calls[1]["config"].seed == 102
def test_gemini_batch_single_list_element(structured_client):
+7 -1
View File
@@ -778,12 +778,18 @@ def test_interactions_batch_seeds(mock_interactions_client):
{"text": "A1", "id": "i1"},
{"text": "A2", "id": "i2"},
]
args = _execute_kwargs(prompt="Single", use_seed=True, seed=[501, 502])
config = _config(use_cache=True, cache_seed=615665156)
args = _execute_kwargs(config=config, prompt="Single", use_seed=True, seed=[[501, 502]], store="new")
output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args)
assert isinstance(output[0], list) and len(output[0]) == 2
assert output[0] == ["A1", "A2"]
assert output[2] == [501, 502]
assert len(mock_interactions_client.calls) == 2
assert mock_interactions_client.calls[0]["generation_config"]["seed"] == 501
assert mock_interactions_client.calls[1]["generation_config"]["seed"] == 502
assert "previous_interaction_id" not in mock_interactions_client.calls[0]
assert "previous_interaction_id" not in mock_interactions_client.calls[1]
def test_interactions_batch_videos(mock_interactions_client):