nodes/silveroxides-ComfyUI_Gemini_Expanded_API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(gemini): preserve image batch order

Browse Source
This commit is contained in:
silveroxides
2026-08-16 07:19:05 +02:00
parent 64d5010a81
commit 44e471c6e8
2 changed files with 88 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini_nodes.py
+42 -21
View File
@@ -283,27 +283,32 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
return tensor
@classmethod
def _flatten_image_inputs(cls, image_inputs: IO.Autogrow.Type | None):
def _ordered_image_frames(cls, image_inputs: IO.Autogrow.Type | None):
if not image_inputs:
return []
return [], []
def image_index(name):
match = re.fullmatch(r"image_(\d+)", name)
return int(match.group(1)) if match else 101
flattened = []
for _, image_batch in sorted(
frames = []
batch_counts = []
for socket_name, image_batch in sorted(
image_inputs.items(), key=lambda item: (image_index(item[0]), item[0])
):
if image_batch is None:
continue
if not torch.is_tensor(image_batch) or image_batch.ndim != 4:
raise ValueError("Gemini image inputs must be BHWC image tensors.")
flattened.extend(
image_batch[index : index + 1]
for index in range(image_batch.shape[0])
raise ValueError(
f"Gemini {socket_name} must be a BHWC image tensor."
)
batch_size = int(image_batch.shape[0])
batch_counts.append((socket_name, batch_size))
frames.extend(
(socket_name, index, image_batch[index])
for index in range(batch_size)
)
return flattened
return frames, batch_counts
@classmethod
def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
@@ -323,10 +328,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
print(f"[WARNING] Cache hashing failed for image: {e}")
return "Error"
image_hashes = tuple(
get_tensor_hash(image)
for image in cls._flatten_image_inputs(image_inputs)
)
image_frames, _ = cls._ordered_image_frames(image_inputs)
image_hashes = tuple(get_tensor_hash(frame) for _, _, frame in image_frames)
# 2. Determine Effective Parameters based on Model
# This ensures we don't cache-miss if an irrelevant parameter changes
@@ -699,25 +702,43 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
return IO.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
# Prepare contents (images + prompt)
images_to_process = cls._flatten_image_inputs(image_inputs)
image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs)
if images_to_process:
if image_frames:
try:
img_parts = []
for img in images_to_process:
img_array = img[0].cpu().numpy()
for _, _, image_frame in image_frames:
img_array = image_frame.cpu().numpy()
img_array = (img_array * 255).astype(np.uint8)
pil_img = Image.fromarray(img_array)
print(f"[DEBUG] Input image format: {pil_img.mode}")
img_byte_arr = BytesIO()
pil_img.save(img_byte_arr, format='PNG')
img_bytes = img_byte_arr.getvalue()
if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified":
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}, "media_resolution": {"level": f"MEDIA_RESOLUTION_{media_resolution.upper()}"}}
img_part = types.Part.from_bytes(
data=img_bytes,
mime_type="image/png",
media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}",
)
else:
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}}
img_part = types.Part.from_bytes(
data=img_bytes,
mime_type="image/png",
)
img_parts.append(img_part)
contents = img_parts + [{"text": padded_prompt}]
batch_summary = ", ".join(
f"{socket_name}={batch_size}"
for socket_name, batch_size in image_batch_counts
)
print(
f"[INFO] Prepared Gemini image inputs: {batch_summary}, "
f"total={len(img_parts)}"
)
contents = [
types.UserContent(
parts=img_parts + [types.Part.from_text(text=padded_prompt)]
)
]
except Exception as e:
print(f"[ERROR] Error processing input image: {str(e)}")
return IO.NodeOutput(f"Error processing input image: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
tests/test_gemini_nodes.py
+46 -11
View File
@@ -98,26 +98,53 @@ def test_image_inputs_flatten_in_numeric_socket_and_batch_order():
[
torch.full((2, 2, 3), 0.1),
torch.full((2, 2, 3), 0.2),
torch.full((2, 2, 3), 0.3),
torch.full((2, 2, 3), 0.4),
torch.full((2, 2, 3), 0.5),
]
)
third_socket = torch.full((1, 2, 2, 3), 0.3)
second_socket = torch.full((1, 2, 2, 3), 0.6)
images = gemini_nodes.SSL_GeminiTextPrompt._flatten_image_inputs(
frames, batch_counts = gemini_nodes.SSL_GeminiTextPrompt._ordered_image_frames(
{
"image_3": third_socket,
"image_2": second_socket,
"image_1": first_batch,
}
)
assert [round(float(image[0, 0, 0, 0]), 1) for image in images] == [0.1, 0.2, 0.3]
assert [(socket, index) for socket, index, _ in frames] == [
("image_1", 0),
("image_1", 1),
("image_1", 2),
("image_1", 3),
("image_1", 4),
("image_2", 0),
]
assert [round(float(frame[0, 0, 0]), 1) for _, _, frame in frames] == [
0.1,
0.2,
0.3,
0.4,
0.5,
0.6,
]
assert batch_counts == [("image_1", 5), ("image_2", 1)]
def test_image_inputs_reject_non_bhwc_tensor_with_socket_name():
with pytest.raises(ValueError, match=r"Gemini image_2 must be a BHWC image tensor"):
gemini_nodes.SSL_GeminiTextPrompt._ordered_image_frames(
{"image_2": torch.zeros((2, 2, 3))}
)
def test_all_ordered_images_are_sent_to_gemini_2_and_3_before_prompt(monkeypatch):
captured = {}
captured = {"call_count": 0}
class FakeModels:
@staticmethod
def generate_content(**kwargs):
captured["call_count"] += 1
captured.update(kwargs)
part = SimpleNamespace(text="ok", inline_data=None)
content = SimpleNamespace(parts=[part])
@@ -140,26 +167,33 @@ def test_all_ordered_images_are_sent_to_gemini_2_and_3_before_prompt(monkeypatch
}
kwargs = _execute_kwargs(config)
kwargs["image_inputs"] = {
"image_3": torch.full((1, 2, 2, 3), 0.3),
"image_2": torch.full((1, 2, 2, 3), 0.6),
"image_1": torch.stack(
[
torch.full((2, 2, 3), 0.1),
torch.full((2, 2, 3), 0.2),
torch.full((2, 2, 3), 0.3),
torch.full((2, 2, 3), 0.4),
torch.full((2, 2, 3), 0.5),
]
),
}
def assert_ordered_parts():
assert len(captured["contents"]) == 4
assert captured["contents"][-1] == {"text": "test prompt"}
assert captured["call_count"] == 1
assert len(captured["contents"]) == 1
parts = captured["contents"][0].parts
assert len(parts) == 7
assert parts[-1].text == "test prompt"
pixels = [
Image.open(BytesIO(part["inline_data"]["data"])).getpixel((0, 0))[0]
for part in captured["contents"][:-1]
Image.open(BytesIO(part.inline_data.data)).getpixel((0, 0))[0]
for part in parts[:-1]
]
assert pixels == [25, 51, 76]
assert pixels == [25, 51, 76, 102, 127, 153]
for model in ("gemini-2.0-flash", "gemini-3-flash-preview"):
captured.clear()
captured["call_count"] = 0
kwargs["model"] = model
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs)
@@ -174,6 +208,7 @@ def test_all_ordered_images_are_sent_to_gemini_2_and_3_before_prompt(monkeypatch
kwargs["config"] = config | {"use_vertexai_env": True, "api_key": ""}
kwargs["model"] = "gemini-2.5-flash"
captured.clear()
captured["call_count"] = 0
output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs)