From ecf4cc34a7d4cfeaa74d0fb77a27c77cdc1daeda Mon Sep 17 00:00:00 2001
From: Aryan185 <aryan@euphotic.io>
Date: Mon, 15 Dec 2025 21:26:51 +0530
Subject: [PATCH] minor bug fixes

---
 __init__.py                  | 1 -
 gemini_diarisation_vertex.py | 5 +++--
 imagen.py                    | 6 +++---
 imagen_edit_vertex.py        | 6 +++---
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/__init__.py b/__init__.py
index 5839e1d..f157209 100644
--- a/__init__.py
+++ b/__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
 
diff --git a/gemini_diarisation_vertex.py b/gemini_diarisation_vertex.py
index fb8b358..b8f7eed 100644
--- a/gemini_diarisation_vertex.py
+++ b/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:
diff --git a/imagen.py b/imagen.py
index 95eb46c..a28be72 100644
--- a/imagen.py
+++ b/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"}
\ No newline at end of file
+NODE_CLASS_MAPPINGS = {"GoogleImagenVertex": GoogleImagenVertex}
+NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenVertex": "Google Imagen (Vertex AI)"}
\ No newline at end of file
diff --git a/imagen_edit_vertex.py b/imagen_edit_vertex.py
index 20ed2fd..cbd87aa 100644
--- a/imagen_edit_vertex.py
+++ b/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)"}
\ No newline at end of file
+NODE_CLASS_MAPPINGS = {"GoogleImagenEditVertex": GoogleImagenEditVertex}
+NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenEditVertex": "Imagen Edit (Vertex AI)"}
\ No newline at end of file