From cf939cd6a6dddf0312878bb274e4ddcb70551618 Mon Sep 17 00:00:00 2001 From: Fill Date: Thu, 15 May 2025 12:25:49 +0900 Subject: [PATCH] tweaks to gemini image gen nodes --- nodes/FL_GeminiImageEditor.py | 27 +++++---- nodes/FL_GeminiImageGenADV.py | 111 +++++++++++++++++++--------------- 2 files changed, 79 insertions(+), 59 deletions(-) diff --git a/nodes/FL_GeminiImageEditor.py b/nodes/FL_GeminiImageEditor.py index 4ceb205..28e50c2 100644 --- a/nodes/FL_GeminiImageEditor.py +++ b/nodes/FL_GeminiImageEditor.py @@ -31,7 +31,7 @@ class FL_GeminiImageEditor: "batch_size": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}), }, "optional": { - "seed": ("INT", {"default": 66666666, "min": 0, "max": 2147483647}), + "seed": ("INT", {"default": 66666666, "min": 0, "max": 66666666}), "image1": ("IMAGE",), "image2": ("IMAGE",), "image3": ("IMAGE",), @@ -453,26 +453,31 @@ class FL_GeminiImageEditor: return await asyncio.gather(*tasks) # Run the async batch processing - import asyncio + # Always create a new event loop for this execution context + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + results = None # Initialize results try: - # Get or create event loop based on platform - try: - loop = asyncio.get_event_loop() - except RuntimeError: - # If no event loop exists, create a new one - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - # Run the batch processing results = loop.run_until_complete(run_batch()) except Exception as e: self._log(f"Error in async processing: {str(e)}") traceback.print_exc() - # Create batch of error images error_imgs = [self._create_error_image(f"Async processing error: {str(e)}")] * batch_size batch_tensor = torch.cat(error_imgs, dim=0) return (batch_tensor, f"Async processing error: {str(e)}") + finally: + # Ensure the loop is closed + if loop and not loop.is_closed(): + loop.close() + + # Process results (ensure results is not None if an error occurred before assignment) + if results is None: + self._log("Async processing did not yield results, possibly due to an earlier error before gather.") + error_imgs = [self._create_error_image("Async processing failed to produce results")] * batch_size + batch_tensor = torch.cat(error_imgs, dim=0) + return (batch_tensor, "Async processing failed to produce results") # Process results all_tensors = [] diff --git a/nodes/FL_GeminiImageGenADV.py b/nodes/FL_GeminiImageGenADV.py index 2d2a390..f90f5b4 100644 --- a/nodes/FL_GeminiImageGenADV.py +++ b/nodes/FL_GeminiImageGenADV.py @@ -25,7 +25,7 @@ class FL_GeminiImageGenADV: def INPUT_TYPES(cls): return { "required": { - "inputcount": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), + "inputcount": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}), "api_key": ("STRING", {"default": os.getenv("GEMINI_API_KEY", ""), "multiline": False}), "model": (["models/gemini-2.0-flash-exp", "models/gemini-2.0-flash-preview-image-generation"], {"default": "models/gemini-2.0-flash-preview-image-generation"}), "temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}), @@ -34,7 +34,7 @@ class FL_GeminiImageGenADV: }, "optional": { "image_1": ("IMAGE", {}), # Moved image_1 to optional. Default will be None if not connected. - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), # Restored full seed range + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), # Restored full seed range # Subsequent image_i and prompt_i will be handled by **kwargs based on inputcount } } @@ -117,28 +117,36 @@ Each pair triggers an asynchronous API call. Results are batched. self._log(f"Created error image: '{error_message}'") return img_tensor - def _process_tensor_to_pil(self, tensor_image: Optional[torch.Tensor], image_name: str = "Image") -> Optional[Image.Image]: + def _process_tensor_to_pil_list(self, tensor_image: Optional[torch.Tensor], image_name_prefix: str = "Image") -> Optional[List[Image.Image]]: if tensor_image is None: - self._log(f"{image_name} is None, skipping PIL conversion.") + self._log(f"{image_name_prefix} input is None, skipping PIL conversion.") return None if not isinstance(tensor_image, torch.Tensor): - self._log(f"{image_name} is not a tensor, skipping. Type: {type(tensor_image)}") + self._log(f"{image_name_prefix} is not a tensor (type: {type(tensor_image)}), skipping.") return None - if tensor_image.ndim == 4 and tensor_image.shape[0] > 0: - img_np = tensor_image[0].cpu().numpy() # Use the first image in the batch - img_np = (img_np * 255).astype(np.uint8) - pil_image = Image.fromarray(img_np) - self._log(f"Converted {image_name} (shape: {tensor_image.shape}) to PIL Image (size: {pil_image.size}).") - return pil_image - elif tensor_image.ndim == 3: # Assume (H, W, C) + + pil_images = [] + if tensor_image.ndim == 4: # Batch of images (B, H, W, C) + if tensor_image.shape[0] == 0: + self._log(f"{image_name_prefix} batch is empty (shape: {tensor_image.shape}).") + return None + for i in range(tensor_image.shape[0]): + img_np = tensor_image[i].cpu().numpy() + img_np = (img_np * 255).astype(np.uint8) + pil_image = Image.fromarray(img_np) + self._log(f"Converted {image_name_prefix} batch item {i} (original shape: {tensor_image.shape}) to PIL Image (size: {pil_image.size}).") + pil_images.append(pil_image) + elif tensor_image.ndim == 3: # Single image (H, W, C) img_np = tensor_image.cpu().numpy() img_np = (img_np * 255).astype(np.uint8) pil_image = Image.fromarray(img_np) - self._log(f"Converted {image_name} (shape: {tensor_image.shape}) to PIL Image (size: {pil_image.size}).") - return pil_image + self._log(f"Converted single {image_name_prefix} (original shape: {tensor_image.shape}) to PIL Image (size: {pil_image.size}).") + pil_images.append(pil_image) else: - self._log(f"Cannot convert {image_name} with shape {tensor_image.shape} to PIL Image.") + self._log(f"Cannot convert {image_name_prefix} with ndim {tensor_image.ndim} (shape: {tensor_image.shape}) to PIL Image(s).") return None + + return pil_images if pil_images else None def _call_gemini_api(self, client_instance, model_name_full, contents, gen_config_obj, retry_count=0, max_retries=3, call_id="0"): try: @@ -289,31 +297,24 @@ Each pair triggers an asynchronous API call. Results are batched. return image_tensor, final_response_text - async def _generate_single_image_async(self, api_key, model_name_full, prompt_text, input_pil_image, temperature, max_retries, seed_val, call_id): + async def _generate_single_image_async(self, api_key, model_name_full, prompt_text, input_pil_images: Optional[List[Image.Image]], temperature, max_retries, seed_val, call_id): # Parameter changed to input_pil_images try: - # Use genai.Client for initialization, similar to FL_GeminiImageEditor - # Ensure API key is passed if required by genai.Client constructor - # Some SDK versions might use genai.configure(api_key=...) globally first, - # then genai.Client() without api_key arg. - # Given the error, direct client instantiation is safer. try: client_instance = genai.Client(api_key=api_key) - except TypeError: # Fallback if genai.Client() doesn't take api_key (older versions might not) - genai.configure(api_key=api_key) # Try global configure - client_instance = genai.Client() # Then instantiate - except AttributeError: # If genai.Client itself is not found, this is a deeper SDK issue. + except TypeError: + genai.configure(api_key=api_key) + client_instance = genai.Client() + except AttributeError: self._log(f"[Call {call_id}] CRITICAL: genai.Client not found. Please check google-genai SDK installation and version.") error_msg = f"Call {call_id} Error: genai.Client not found." return self._create_error_image(error_msg), error_msg, call_id - actual_seed = seed_val if seed_val != 0 else random.randint(1, 0xffffffffffffffff) self._log(f"[Call {call_id}] Using seed: {actual_seed} for prompt: '{prompt_text[:50]}...'") - # Use types.GenerateContentConfig as in FL_GeminiImageEditor gen_config_params = { "temperature": temperature, - "response_modalities": ['Text', 'Image'] # From FL_GeminiImageEditor + "response_modalities": ['Text', 'Image'] } if actual_seed != 0: gen_config_params["seed"] = actual_seed @@ -325,9 +326,11 @@ Each pair triggers an asynchronous API call. Results are batched. if current_seed_in_config != actual_seed: self._log(f"[Call {call_id}] Warning: Seed {actual_seed} was specified. GenerateContentConfig has seed: {current_seed_in_config}. Ensure model supports seed via this config.") - contents = [prompt_text] - if input_pil_image: - contents.append(input_pil_image) + contents = [prompt_text] # Start with the prompt + if input_pil_images: # If the list of PIL images is provided and not empty + for pil_img in input_pil_images: + if pil_img: # Ensure the image itself is not None + contents.append(pil_img) loop = asyncio.get_event_loop() response = await loop.run_in_executor( @@ -343,44 +346,56 @@ Each pair triggers an asynchronous API call. Results are batched. error_msg = f"Call {call_id} Error: {str(e)}" return self._create_error_image(error_msg), error_msg, call_id - def generate_images_advanced(self, inputcount, api_key, model, temperature, max_retries, prompt_1, image_1=None, seed=0, **kwargs): # image_1 now has default None + def generate_images_advanced(self, inputcount, api_key, model, temperature, max_retries, prompt_1, image_1=None, seed=0, **kwargs): self.log_messages = [] if not api_key: error_msg = "API key not provided." self._log(error_msg) - # Ensure error image is created with the correct (updated) default size if needed error_img_instance = self._create_error_image(error_msg) + # API key error should return 'inputcount' error images if we can determine it, + # otherwise, a single error image is a reasonable fallback. + # Since each slot is one API call now, inputcount is the number of expected results. return ([error_img_instance] * inputcount, error_msg) loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) tasks = [] - for i in range(1, inputcount + 1): - current_prompt = prompt_1 if i == 1 else kwargs.get(f"prompt_{i}", f"Default prompt for image {i}") - # Handle image_1 being optional, and subsequent image_i from kwargs - current_image_tensor = None - if i == 1: - current_image_tensor = image_1 # This could be None if not provided + for slot_idx in range(1, inputcount + 1): + current_prompt = prompt_1 if slot_idx == 1 else kwargs.get(f"prompt_{slot_idx}", f"Default prompt for image {slot_idx}") + + current_image_tensor_for_slot = None + if slot_idx == 1: + current_image_tensor_for_slot = image_1 else: - current_image_tensor = kwargs.get(f"image_{i}") # This could also be None + current_image_tensor_for_slot = kwargs.get(f"image_{slot_idx}") - current_pil_image = self._process_tensor_to_pil(current_image_tensor, f"InputImage{i}") # This function already handles None + # This will return a list of PIL images if current_image_tensor_for_slot is a batch or single image, + # or None if it's None or invalid. + pil_images_for_this_slot = self._process_tensor_to_pil_list(current_image_tensor_for_slot, f"InputSlot{slot_idx}") - # Increment seed for each task if a non-zero base seed is provided - current_seed = seed + (i-1) if seed != 0 else 0 + # Each slot_idx corresponds to one API call. + # The seed is incremented per slot_idx (per API call). + current_task_seed = seed + (slot_idx - 1) if seed != 0 else 0 + task_call_id = str(slot_idx) # Simplified call_id tasks.append(self._generate_single_image_async( - api_key, model, current_prompt, current_pil_image, - temperature, max_retries, current_seed, str(i) + api_key, model, current_prompt, pil_images_for_this_slot, # Pass the list of PIL images + temperature, max_retries, current_task_seed, task_call_id )) + if not tasks: + self._log("No tasks were created. This might indicate an issue with inputcount or logic.") + return ([self._create_error_image("No tasks generated")], "No tasks generated") + + results_with_id = [] try: results_with_id = loop.run_until_complete(asyncio.gather(*tasks)) finally: - loop.close() - - # Sort results by call_id to maintain order, though gather usually preserves submission order + if loop and not loop.is_closed(): + loop.close() + + # Sort results by call_id (which is now just the string of slot_idx) results_with_id.sort(key=lambda x: int(x[2])) output_images = []