From 88faf138c3808b0da67758a51d98c5066e65f30c Mon Sep 17 00:00:00 2001
From: Ganapathy Hari Narayan
 <ganhari123@gmail.com / ganapathy.hari.narayan@gmail.com>
Date: Wed, 14 May 2025 16:56:35 -0700
Subject: [PATCH] Added GCP bucket support for image uploads, added image +
 text -> text.

Added some steps and documentation
---
 comfy_api_nodes/nodes_vertex_gemini.py | 79 ++++++++++++++++++++++++--
 requirements.txt                       |  1 +
 2 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/comfy_api_nodes/nodes_vertex_gemini.py b/comfy_api_nodes/nodes_vertex_gemini.py
index 2349ac445..af80db02e 100644
--- a/comfy_api_nodes/nodes_vertex_gemini.py
+++ b/comfy_api_nodes/nodes_vertex_gemini.py
@@ -1,11 +1,43 @@
+import os
+import uuid
 from inspect import cleandoc
 from google import genai
-from google.genai.types import HttpOptions
+from google.genai.types import HttpOptions, Part
+from google.cloud import storage
 
 from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
 from comfy_api_nodes.apinode_utils import validate_string
 from server import PromptServer
 
+### Documentation ###
+'''
+export GOOGLE_CLOUD_PROJECT="<NAME OF GCP PROJECT>"
+export GOOGLE_CLOUD_LOCATION="<NAME OF REGION e.g. us-central1>"
+export GOOGLE_GENAI_USE_VERTEXAI=True Use the vertex api
+
+How to create a service account:
+1. On GCP console go to IAM & Admin
+2. Select Service accounts
+3. on the top select + create service account
+4. put in the name of service account and select Create and continue
+5. In "Grant this service account access to project" select "Vertex AI User" & "Storage Object Admin"
+6. Once created go to the list of service account and select your service account
+7. Click on keys and click Add key (Json)
+8. download the key and keep it in secure place on your system
+9. use this key to auth by exporting var below:
+export GOOGLE_APPLICATION_CREDENTIALS="/Path/to/Key.json"
+'''
+
+BUCKET_NAME = "comfyui-interview-temp"
+
+def get_model_list():
+    return list([
+        "gemini-2.0-flash-001",
+        "gemini-2.0-flash-lite",
+        "gemini-2.5-pro-preview-05-06",
+        "gemini-2.5-flash-preview-04-17"
+    ])
+
 
 class VertexGeminiAPI(ComfyNodeABC):
     """
@@ -28,36 +60,71 @@ class VertexGeminiAPI(ComfyNodeABC):
                     },
                 ),
                 "model": (
-                    IO.STRING,
+                    get_model_list(),
                     {
                         "default": "gemini-2.0-flash-001",
-                        "tooltip": "The gemini model to use"
+                        "tooltip": "Select the model you would like to use"
                     }
                 )
             },
+            "optional": {
+                "image_path": (
+                    IO.STRING,
+                    {
+                        "default": None,
+                        "tooltip": "Optional reference path to an image for inference.",
+                    }
+                )
+            },
+            "hidden": {
+                "unique_id": "UNIQUE_ID",
+            },
         }
 
     RETURN_TYPES = (IO.STRING,)
     FUNCTION = "api_call"
     CATEGORY = "api node/text/gemini/vertex"
     DESCRIPTION = cleandoc(__doc__ or "")
+    OUTPUT_NODE = True
     API_NODE = True
 
     def api_call(
         self,
         prompt,
         model="gemini-2.0-flash-001",
+        image_path=None,
+        unique_id=None,
         **kwargs
     ):
         validate_string(prompt, strip_whitespace=False)
         client = genai.Client(http_options=HttpOptions(api_version="v1"))
+        contents = [prompt]
+        if image_path:
+            storage_client = storage.Client()
+
+            # Define bucket and file info
+            bucket_name = "comfyui-interview-temp"
+            source_file = image_path              # local path
+            file_name = os.path.basename(source_file)
+            destination_blob = f"{uuid.uuid4()}/{file_name}"         # name in bucket
+            # Upload
+            bucket = storage_client.bucket(bucket_name)
+            blob = bucket.blob(destination_blob)
+            blob.upload_from_filename(source_file)
+
+            print(f"Uploaded {source_file} to gs://{bucket_name}/{destination_blob}")
+            image_part = Part.from_uri(
+                file_uri=f"gs://{bucket_name}/{destination_blob}",
+                mime_type="image/jpeg"
+            )
+            contents.append(image_part)
         response = client.models.generate_content(
             model=model,
-            contents=prompt,
+            contents=contents,
         )
         print(response.text)
-        PromptServer.instance.send_progress_text(response.text)
-        return response.text
+        PromptServer.instance.send_progress_text(response.text, node_id=unique_id)
+        return (response.text,)
 
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
diff --git a/requirements.txt b/requirements.txt
index ef3dbc9f6..1e2bfb7aa 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -18,6 +18,7 @@ scipy
 tqdm
 psutil
 google-genai
+google-cloud-storage
 
 #non essential dependencies:
 kornia>=0.7.1