minor bug fixes

__init__.py

@@ -4,7 +4,6 @@ from .nano_banana_vertex import NODE_CLASS_MAPPINGS as NANO_BANANA_MAPPINGS, NOD
 from .gemini_tts_vertex import NODE_CLASS_MAPPINGS as GEMINI_TTS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_TTS_DISPLAY
 from .veo_vertex import NODE_CLASS_MAPPINGS as VEO_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as VEO_DISPLAY
 from .gemini_diarisation_vertex import NODE_CLASS_MAPPINGS as GEMINI_DIARISATION_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_DIARISATION_DISPLAY
-
 from .imagen_edit_vertex import NODE_CLASS_MAPPINGS as IMAGEN_EDIT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGEN_EDIT_DISPLAY
 from .imagen import NODE_CLASS_MAPPINGS as IMAGEN_IMAGE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGEN_IMAGE_DISPLAY
 

gemini_diarisation_vertex.py

@@ -41,6 +41,7 @@ class GeminiDiarisationNode:
             "optional": {
                 "thinking": ("BOOLEAN", {"default": False}),
                 "thinking_budget": ("INT", {"default": 0, "min": -1, "max": 24576, "step": 1}),
+                "audio_timestamp": ("BOOLEAN", {"default": False})
             }
         }
     
@@ -229,7 +230,7 @@ class GeminiDiarisationNode:
     
     def diarise(self, audio, project_id, location, service_account, model, 
                 seed, temperature, num_speakers,
-                thinking=False, thinking_budget=0):
+                thinking=False, thinking_budget=0, audio_timestamp=False):
         
         audio_data, sr = self.extract_audio_data(audio)
         audio_data = self.normalize_audio(audio_data)
@@ -244,7 +245,7 @@ class GeminiDiarisationNode:
         
         diarization_config = types.GenerateContentConfig(
             temperature=temperature,
-            audio_timestamp=True,
+            audio_timestamp=audio_timestamp,
         )
         
         if seed >= 0:

imagen.py

@@ -5,7 +5,7 @@ from PIL import Image
 from google import genai
 from google.genai import types
 
-class GoogleImagenNode:
+class GoogleImagenVertex:
     
     @classmethod
     def INPUT_TYPES(cls):
@@ -107,5 +107,5 @@ class GoogleImagenNode:
     def IS_CHANGED(cls, **kwargs):
         return f"{kwargs.get('prompt', '')}-{kwargs.get('model', '')}-{kwargs.get('number_of_images', 1)}-{kwargs.get('aspect_ratio', '1:1')}-{kwargs.get('image_size', '1K')}"
 
-NODE_CLASS_MAPPINGS = {"GoogleImagenNode": GoogleImagenNode}
-NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenNode": "Google Imagen Generator"}
+NODE_CLASS_MAPPINGS = {"GoogleImagenVertex": GoogleImagenVertex}
+NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenVertex": "Google Imagen (Vertex AI)"}

imagen_edit_vertex.py

@@ -10,7 +10,7 @@ from google import genai
 from google.genai import types
 from google.genai.types import RawReferenceImage, MaskReferenceImage
 
-class GoogleImagenEditNode:
+class GoogleImagenEditVertex:
     
     @classmethod
     def INPUT_TYPES(cls):
@@ -168,5 +168,5 @@ class GoogleImagenEditNode:
     def IS_CHANGED(cls, **kwargs):
         return f"{kwargs.get('prompt', '')}-{kwargs.get('negative_prompt', '')}-{kwargs.get('edit_mode', '')}-{kwargs.get('number_of_images', 1)}-{kwargs.get('seed', 12345)}-{kwargs.get('base_steps', 50)}"
 
-NODE_CLASS_MAPPINGS = {"GoogleImagenEditNode": GoogleImagenEditNode}
-NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenEditNode": "Imagen Edit (Vertex AI)"}
+NODE_CLASS_MAPPINGS = {"GoogleImagenEditVertex": GoogleImagenEditVertex}
+NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenEditVertex": "Imagen Edit (Vertex AI)"}