From 2b4725a02e3c6d0a2162e26f434dc4a31ab6c975 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 21 Oct 2026 19:16:06 +0200 Subject: [PATCH] Fix list execution, interaction chaining, and interrupt handling --- gemini_nodes.py | 252 +++++++++++++++++++++++-------- interactions_nodes.py | 231 ++++++++++++++++++++-------- tests/test_gemini_nodes.py | 44 +++++- tests/test_interactions_nodes.py | 62 +++++++- 4 files changed, 455 insertions(+), 134 deletions(-) diff --git a/gemini_nodes.py b/gemini_nodes.py index 729e352..c7dd465 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -8,6 +8,7 @@ import cv2 from PIL import Image from io import BytesIO import folder_paths # type: ignore[reportMissingImports] +import comfy.model_management # type: ignore[reportMissingImports] from comfy_api.latest import ComfyExtension, UI, IO, Types, VideoFromComponents # type: ignore[reportMissingImports] from google import genai from google.genai import errors, types @@ -21,7 +22,6 @@ import random import hashlib import warnings from typing import Any, Tuple -from concurrent.futures import ThreadPoolExecutor class GetKeyAPI(IO.ComfyNode): @classmethod @@ -436,6 +436,46 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): ) return frames, batch_counts + @classmethod + def _prepare_image_inputs(cls, image_inputs: IO.Autogrow.Type | None) -> dict[str, Any] | None: + if not image_inputs: + return None + image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) + if not image_frames: + return { + "image_hashes": (), + "png_bytes_list": [], + "png_hashes": [], + "batch_counts": [], + } + tensor_hashes = [] + png_bytes_list = [] + png_hashes = [] + for _, _, image_frame in image_frames: + np_frame = np.ascontiguousarray(image_frame.cpu().numpy()) + tensor_hashes.append(hashlib.sha256(np_frame.tobytes()).hexdigest()) + img_array = (np_frame * 255).astype(np.uint8) + pil_img = Image.fromarray(img_array) + img_byte_arr = BytesIO() + pil_img.save(img_byte_arr, format='PNG') + img_bytes = img_byte_arr.getvalue() + png_bytes_list.append(img_bytes) + png_hashes.append(hashlib.sha256(img_bytes).hexdigest()) + 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(png_bytes_list)}" + ) + return { + "image_hashes": tuple(tensor_hashes), + "png_bytes_list": png_bytes_list, + "png_hashes": png_hashes, + "batch_counts": image_batch_counts, + } + @classmethod def _serialize_video(cls, video: IO.Video.Type | None, fps: int | None = None, pad_at_start: bool = False, duration_aware_padding: bool = False): @@ -591,7 +631,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, timeout=30, include_thoughts=False, thinking_level=None, media_resolution=None, retry_pattern="", max_retries=3, use_cache=False, - cache_ttl_minutes=60, cache_seed=0, image_size="None", schema_identity=None): + cache_ttl_minutes=60, cache_seed=0, image_size="None", schema_identity=None, + precomputed_image_hashes=None): # 1. Hashing Images def get_tensor_hash(tensor): @@ -603,8 +644,11 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): print(f"[WARNING] Cache hashing failed for image: {e}") return "Error" - image_frames, _ = cls._ordered_image_frames(image_inputs) - image_hashes = tuple(get_tensor_hash(frame) for _, _, frame in image_frames) + if precomputed_image_hashes is not None: + image_hashes = precomputed_image_hashes + else: + 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 @@ -930,6 +974,24 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): def _reject_json_constant(value): raise ValueError(f"Structured response contains invalid JSON constant: {value}.") + @classmethod + def _check_interrupt(cls): + comfy.model_management.throw_exception_if_processing_interrupted() + + @classmethod + def _wait_for_result(cls, result_queue: "queue.Queue[Tuple[str, Any]]", timeout: float) -> Tuple[str, Any]: + deadline = time.monotonic() + float(timeout) + while True: + cls._check_interrupt() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise queue.Empty() + try: + return result_queue.get(timeout=min(0.1, remaining)) + except queue.Empty: + if time.monotonic() >= deadline: + raise + @classmethod 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, @@ -939,8 +1001,11 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): image_size="None", video: GeminiVideoConfig.Type | None = None, image_inputs: IO.Autogrow.Type | None = None, - response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput: + response_schema: ResponseSchema.Type | None = None, + prepared_images: dict[str, Any] | None = None, + precomputed_video: Tuple[bytes | None, str | None, str | None, int | None] | None = None) -> IO.NodeOutput: + cls._check_interrupt() print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}") use_cache = bool(config.get("use_cache", False)) cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60)) @@ -954,22 +1019,26 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): raise ValueError("Response format cannot be combined with image generation. Disable include_images.") schema_identity = json.dumps(response_schema, ensure_ascii=False, separators=(",", ":"), allow_nan=False) schema_snapshot = json.loads(schema_identity) - video_input = video.get("video") if video is not None else None - video_fps = int(video.get("fps", 1)) if video is not None else None - video_pad_at_start = bool(video.get("pad_at_start", False)) if video is not None else False - video_duration_aware_padding = bool(video.get("duration_aware_padding", False)) if video is not None else False - try: - video_bytes, video_mime_type = cls._serialize_video( - video_input, - video_fps, - video_pad_at_start, - video_duration_aware_padding, - ) - except Exception as e: - print(f"[ERROR] Error processing input video: {e}") - output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 - return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "") - video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None + if precomputed_video is not None: + video_bytes, video_mime_type, video_hash, video_fps = precomputed_video + else: + video_input = video.get("video") if video is not None else None + video_fps = int(video.get("fps", 1)) if video is not None else None + video_pad_at_start = bool(video.get("pad_at_start", False)) if video is not None else False + video_duration_aware_padding = bool(video.get("duration_aware_padding", False)) if video is not None else False + try: + video_bytes, video_mime_type = cls._serialize_video( + video_input, + video_fps, + video_pad_at_start, + video_duration_aware_padding, + ) + except Exception as e: + print(f"[ERROR] Error processing input video: {e}") + output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 + return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "") + video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None + precomputed_image_hashes = prepared_images["image_hashes"] if prepared_images is not None else None fingerprint, cached = cls._compute_fingerprint_and_check_cache( config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed, @@ -980,6 +1049,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): retry_pattern, max_retries, use_cache, cache_ttl_minutes, cache_seed, image_size=image_size, schema_identity=schema_identity, + precomputed_image_hashes=precomputed_image_hashes, ) if cached is not None: @@ -1143,7 +1213,6 @@ 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 + video + prompt) - image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) context_image_hashes = [] media_parts = [] context_inline_bytes = ( @@ -1152,42 +1221,60 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): + len(padded_system_instruction.encode("utf-8")) ) - if image_frames: - try: - img_parts = [] - 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) - img_byte_arr = BytesIO() - pil_img.save(img_byte_arr, format='PNG') - img_bytes = img_byte_arr.getvalue() - context_inline_bytes += len(img_bytes) - context_image_hashes.append(hashlib.sha256(img_bytes).hexdigest()) - if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": - img_part = types.Part.from_bytes( - data=img_bytes, - mime_type="image/png", - media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}", - ) - else: - img_part = types.Part.from_bytes( - data=img_bytes, - mime_type="image/png", - ) - img_parts.append(img_part) - 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)}" - ) - media_parts.extend(img_parts) - 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, "", "") + if prepared_images is not None: + for img_bytes, png_hash in zip(prepared_images["png_bytes_list"], prepared_images["png_hashes"]): + context_inline_bytes += len(img_bytes) + context_image_hashes.append(png_hash) + if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": + img_part = types.Part.from_bytes( + data=img_bytes, + mime_type="image/png", + media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}", + ) + else: + img_part = types.Part.from_bytes( + data=img_bytes, + mime_type="image/png", + ) + media_parts.append(img_part) + else: + image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) + if image_frames: + try: + img_parts = [] + 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) + img_byte_arr = BytesIO() + pil_img.save(img_byte_arr, format='PNG') + img_bytes = img_byte_arr.getvalue() + context_inline_bytes += len(img_bytes) + context_image_hashes.append(hashlib.sha256(img_bytes).hexdigest()) + if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": + img_part = types.Part.from_bytes( + data=img_bytes, + mime_type="image/png", + media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}", + ) + else: + img_part = types.Part.from_bytes( + data=img_bytes, + mime_type="image/png", + ) + img_parts.append(img_part) + 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)}" + ) + media_parts.extend(img_parts) + 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, "", "") if video_bytes is not None: if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": @@ -1401,7 +1488,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): thoughts_output = "" try: - status, result = result_queue.get(timeout=timeout) + status, result = cls._wait_for_result(result_queue, timeout) elapsed_time = time.time() - start_time if status == "success": text_output, image_tensor, thoughts_output = result @@ -1414,6 +1501,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): try: compiled_pattern = re.compile(retry_pattern, re.IGNORECASE) if compiled_pattern.search(text_output): + cls._check_interrupt() retry_attempt += 1 print(f"[INFO] Retry pattern matched in response. Retry attempt {retry_attempt}/{max_retries}") @@ -1448,6 +1536,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): try: compiled_pattern = re.compile(retry_pattern, re.IGNORECASE) if compiled_pattern.search(text_output): + cls._check_interrupt() retry_attempt += 1 print(f"[INFO] Retry pattern matched in error. Retry attempt {retry_attempt}/{max_retries}") @@ -1583,7 +1672,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): img_size_val = str(_unwrap_scalar(image_size, "None")) resp_schema_val = _unwrap_scalar(response_schema, None) - def make_item_kwargs(idx): + 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 @@ -1628,6 +1717,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): "video": v, "image_inputs": image_inputs, "response_schema": resp_schema_val, + "prepared_images": prepared_imgs, + "precomputed_video": precomp_vid, } if batch_count == 1: @@ -1640,13 +1731,40 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): [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] + print(f"[INFO] Preparing sequential list execution for {batch_count} items (model: {model_val})") + prepared_images = None + if image_inputs: + try: + prepared_images = cls._prepare_image_inputs(image_inputs) + except Exception: + prepared_images = None + + precomputed_video = None + if len(video_list) == 1 and video_list[0] is not None: + try: + v0 = video_list[0] + v_input = v0.get("video") + v_fps = int(v0.get("fps", 1)) + v_pad = bool(v0.get("pad_at_start", False)) + v_dur = bool(v0.get("duration_aware_padding", False)) + v_bytes, v_mime = cls._serialize_video(v_input, v_fps, v_pad, v_dur) + v_hash = hashlib.sha256(v_bytes).hexdigest() if v_bytes is not None else None + precomputed_video = (v_bytes, v_mime, v_hash, v_fps) + except Exception: + precomputed_video = None + + batch_results = [] + for idx in range(batch_count): + cls._check_interrupt() + batch_results.append( + cls._execute_single( + **make_item_kwargs( + idx, + prepared_imgs=prepared_images, + precomp_vid=precomputed_video, + ) + ) + ) return IO.NodeOutput( [r[0] for r in batch_results], diff --git a/interactions_nodes.py b/interactions_nodes.py index 65af001..4887262 100644 --- a/interactions_nodes.py +++ b/interactions_nodes.py @@ -9,6 +9,7 @@ from PIL import Image from io import BytesIO import base64 import folder_paths # type: ignore[reportMissingImports] +import comfy.model_management # type: ignore[reportMissingImports] from comfy_api.latest import ComfyExtension, UI, IO, Types, VideoFromComponents # type: ignore[reportMissingImports] from google import genai from google.genai import errors, types @@ -23,7 +24,6 @@ import hashlib import warnings from typing import Any, Tuple from datetime import datetime -from concurrent.futures import ThreadPoolExecutor from .gemini_nodes import ( GetKeyAPI, @@ -461,6 +461,41 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): ) return frames, batch_counts + @classmethod + def _prepare_image_inputs(cls, image_inputs: IO.Autogrow.Type | None) -> dict[str, Any] | None: + if not image_inputs: + return None + image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) + if not image_frames: + return { + "image_hashes": (), + "png_bytes_list": [], + "batch_counts": [], + } + tensor_hashes = [] + png_bytes_list = [] + for _, _, image_frame in image_frames: + np_frame = np.ascontiguousarray(image_frame.cpu().numpy()) + tensor_hashes.append(hashlib.sha256(np_frame.tobytes()).hexdigest()) + img_array = (np_frame * 255).astype(np.uint8) + pil_img = Image.fromarray(img_array) + img_byte_arr = BytesIO() + pil_img.save(img_byte_arr, format='PNG') + png_bytes_list.append(img_byte_arr.getvalue()) + 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(png_bytes_list)}" + ) + return { + "image_hashes": tuple(tensor_hashes), + "png_bytes_list": png_bytes_list, + "batch_counts": image_batch_counts, + } + @classmethod def _serialize_video(cls, video: IO.Video.Type | None, fps: int | None = None, pad_at_start: bool = False, duration_aware_padding: bool = False): @@ -592,7 +627,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): retry_pattern="", max_retries=3, use_cache=False, cache_ttl_minutes=60, cache_seed=0, image_size="None", store=True, previous_interaction_id="", safety_settings=None, - schema_identity=None): + schema_identity=None, precomputed_image_hashes=None): def get_tensor_hash(tensor): if tensor is None: @@ -603,8 +638,11 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): print(f"[WARNING] Cache hashing failed for image: {e}") return "Error" - image_frames, _ = cls._ordered_image_frames(image_inputs) - image_hashes = tuple(get_tensor_hash(frame) for _, _, frame in image_frames) + if precomputed_image_hashes is not None: + image_hashes = precomputed_image_hashes + else: + image_frames, _ = cls._ordered_image_frames(image_inputs) + image_hashes = tuple(get_tensor_hash(frame) for _, _, frame in image_frames) eff_include_images = False eff_aspect_ratio = "None" @@ -744,6 +782,24 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): def _reject_json_constant(value): raise ValueError(f"Structured response contains invalid JSON constant: {value}.") + @classmethod + def _check_interrupt(cls): + comfy.model_management.throw_exception_if_processing_interrupted() + + @classmethod + def _wait_for_result(cls, result_queue: "queue.Queue[Tuple[str, Any]]", timeout: float) -> Tuple[str, Any]: + deadline = time.monotonic() + float(timeout) + while True: + cls._check_interrupt() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise queue.Empty() + try: + return result_queue.get(timeout=min(0.1, remaining)) + except queue.Empty: + if time.monotonic() >= deadline: + raise + @classmethod 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, @@ -753,8 +809,11 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): 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: + response_schema: ResponseSchema.Type | None = None, + prepared_images: dict[str, Any] | None = None, + precomputed_video: Tuple[bytes | None, str | None, str | None, int | None] | None = None) -> IO.NodeOutput: + cls._check_interrupt() print(f"[INFO] SSL_GeminiInteractionsTextPrompt execute called, model: {model}") # If model is an image generation model, reroute to SSL_GeminiTextPrompt (GenerateContent) @@ -813,24 +872,28 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): schema_identity = json.dumps(response_schema, ensure_ascii=False, separators=(",", ":"), allow_nan=False) schema_snapshot = json.loads(schema_identity) - video_input = video.get("video") if video is not None else None - video_fps = int(video.get("fps", 1)) if video is not None else None - video_pad_at_start = bool(video.get("pad_at_start", False)) if video is not None else False - video_duration_aware_padding = bool(video.get("duration_aware_padding", False)) if video is not None else False - try: - video_bytes, video_mime_type = cls._serialize_video( - video_input, - video_fps, - video_pad_at_start, - video_duration_aware_padding, - ) - except Exception as e: - print(f"[ERROR] Error processing input video: {e}") - output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 - return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "", "") + if precomputed_video is not None: + video_bytes, video_mime_type, video_hash, video_fps = precomputed_video + else: + video_input = video.get("video") if video is not None else None + video_fps = int(video.get("fps", 1)) if video is not None else None + video_pad_at_start = bool(video.get("pad_at_start", False)) if video is not None else False + video_duration_aware_padding = bool(video.get("duration_aware_padding", False)) if video is not None else False + try: + video_bytes, video_mime_type = cls._serialize_video( + video_input, + video_fps, + video_pad_at_start, + video_duration_aware_padding, + ) + except Exception as e: + print(f"[ERROR] Error processing input video: {e}") + output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 + return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "", "") - video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None + video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None resolved_safety_settings = cls._resolve_safety_settings(config) + precomputed_image_hashes = prepared_images["image_hashes"] if prepared_images is not None else None fingerprint, cached = cls._compute_fingerprint_and_check_cache( config, prompt, system_instruction, model, max_output_tokens, include_images, aspect_ratio, bypass_mode, use_seed, seed, @@ -844,6 +907,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): previous_interaction_id=previous_interaction_id, safety_settings=resolved_safety_settings, schema_identity=schema_identity, + precomputed_image_hashes=precomputed_image_hashes, ) if cached is not None: @@ -1014,37 +1078,48 @@ class SSL_GeminiInteractionsTextPrompt(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 multimodal input items - image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) media_items = [] - if image_frames: - try: - 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) - img_byte_arr = BytesIO() - pil_img.save(img_byte_arr, format='PNG') - img_byte_arr.seek(0) - img_item = { - "type": "image", - "data": img_byte_arr, - "mime_type": "image/png", - } - if media_resolution and media_resolution != "unspecified": - img_item["resolution"] = media_resolution.lower() - media_items.append(img_item) - 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(media_items)}" - ) - 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, "", "", "") + if prepared_images is not None: + for img_bytes in prepared_images["png_bytes_list"]: + img_item = { + "type": "image", + "data": BytesIO(img_bytes), + "mime_type": "image/png", + } + if media_resolution and media_resolution != "unspecified": + img_item["resolution"] = media_resolution.lower() + media_items.append(img_item) + else: + image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) + if image_frames: + try: + 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) + img_byte_arr = BytesIO() + pil_img.save(img_byte_arr, format='PNG') + img_byte_arr.seek(0) + img_item = { + "type": "image", + "data": img_byte_arr, + "mime_type": "image/png", + } + if media_resolution and media_resolution != "unspecified": + img_item["resolution"] = media_resolution.lower() + media_items.append(img_item) + 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(media_items)}" + ) + 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, "", "", "") if video_bytes is not None: video_stream = BytesIO(video_bytes) @@ -1276,7 +1351,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): interaction_id_output = "" try: - status, result = result_queue.get(timeout=timeout) + status, result = cls._wait_for_result(result_queue, timeout) if status == "success": text_output, image_tensor, thoughts_output, interaction_id_output = result structured_output = text_output if schema_snapshot is not None else "" @@ -1286,6 +1361,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): try: compiled_pattern = re.compile(retry_pattern, re.IGNORECASE) if compiled_pattern.search(text_output): + cls._check_interrupt() retry_attempt += 1 print(f"[INFO] Retry pattern matched in response. Retry attempt {retry_attempt}/{max_retries}") actual_seed = cls._handle_seed(True, 0) @@ -1308,6 +1384,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): try: compiled_pattern = re.compile(retry_pattern, re.IGNORECASE) if compiled_pattern.search(text_output): + cls._check_interrupt() retry_attempt += 1 print(f"[INFO] Retry pattern matched in error. Retry attempt {retry_attempt}/{max_retries}") actual_seed = cls._handle_seed(True, 0) @@ -1438,7 +1515,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): store_val = bool(_unwrap_scalar(store, True)) resp_schema_val = _unwrap_scalar(response_schema, None) - def make_item_kwargs(idx): + def make_item_kwargs(idx, chained_prev_id="", 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 @@ -1453,7 +1530,9 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): 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, "") + prev_id = str(_get_aligned(previous_interaction_id, idx, "") or "").strip() + if not prev_id and store_val and chained_prev_id: + prev_id = chained_prev_id return { "config": config_val, @@ -1486,6 +1565,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): "video": v, "image_inputs": image_inputs, "response_schema": resp_schema_val, + "prepared_images": prepared_imgs, + "precomputed_video": precomp_vid, } if batch_count == 1: @@ -1499,13 +1580,43 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): [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] + print(f"[INFO] Preparing sequential list execution for {batch_count} items (model: {model_val})") + prepared_images = None + if image_inputs and model_val not in cls.IMAGE_MODELS: + try: + prepared_images = cls._prepare_image_inputs(image_inputs) + except Exception: + prepared_images = None + + precomputed_video = None + if len(video_list) == 1 and video_list[0] is not None and model_val not in cls.IMAGE_MODELS: + try: + v0 = video_list[0] + v_input = v0.get("video") + v_fps = int(v0.get("fps", 1)) + v_pad = bool(v0.get("pad_at_start", False)) + v_dur = bool(v0.get("duration_aware_padding", False)) + v_bytes, v_mime = cls._serialize_video(v_input, v_fps, v_pad, v_dur) + v_hash = hashlib.sha256(v_bytes).hexdigest() if v_bytes is not None else None + precomputed_video = (v_bytes, v_mime, v_hash, v_fps) + except Exception: + precomputed_video = None + + batch_results = [] + chained_interaction_id = "" + for idx in range(batch_count): + cls._check_interrupt() + item_res = cls._execute_single( + **make_item_kwargs( + idx, + chained_prev_id=chained_interaction_id, + prepared_imgs=prepared_images, + precomp_vid=precomputed_video, + ) + ) + batch_results.append(item_res) + if store_val and item_res[5]: + chained_interaction_id = str(item_res[5]).strip() return IO.NodeOutput( [r[0] for r in batch_results], diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index 77e256d..0017eb9 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -1697,7 +1697,7 @@ def test_gemini_standard_auto_safety_resolution(structured_client): assert all(s.method is None for s in sent_safety) -def test_gemini_batch_prompts(structured_client): +def test_gemini_batch_prompts(structured_client, capsys): structured_client.responses = [ [("final", "Response 1")], [("final", "Response 2")], @@ -1705,6 +1705,7 @@ def test_gemini_batch_prompts(structured_client): ] args = _execute_kwargs(_config()) args["prompt"] = ["Prompt one", "Prompt two", "Prompt three"] + args["image_inputs"] = {"image_1": torch.full((1, 2, 2, 3), 0.5)} output = gemini_nodes.SSL_GeminiTextPrompt.execute(**args) assert isinstance(output[0], list) @@ -1712,6 +1713,47 @@ def test_gemini_batch_prompts(structured_client): 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 + assert len(structured_client.clients) == 1 + captured_out = capsys.readouterr().out + assert captured_out.count("Prepared Gemini image inputs: image_1=1, total=1") == 1 + + +def test_gemini_interrupt_aborts_in_flight_call(monkeypatch): + import time + import comfy.model_management + + started = threading.Event() + + class SlowModels: + def generate_content(self, **kwargs): + started.set() + time.sleep(5.0) + return SimpleNamespace(candidates=[]) + + class SlowClient: + def __init__(self, **kwargs): + self.models = SlowModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", SlowClient) + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + comfy.model_management.interrupt_current_processing(False) + + def trigger_interrupt(): + started.wait(timeout=2.0) + time.sleep(0.05) + comfy.model_management.interrupt_current_processing(True) + + interrupter = threading.Thread(target=trigger_interrupt, daemon=True) + interrupter.start() + + args = _execute_kwargs(_config()) + args["timeout"] = 120 + t0 = time.monotonic() + with pytest.raises(comfy.model_management.InterruptProcessingException): + gemini_nodes.SSL_GeminiTextPrompt.execute(**args) + elapsed = time.monotonic() - t0 + assert elapsed < 2.0 + comfy.model_management.interrupt_current_processing(False) def test_gemini_batch_seeds(structured_client): diff --git a/tests/test_interactions_nodes.py b/tests/test_interactions_nodes.py index 3e78a95..c1849b1 100644 --- a/tests/test_interactions_nodes.py +++ b/tests/test_interactions_nodes.py @@ -524,7 +524,7 @@ def test_caching_with_cache_seed_and_usage(mock_interactions_client): assert len(mock_interactions_client.calls) == 1 -def test_interactions_store_node(): +def test_interactions_store_node(mock_interactions_client): schema = interactions_nodes.SSL_GeminiInteractionsStore.define_schema() assert schema.node_id == "SSL_GeminiInteractionsStore" assert schema.category == "API/Gemini/Interactions" @@ -685,19 +685,69 @@ 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): +def test_interactions_batch_prompts(mock_interactions_client, capsys): mock_interactions_client.responses = [ {"text": "Answer 1", "id": "int_b1"}, {"text": "Answer 2", "id": "int_b2"}, + {"text": "Answer 3", "id": "int_b3"}, ] - args = _execute_kwargs(prompt=["Prompt A", "Prompt B"]) + args = _execute_kwargs( + prompt=["Prompt A", "Prompt B", "Prompt C"], + store=True, + previous_interaction_id="", + image_inputs={"image_1": torch.full((1, 2, 2, 3), 0.5)}, + ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert isinstance(output[0], list) - assert output[0] == ["Answer 1", "Answer 2"] + assert output[0] == ["Answer 1", "Answer 2", "Answer 3"] assert isinstance(output[5], list) - assert output[5] == ["int_b1", "int_b2"] - assert len(mock_interactions_client.calls) == 2 + assert output[5] == ["int_b1", "int_b2", "int_b3"] + assert len(mock_interactions_client.calls) == 3 + assert len(mock_interactions_client.clients) == 1 + assert "previous_interaction_id" not in mock_interactions_client.calls[0] + assert mock_interactions_client.calls[1]["previous_interaction_id"] == "int_b1" + assert mock_interactions_client.calls[2]["previous_interaction_id"] == "int_b2" + captured_out = capsys.readouterr().out + assert captured_out.count("Prepared Gemini image inputs: image_1=1, total=1") == 1 + + +def test_interactions_interrupt_aborts_in_flight_call(mock_interactions_client, monkeypatch): + import time + import comfy.model_management + + started = threading.Event() + + class SlowInteractions: + def create(self, **kwargs): + started.set() + time.sleep(5.0) + return SimpleNamespace(id="slow", output_text="slow") + + class SlowClient: + def __init__(self, **kwargs): + self.interactions = SlowInteractions() + + monkeypatch.setattr(interactions_nodes.genai, "Client", SlowClient) + interactions_nodes.SSL_GeminiInteractionsTextPrompt._client_cache.clear() + comfy.model_management.interrupt_current_processing(False) + + def trigger_interrupt(): + started.wait(timeout=2.0) + time.sleep(0.05) + comfy.model_management.interrupt_current_processing(True) + + interrupter = threading.Thread(target=trigger_interrupt, daemon=True) + interrupter.start() + + t0 = time.monotonic() + with pytest.raises(comfy.model_management.InterruptProcessingException): + interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute( + **_execute_kwargs(timeout=120) + ) + elapsed = time.monotonic() - t0 + assert elapsed < 2.0 + comfy.model_management.interrupt_current_processing(False) def test_interactions_batch_seeds(mock_interactions_client):