From 44e471c6e8212a66f4c839c3e95dfbbc66723578 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 16 Aug 2026 07:19:05 +0200 Subject: [PATCH] fix(gemini): preserve image batch order --- gemini_nodes.py | 63 +++++++++++++++++++++++++------------- tests/test_gemini_nodes.py | 57 +++++++++++++++++++++++++++------- 2 files changed, 88 insertions(+), 32 deletions(-) diff --git a/gemini_nodes.py b/gemini_nodes.py index 6435b69..bda3ab5 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -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) diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index db63697..ce0479b 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -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)