mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-04 09:37:07 +08:00
Added GCP bucket support for image uploads, added image + text -> text.
Added some steps and documentation
This commit is contained in:
parent
670ab2a4b3
commit
88faf138c3
@ -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
|
||||
|
||||
@ -18,6 +18,7 @@ scipy
|
||||
tqdm
|
||||
psutil
|
||||
google-genai
|
||||
google-cloud-storage
|
||||
|
||||
#non essential dependencies:
|
||||
kornia>=0.7.1
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user