tweaks to gemini image gen nodes

This commit is contained in:
Fill
2025-05-15 12:25:49 +09:00
parent c3e85fffd8
commit cf939cd6a6
2 changed files with 79 additions and 59 deletions
+16 -11
View File
@@ -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 = []
+63 -48
View File
@@ -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 = []