主分支上传
This commit is contained in:
tatookan
@@ -0,0 +1,106 @@
|
||||
# ComfyUI Gemini Flash 节点
|
||||
|
||||
这是一个用于ComfyUI的Google Gemini API集成节点，支持文本生成和图像生成功能。通过此节点，您可以在ComfyUI工作流中直接使用Google的Gemini 2.0系列模型。
|
||||
|
||||
## 功能特点
|
||||
|
||||
- 支持Gemini 2.0系列模型（gemini-2.0-flash, gemini-2.0-flash-exp, gemini-2.0-pro）
|
||||
- 支持文本到文本生成
|
||||
- 支持图像到文本生成（图像理解）
|
||||
- 支持文本到图像生成（仅在其他.py文件中实现）
|
||||
- 内置代理支持，方便中国用户使用
|
||||
- 自动依赖检查和安装
|
||||
- 完善的错误处理和日志记录
|
||||
|
||||
## 安装方法
|
||||
|
||||
1. 确保您已经安装了ComfyUI
|
||||
2. 将此仓库克隆或下载到ComfyUI的`custom_nodes`目录中：
|
||||
```
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/your-username/comfyui_gemini_flash.git
|
||||
```
|
||||
3. 安装依赖：
|
||||
```
|
||||
pip install -r comfyui_gemini_flash/requirements.txt
|
||||
```
|
||||
4. 重启ComfyUI
|
||||
|
||||
## 使用方法
|
||||
|
||||
### 1. 配置API密钥
|
||||
|
||||
首先，您需要获取Google Gemini API密钥：
|
||||
1. 访问[Google AI Studio](https://makersuite.google.com/app/apikey)
|
||||
2. 创建一个API密钥
|
||||
3. 在ComfyUI中使用`💠SSL/API/Gemini/API Key Config`节点输入您的API密钥
|
||||
|
||||
### 2. 文本生成
|
||||
|
||||
使用`💠SSL/API/Gemini/Text Prompt`节点进行文本生成：
|
||||
|
||||
- 连接API密钥配置节点到`config`输入
|
||||
- 在`prompt`中输入您的提示文本
|
||||
- 调整生成参数（温度、top_p、top_k等）
|
||||
- 如果需要，可以连接图像到`input_image`输入，实现图像理解功能
|
||||
|
||||
### 3. 代理设置
|
||||
|
||||
如果您在中国或其他需要代理的地区：
|
||||
|
||||
1. 将`use_proxy`设置为`True`
|
||||
2. 设置`proxy_host`（默认为127.0.0.1）
|
||||
3. 设置`proxy_port`（默认为7890）
|
||||
|
||||
## 参数说明
|
||||
|
||||
### API密钥配置节点
|
||||
|
||||
- `api_key`: Google Gemini API密钥
|
||||
|
||||
### 文本生成节点
|
||||
|
||||
#### 必填参数
|
||||
|
||||
- `config`: API密钥配置
|
||||
- `prompt`: 提示文本
|
||||
- `model`: 选择模型（gemini-2.0-flash, gemini-2.0-flash-exp, gemini-2.0-pro）
|
||||
- `temperature`: 生成温度（0.0-1.0），控制创意程度
|
||||
- `top_p`: 核采样参数（0.0-1.0）
|
||||
- `top_k`: 考虑的候选词数量（1-100）
|
||||
- `max_output_tokens`: 最大输出标记数（1-8192）
|
||||
- `include_images`: 是否在响应中包含图像（True/False）
|
||||
|
||||
#### 可选参数
|
||||
|
||||
- `input_image`: 输入图像（用于图像理解）
|
||||
- `use_proxy`: 是否使用代理（True/False）
|
||||
- `proxy_host`: 代理主机地址
|
||||
- `proxy_port`: 代理端口
|
||||
|
||||
## 输出
|
||||
|
||||
文本生成节点输出：
|
||||
- `text`: 生成的文本
|
||||
- `image`: 如果启用了图像生成，则输出图像
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 使用此节点需要稳定的网络连接或有效的代理设置
|
||||
- API请求可能会受到Google服务器负载的影响
|
||||
- 大型请求可能需要更长的处理时间
|
||||
- 图像生成功能需要使用支持图像生成的模型（如gemini-2.0-flash-exp）
|
||||
|
||||
## 故障排除
|
||||
|
||||
- 如果遇到网络连接问题，请检查代理设置
|
||||
- 如果API请求失败，请检查API密钥是否有效
|
||||
- 如果依赖安装失败，请手动安装所需的依赖包
|
||||
|
||||
## 许可证
|
||||
|
||||
[在此添加许可证信息]
|
||||
|
||||
## 致谢
|
||||
|
||||
感谢Google提供的Gemini API服务。
|
||||
@@ -0,0 +1,5 @@
|
||||
import os
|
||||
import folder_paths
|
||||
from .gemini_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
+418
@@ -0,0 +1,418 @@
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import io
|
||||
import folder_paths
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
import time
|
||||
import traceback
|
||||
import threading
|
||||
import queue
|
||||
import sys
|
||||
import importlib
|
||||
import subprocess
|
||||
import random
|
||||
|
||||
def check_and_install_dependencies():
|
||||
required_packages = {
|
||||
'requests': 'requests',
|
||||
'pysocks': 'PySocks',
|
||||
}
|
||||
|
||||
for module_name, package_name in required_packages.items():
|
||||
try:
|
||||
importlib.import_module(module_name)
|
||||
except ImportError:
|
||||
print(f"[WARNING] 未找到 {package_name}，尝试安装...")
|
||||
try:
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
|
||||
print(f"[INFO] 成功安装 {package_name}")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] 安装 {package_name} 失败: {str(e)}")
|
||||
|
||||
try:
|
||||
check_and_install_dependencies()
|
||||
except Exception as e:
|
||||
print(f"[WARNING] 检查依赖时出错: {str(e)}")
|
||||
|
||||
class SSL_GeminiAPIKeyConfig:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"api_key": ("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GEMINI_CONFIG",)
|
||||
RETURN_NAMES = ("config",)
|
||||
FUNCTION = "configure"
|
||||
CATEGORY = "💠SSL/API/Gemini"
|
||||
|
||||
def configure(self, api_key):
|
||||
config = {"api_key": api_key}
|
||||
return (config,)
|
||||
|
||||
|
||||
class SSL_GeminiTextPrompt:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"config": ("GEMINI_CONFIG",),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"model": (["gemini-2.0-flash", "gemini-2.0-flash-exp", "gemini-2.0-pro"], {"default": "gemini-2.0-flash"}),
|
||||
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"top_k": ("INT", {"default": 40, "min": 1, "max": 100, "step": 1}),
|
||||
"max_output_tokens": ("INT", {"default": 8192, "min": 1, "max": 8192, "step": 1}),
|
||||
"include_images": (["True", "False"], {"default": "True"}),
|
||||
},
|
||||
"optional": {
|
||||
"input_image": ("IMAGE",),
|
||||
"use_proxy": (["False", "True"], {"default": "False"}),
|
||||
"proxy_host": ("STRING", {"default": "127.0.0.1"}),
|
||||
"proxy_port": ("INT", {"default": 7890, "min": 1, "max": 65535}),
|
||||
"use_seed": (["False", "True"], {"default": "False"}),
|
||||
"seed_type": (["random", "fixed"], {"default": "random"}),
|
||||
"seed_value": ("INT", {"default": 66666, "min": 0, "max": 2147483647}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "IMAGE")
|
||||
RETURN_NAMES = ("text", "image")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "💠SSL/API/Gemini"
|
||||
|
||||
def save_binary_file(self, data, mime_type):
|
||||
ext = ".bin"
|
||||
if mime_type == "image/png":
|
||||
ext = ".png"
|
||||
elif mime_type == "image/jpeg":
|
||||
ext = ".jpg"
|
||||
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
gemini_dir = os.path.join(output_dir, "gemini_outputs")
|
||||
os.makedirs(gemini_dir, exist_ok=True)
|
||||
|
||||
file_name = os.path.join(gemini_dir, f"gemini_output_{uuid.uuid4()}{ext}")
|
||||
|
||||
with open(file_name, "wb") as f:
|
||||
f.write(data)
|
||||
|
||||
return file_name
|
||||
|
||||
def generate_empty_image(self, width=64, height=64):
|
||||
empty_image = np.ones((height, width, 3), dtype=np.float32) * 0.2
|
||||
tensor = torch.from_numpy(empty_image).unsqueeze(0)
|
||||
return tensor
|
||||
|
||||
def generate(self, config, prompt, model, temperature, top_p, top_k, max_output_tokens, include_images,
|
||||
input_image=None, use_proxy="False", proxy_host="127.0.0.1", proxy_port=7890,
|
||||
use_seed="False", seed_type="random", seed_value=42):
|
||||
original_http_proxy = os.environ.get('HTTP_PROXY')
|
||||
original_https_proxy = os.environ.get('HTTPS_PROXY')
|
||||
original_http_proxy_lower = os.environ.get('http_proxy')
|
||||
original_https_proxy_lower = os.environ.get('https_proxy')
|
||||
|
||||
print(f"[INFO] 开始生成，模型: {model}, 温度: {temperature}")
|
||||
|
||||
if use_seed == "True":
|
||||
if seed_type == "random":
|
||||
actual_seed = random.randint(0, 2147483647)
|
||||
else:
|
||||
actual_seed = seed_value
|
||||
|
||||
random.seed(actual_seed)
|
||||
np.random.seed(actual_seed)
|
||||
torch.manual_seed(actual_seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(actual_seed)
|
||||
else:
|
||||
actual_seed = None
|
||||
|
||||
try:
|
||||
if use_proxy == "True":
|
||||
if not proxy_host.startswith(('http://', 'https://')):
|
||||
proxy_url = f"http://{proxy_host}:{proxy_port}"
|
||||
else:
|
||||
proxy_url = f"{proxy_host}:{proxy_port}"
|
||||
|
||||
os.environ['HTTP_PROXY'] = proxy_url
|
||||
os.environ['HTTPS_PROXY'] = proxy_url
|
||||
os.environ['http_proxy'] = proxy_url
|
||||
os.environ['https_proxy'] = proxy_url
|
||||
os.environ['REQUESTS_CA_BUNDLE'] = ''
|
||||
|
||||
print(f"[INFO] 已启用代理: {proxy_url}")
|
||||
|
||||
print(f"[INFO] 初始化Gemini客户端")
|
||||
try:
|
||||
client_options = {}
|
||||
|
||||
if use_proxy == "True":
|
||||
try:
|
||||
import google.api_core.http_client
|
||||
import google.auth.transport.requests
|
||||
|
||||
try:
|
||||
import requests
|
||||
from requests.adapters import HTTPAdapter
|
||||
|
||||
class ProxyAdapter(HTTPAdapter):
|
||||
def __init__(self, proxy_url, **kwargs):
|
||||
self.proxy_url = proxy_url
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def add_headers(self, request, **kwargs):
|
||||
super().add_headers(request, **kwargs)
|
||||
|
||||
session = requests.Session()
|
||||
proxies = {
|
||||
"http": proxy_url,
|
||||
"https": proxy_url
|
||||
}
|
||||
session.proxies.update(proxies)
|
||||
|
||||
adapter = ProxyAdapter(proxy_url, max_retries=3)
|
||||
session.mount('http://', adapter)
|
||||
session.mount('https://', adapter)
|
||||
|
||||
session.verify = False
|
||||
|
||||
http_client = google.api_core.http_client.RequestsHttpClient(session=session)
|
||||
client_options["http_client"] = http_client
|
||||
|
||||
print(f"[INFO] 已使用requests库设置代理到HTTP客户端")
|
||||
except Exception as proxy_error:
|
||||
print(f"[WARNING] 设置HTTP客户端代理失败: {str(proxy_error)}")
|
||||
except ImportError as e:
|
||||
print(f"[WARNING] 导入Google API HTTP客户端库失败: {str(e)}")
|
||||
|
||||
client = genai.Client(api_key=config["api_key"], **client_options)
|
||||
print(f"[INFO] Gemini客户端初始化成功")
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Gemini客户端初始化失败: {str(e)}")
|
||||
return (f"Gemini客户端初始化失败: {str(e)}", self.generate_empty_image())
|
||||
|
||||
try:
|
||||
import socket
|
||||
test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
test_socket.settimeout(5)
|
||||
|
||||
test_host = "generativelanguage.googleapis.com"
|
||||
|
||||
if use_proxy == "True":
|
||||
try:
|
||||
import socks
|
||||
test_socket = socks.socksocket()
|
||||
test_socket.set_proxy(socks.HTTP, proxy_host, proxy_port)
|
||||
test_socket.settimeout(5)
|
||||
except ImportError:
|
||||
print(f"[WARNING] 未安装PySocks库，无法通过代理测试连接")
|
||||
|
||||
test_socket.connect((test_host, 443))
|
||||
test_socket.close()
|
||||
network_ok = True
|
||||
except Exception as e:
|
||||
network_ok = False
|
||||
print(f"[WARNING] 网络连接测试失败: {str(e)}")
|
||||
|
||||
if use_proxy == "True":
|
||||
try:
|
||||
import subprocess
|
||||
cmd = f"curl -x {proxy_url} -s -o /dev/null -w '%{{http_code}}' https://{test_host}"
|
||||
try:
|
||||
result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=10)
|
||||
if result.returncode == 0 and result.stdout.strip() in ['200', '301', '302']:
|
||||
network_ok = True
|
||||
else:
|
||||
print(f"[WARNING] curl代理测试失败，状态码: {result.stdout.strip() if result.stdout else 'N/A'}")
|
||||
except Exception as curl_error:
|
||||
print(f"[WARNING] curl测试失败: {str(curl_error)}")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
contents = []
|
||||
|
||||
if input_image is not None:
|
||||
try:
|
||||
img_array = input_image[0].cpu().numpy()
|
||||
img_array = (img_array * 255).astype(np.uint8)
|
||||
pil_img = Image.fromarray(img_array)
|
||||
|
||||
img_byte_arr = io.BytesIO()
|
||||
pil_img.save(img_byte_arr, format='PNG')
|
||||
img_bytes = img_byte_arr.getvalue()
|
||||
|
||||
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}}
|
||||
txt_part = {"text": prompt}
|
||||
|
||||
contents = [img_part, txt_part]
|
||||
except Exception as e:
|
||||
print(f"[ERROR] 处理输入图像时出错: {str(e)}")
|
||||
return (f"处理输入图像时出错: {str(e)}", self.generate_empty_image())
|
||||
else:
|
||||
contents = prompt
|
||||
|
||||
response_modalities = ["text"]
|
||||
if include_images == "True":
|
||||
response_modalities.append("image")
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
max_output_tokens=max_output_tokens,
|
||||
response_modalities=response_modalities,
|
||||
)
|
||||
|
||||
if use_seed == "True" and actual_seed is not None:
|
||||
try:
|
||||
generate_content_config.seed = actual_seed
|
||||
except Exception as seed_error:
|
||||
print(f"[WARNING] 无法设置种子到API请求: {str(seed_error)}")
|
||||
|
||||
text_output = ""
|
||||
image_tensor = None
|
||||
|
||||
try:
|
||||
print(f"[INFO] 发送API请求到Gemini")
|
||||
|
||||
start_time = time.time()
|
||||
timeout = 30
|
||||
|
||||
def api_call():
|
||||
try:
|
||||
api_response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=contents,
|
||||
config=generate_content_config,
|
||||
)
|
||||
result_queue.put(("success", api_response))
|
||||
except Exception as e:
|
||||
print(f"[ERROR] API线程中出错: {str(e)}")
|
||||
result_queue.put(("error", e))
|
||||
|
||||
result_queue = queue.Queue()
|
||||
|
||||
api_thread = threading.Thread(target=api_call)
|
||||
api_thread.daemon = True
|
||||
api_thread.start()
|
||||
|
||||
try:
|
||||
status, result = result_queue.get(timeout=timeout)
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
if status == "success":
|
||||
response = result
|
||||
print(f"[INFO] API请求成功完成，耗时: {elapsed_time:.2f}秒")
|
||||
else:
|
||||
print(f"[ERROR] API请求线程中出错，耗时: {elapsed_time:.2f}秒，错误: {str(result)}")
|
||||
error_str = str(result).lower()
|
||||
if any(term in error_str for term in ["timeout", "connection", "network", "socket", "连接", "网络"]):
|
||||
if not network_ok and use_proxy != "True":
|
||||
return (f"API请求失败: {str(result)}。网络连接测试失败，建议启用代理。", self.generate_empty_image())
|
||||
raise result
|
||||
except queue.Empty:
|
||||
elapsed_time = time.time() - start_time
|
||||
print(f"[ERROR] API请求超时，已等待: {elapsed_time:.2f}秒")
|
||||
|
||||
timeout_msg = f"Gemini API请求超时，已等待{timeout}秒。"
|
||||
if not network_ok:
|
||||
if use_proxy == "True":
|
||||
timeout_msg += f"网络连接测试失败，当前使用的代理({proxy_host}:{proxy_port})可能无效，请检查代理设置。"
|
||||
else:
|
||||
timeout_msg += "网络连接测试失败，建议启用代理。"
|
||||
else:
|
||||
timeout_msg += "网络连接测试成功，但API请求仍然超时，可能是服务器繁忙或请求内容过大。"
|
||||
|
||||
return (timeout_msg, self.generate_empty_image())
|
||||
|
||||
print(f"[INFO] 收到API响应")
|
||||
|
||||
if hasattr(response, 'candidates') and response.candidates:
|
||||
for part in response.candidates[0].content.parts:
|
||||
if hasattr(part, 'text') and part.text is not None:
|
||||
text_output += part.text
|
||||
|
||||
elif hasattr(part, 'inline_data') and part.inline_data is not None:
|
||||
try:
|
||||
inline_data = part.inline_data
|
||||
mime_type = inline_data.mime_type
|
||||
data = inline_data.data
|
||||
|
||||
image_path = self.save_binary_file(data, mime_type)
|
||||
|
||||
img = Image.open(image_path)
|
||||
|
||||
if img.mode != 'RGB':
|
||||
img = img.convert('RGB')
|
||||
|
||||
img_array = np.array(img).astype(np.float32) / 255.0
|
||||
image_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
||||
except Exception as e:
|
||||
print(f"[ERROR] 图像处理错误: {str(e)}")
|
||||
text_output += f"\n图像处理错误: {str(e)}"
|
||||
except Exception as e:
|
||||
print(f"[ERROR] API调用错误: {str(e)}")
|
||||
text_output = f"API调用错误: {str(e)}"
|
||||
|
||||
if image_tensor is None:
|
||||
image_tensor = self.generate_empty_image()
|
||||
|
||||
if use_seed == "True" and actual_seed is not None:
|
||||
seed_info = f"\n\n[种子信息: {actual_seed}]"
|
||||
text_output += seed_info
|
||||
|
||||
return (text_output, image_tensor)
|
||||
finally:
|
||||
if original_http_proxy:
|
||||
os.environ['HTTP_PROXY'] = original_http_proxy
|
||||
else:
|
||||
if 'HTTP_PROXY' in os.environ:
|
||||
os.environ.pop('HTTP_PROXY')
|
||||
|
||||
if original_https_proxy:
|
||||
os.environ['HTTPS_PROXY'] = original_https_proxy
|
||||
else:
|
||||
if 'HTTPS_PROXY' in os.environ:
|
||||
os.environ.pop('HTTPS_PROXY')
|
||||
|
||||
if original_http_proxy_lower:
|
||||
os.environ['http_proxy'] = original_http_proxy_lower
|
||||
else:
|
||||
if 'http_proxy' in os.environ:
|
||||
os.environ.pop('http_proxy')
|
||||
|
||||
if original_https_proxy_lower:
|
||||
os.environ['https_proxy'] = original_https_proxy_lower
|
||||
else:
|
||||
if 'https_proxy' in os.environ:
|
||||
os.environ.pop('https_proxy')
|
||||
|
||||
if 'REQUESTS_CA_BUNDLE' in os.environ:
|
||||
os.environ.pop('REQUESTS_CA_BUNDLE')
|
||||
|
||||
try:
|
||||
import requests
|
||||
if hasattr(requests, 'Session'):
|
||||
clean_session = requests.Session()
|
||||
if hasattr(requests, 'session'):
|
||||
requests.session = lambda: clean_session
|
||||
except:
|
||||
pass
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SSL_GeminiAPIKeyConfig": SSL_GeminiAPIKeyConfig,
|
||||
"SSL_GeminiTextPrompt": SSL_GeminiTextPrompt,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SSL_GeminiAPIKeyConfig": "💠SSL Gemini API Key",
|
||||
"SSL_GeminiTextPrompt": "💠SSL Gemini Text Prompt",
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
google-genai
|
||||
numpy
|
||||
torch
|
||||
Reference in New Issue
Block a user