554 lines
20 KiB
Python
554 lines
20 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import requests
|
|
from PIL import Image
|
|
|
|
try:
|
|
from comfy_execution.graph_utils import ExecutionBlocker
|
|
except Exception:
|
|
ExecutionBlocker = None
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
NODE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
PROTOCOL_OPENAI_CHAT = "openai_chat_completions"
|
|
PROTOCOL_OPENAI_RESPONSES = "openai_responses"
|
|
PROTOCOL_ANTHROPIC_MESSAGES = "anthropic_messages"
|
|
PROTOCOL_GEMINI_GENERATE_CONTENT = "gemini_generatecontent"
|
|
|
|
API_PROTOCOLS = [
|
|
PROTOCOL_OPENAI_CHAT,
|
|
PROTOCOL_OPENAI_RESPONSES,
|
|
PROTOCOL_ANTHROPIC_MESSAGES,
|
|
PROTOCOL_GEMINI_GENERATE_CONTENT,
|
|
]
|
|
|
|
PROTOCOL_TOOLTIP = (
|
|
"Choose the JSON protocol sent to the server. "
|
|
"openai_chat_completions -> /chat/completions; "
|
|
"openai_responses -> /responses; "
|
|
"anthropic_messages -> /messages; "
|
|
"gemini_generatecontent -> /models/{model}:generateContent. "
|
|
"Select the format expected by your API service."
|
|
)
|
|
|
|
MODEL_TOOLTIP = (
|
|
"Model name passed through unchanged, for example gpt-4o, claude-3-5-sonnet-latest, "
|
|
"gemini-2.5-flash, or a model alias configured by your API service."
|
|
)
|
|
|
|
BASE_URL_TOOLTIP = (
|
|
"Base URL for the API service. The node appends the endpoint for the selected protocol. "
|
|
"Examples: https://api.openai.com/v1, https://api.anthropic.com/v1, "
|
|
"or https://generativelanguage.googleapis.com/v1beta. "
|
|
"For Gemini you may also provide a full URL ending in :generateContent or include {model}."
|
|
)
|
|
|
|
KEY_FILE_TOOLTIP = (
|
|
"File name only, read from this node directory. Do not enter an API key here. "
|
|
"Default: api_key.txt. This keeps secrets out of workflow JSON files. "
|
|
"Absolute paths and ../ are rejected."
|
|
)
|
|
|
|
DEBUG_OUTPUT_TOOLTIP = (
|
|
"Second output contains request/response debug JSON. Auth headers are redacted, "
|
|
"and large image payloads may be omitted."
|
|
)
|
|
|
|
|
|
def read_api_key(api_key_file: str = "api_key.txt") -> str:
|
|
key_filename = os.path.basename(api_key_file.strip() or "api_key.txt")
|
|
if key_filename != (api_key_file.strip() or "api_key.txt"):
|
|
raise ValueError("api_key_file must be a file name in this node directory, not a path.")
|
|
|
|
key_path = os.path.join(NODE_DIR, key_filename)
|
|
if not os.path.exists(key_path):
|
|
with open(key_path, "w", encoding="utf-8") as key_file:
|
|
key_file.write("YOUR_API_KEY_HERE")
|
|
log.warning("API key file not found. Created template at: %s", key_path)
|
|
return ""
|
|
|
|
with open(key_path, "r", encoding="utf-8") as key_file:
|
|
key = key_file.read().strip()
|
|
if not key or key == "YOUR_API_KEY_HERE":
|
|
log.warning("API key is missing or placeholder in %s", key_path)
|
|
return ""
|
|
return key
|
|
|
|
|
|
def tensor_to_pil(image: Any) -> Image.Image:
|
|
if image.ndim == 4:
|
|
image = image[0]
|
|
if image.ndim != 3:
|
|
raise ValueError("IMAGE tensor must be HWC or NHWC.")
|
|
|
|
image_np = image.detach().cpu().numpy()
|
|
image_np = np.clip(image_np * 255.0, 0, 255).astype(np.uint8)
|
|
if image_np.shape[-1] == 1:
|
|
return Image.fromarray(image_np[:, :, 0], mode="L")
|
|
if image_np.shape[-1] == 3:
|
|
return Image.fromarray(image_np, mode="RGB")
|
|
if image_np.shape[-1] == 4:
|
|
return Image.fromarray(image_np, mode="RGBA")
|
|
raise ValueError(f"Unsupported IMAGE channel count: {image_np.shape[-1]}")
|
|
|
|
|
|
def tensor_to_png_base64(image: Any) -> str:
|
|
pil_image = tensor_to_pil(image)
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
|
|
|
|
def collect_image_b64(*images: Optional[Any]) -> List[str]:
|
|
return [tensor_to_png_base64(image) for image in images if image is not None]
|
|
|
|
|
|
def normalize_base_url(api_base_url: str, default_base_url: str) -> str:
|
|
base_url = (api_base_url or "").strip() or default_base_url
|
|
return base_url.rstrip("/")
|
|
|
|
|
|
def redact_headers(headers: Dict[str, str]) -> Dict[str, str]:
|
|
redacted = dict(headers)
|
|
for key in list(redacted.keys()):
|
|
if key.lower() in {"authorization", "x-api-key", "x-goog-api-key"}:
|
|
redacted[key] = "***"
|
|
return redacted
|
|
|
|
|
|
def trim_debug_payload(payload: Any) -> Any:
|
|
text = json.dumps(payload, ensure_ascii=False)
|
|
if len(text) <= 8000:
|
|
return payload
|
|
return {"notice": "Payload omitted because it is large; image data was included in the request."}
|
|
|
|
|
|
def first_text_from_openai_chat(response_json: Dict[str, Any]) -> str:
|
|
return response_json["choices"][0]["message"].get("content") or ""
|
|
|
|
|
|
def first_text_from_openai_responses(response_json: Dict[str, Any]) -> str:
|
|
if response_json.get("output_text"):
|
|
return response_json["output_text"]
|
|
|
|
chunks: List[str] = []
|
|
for item in response_json.get("output", []):
|
|
for content in item.get("content", []):
|
|
if content.get("type") in {"output_text", "text"} and content.get("text"):
|
|
chunks.append(content["text"])
|
|
return "".join(chunks)
|
|
|
|
|
|
def first_text_from_anthropic(response_json: Dict[str, Any]) -> str:
|
|
chunks: List[str] = []
|
|
for item in response_json.get("content", []):
|
|
if item.get("type") == "text" and item.get("text"):
|
|
chunks.append(item["text"])
|
|
return "".join(chunks)
|
|
|
|
|
|
def first_text_from_gemini(response_json: Dict[str, Any]) -> str:
|
|
chunks: List[str] = []
|
|
for candidate in response_json.get("candidates", []):
|
|
for part in candidate.get("content", {}).get("parts", []):
|
|
if part.get("text"):
|
|
chunks.append(part["text"])
|
|
return "".join(chunks)
|
|
|
|
|
|
class FOKMultiProtocolChatVisionAPI:
|
|
"""
|
|
Sends text and up to four images to OpenAI, Anthropic, Gemini, or compatible API services.
|
|
"""
|
|
|
|
CATEGORY = "FOK API Tools/API"
|
|
DESCRIPTION = (
|
|
"Send text and up to four images to OpenAI Chat Completions, OpenAI Responses, "
|
|
"Anthropic Messages, Gemini GenerateContent, or another compatible API service. "
|
|
"API keys are loaded from a local file instead of being stored in the workflow."
|
|
)
|
|
FUNCTION = "execute"
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("response_text", "debug_info")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_protocol": (
|
|
API_PROTOCOLS,
|
|
{"default": PROTOCOL_OPENAI_CHAT, "tooltip": PROTOCOL_TOOLTIP},
|
|
),
|
|
"model": ("STRING", {"default": "gpt-4o", "tooltip": MODEL_TOOLTIP}),
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"tooltip": "User prompt sent with the optional images. Connect a text node here.",
|
|
},
|
|
),
|
|
"api_base_url": (
|
|
"STRING",
|
|
{
|
|
"multiline": False,
|
|
"default": "https://api.openai.com/v1",
|
|
"tooltip": BASE_URL_TOOLTIP,
|
|
},
|
|
),
|
|
"on_error": (
|
|
["skip", "stop", "raise"],
|
|
{
|
|
"default": "skip",
|
|
"tooltip": (
|
|
"Failure behavior. skip returns the original prompt; stop blocks downstream nodes "
|
|
"when ExecutionBlocker is available; raise throws the error and stops execution."
|
|
),
|
|
},
|
|
),
|
|
"timeout": (
|
|
"INT",
|
|
{
|
|
"default": 120,
|
|
"min": 10,
|
|
"max": 600,
|
|
"step": 10,
|
|
"tooltip": "HTTP request timeout in seconds. Increase for slow models or image-heavy requests.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"system_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": (
|
|
"Optional system/instruction text. Mapped to system, instructions, "
|
|
"or systemInstruction depending on the selected protocol."
|
|
),
|
|
},
|
|
),
|
|
"image_1": ("IMAGE", {"tooltip": "Optional first image. Sent as PNG base64."}),
|
|
"image_2": ("IMAGE", {"tooltip": "Optional second image. Sent as PNG base64."}),
|
|
"image_3": ("IMAGE", {"tooltip": "Optional third image. Sent as PNG base64."}),
|
|
"image_4": ("IMAGE", {"tooltip": "Optional fourth image. Sent as PNG base64."}),
|
|
"max_tokens": (
|
|
"INT",
|
|
{
|
|
"default": 2048,
|
|
"min": 1,
|
|
"max": 32768,
|
|
"step": 1,
|
|
"tooltip": (
|
|
"Maximum text tokens to generate. Sent as max_tokens for Chat/Anthropic, "
|
|
"max_output_tokens for Responses, and maxOutputTokens for Gemini."
|
|
),
|
|
},
|
|
),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.05,
|
|
"tooltip": "Sampling temperature. Lower is more deterministic; higher is more varied.",
|
|
},
|
|
),
|
|
"api_key_file": (
|
|
"STRING",
|
|
{
|
|
"multiline": False,
|
|
"default": "api_key.txt",
|
|
"tooltip": KEY_FILE_TOOLTIP,
|
|
},
|
|
),
|
|
"anthropic_version": (
|
|
"STRING",
|
|
{
|
|
"multiline": False,
|
|
"default": "2023-06-01",
|
|
"tooltip": (
|
|
"Anthropic API version header used only for anthropic_messages. "
|
|
"Leave default unless your API service requires another value."
|
|
),
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
def execute(
|
|
self,
|
|
api_protocol: str,
|
|
model: str,
|
|
prompt: str,
|
|
api_base_url: str,
|
|
on_error: str,
|
|
timeout: int,
|
|
system_prompt: str = "",
|
|
image_1: Optional[Any] = None,
|
|
image_2: Optional[Any] = None,
|
|
image_3: Optional[Any] = None,
|
|
image_4: Optional[Any] = None,
|
|
max_tokens: int = 2048,
|
|
temperature: float = 0.7,
|
|
api_key_file: str = "api_key.txt",
|
|
anthropic_version: str = "2023-06-01",
|
|
) -> Tuple[Any, str]:
|
|
debug_data: Dict[str, Any] = {"request": {}, "response": {}}
|
|
|
|
try:
|
|
if api_protocol not in API_PROTOCOLS:
|
|
raise ValueError(f"Unsupported api_protocol: {api_protocol}")
|
|
if not model.strip():
|
|
raise ValueError("Model ID is required.")
|
|
|
|
key = read_api_key(api_key_file)
|
|
image_b64_list = collect_image_b64(image_1, image_2, image_3, image_4)
|
|
|
|
url, headers, payload = self.build_request(
|
|
api_protocol=api_protocol,
|
|
model=model.strip(),
|
|
prompt=prompt,
|
|
system_prompt=system_prompt,
|
|
api_base_url=api_base_url,
|
|
api_key=key,
|
|
image_b64_list=image_b64_list,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
anthropic_version=anthropic_version,
|
|
)
|
|
|
|
debug_data["request"] = {
|
|
"url": url,
|
|
"method": "POST",
|
|
"headers": redact_headers(headers),
|
|
"payload": trim_debug_payload(payload),
|
|
}
|
|
|
|
response = requests.post(url, headers=headers, json=payload, timeout=timeout)
|
|
debug_data["response"]["status_code"] = response.status_code
|
|
response.raise_for_status()
|
|
|
|
response_json = response.json()
|
|
debug_data["response"]["body"] = response_json
|
|
response_text = self.parse_response(api_protocol, response_json)
|
|
if not response_text:
|
|
raise RuntimeError("API returned an empty text response.")
|
|
return (response_text, json.dumps(debug_data, indent=2, ensure_ascii=False))
|
|
|
|
except Exception as exc:
|
|
log.error("%s failed: %s", self.__class__.__name__, exc)
|
|
if isinstance(exc, requests.exceptions.RequestException) and exc.response is not None:
|
|
debug_data["response"]["status_code"] = exc.response.status_code
|
|
debug_data["response"]["body"] = exc.response.text
|
|
debug_data["error"] = str(exc)
|
|
debug_json = json.dumps(debug_data, indent=2, ensure_ascii=False)
|
|
|
|
if on_error == "stop":
|
|
if ExecutionBlocker is None:
|
|
return ("", debug_json)
|
|
return (ExecutionBlocker(None), debug_json)
|
|
if on_error == "skip":
|
|
return (prompt, debug_json)
|
|
raise
|
|
|
|
def build_request(
|
|
self,
|
|
api_protocol: str,
|
|
model: str,
|
|
prompt: str,
|
|
system_prompt: str,
|
|
api_base_url: str,
|
|
api_key: str,
|
|
image_b64_list: List[str],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
anthropic_version: str,
|
|
) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
|
|
if api_protocol == PROTOCOL_OPENAI_CHAT:
|
|
return self.build_openai_chat_request(
|
|
model, prompt, system_prompt, api_base_url, api_key, image_b64_list, max_tokens, temperature
|
|
)
|
|
if api_protocol == PROTOCOL_OPENAI_RESPONSES:
|
|
return self.build_openai_responses_request(
|
|
model, prompt, system_prompt, api_base_url, api_key, image_b64_list, max_tokens, temperature
|
|
)
|
|
if api_protocol == PROTOCOL_ANTHROPIC_MESSAGES:
|
|
return self.build_anthropic_messages_request(
|
|
model,
|
|
prompt,
|
|
system_prompt,
|
|
api_base_url,
|
|
api_key,
|
|
image_b64_list,
|
|
max_tokens,
|
|
temperature,
|
|
anthropic_version,
|
|
)
|
|
return self.build_gemini_generate_content_request(
|
|
model, prompt, system_prompt, api_base_url, api_key, image_b64_list, max_tokens, temperature
|
|
)
|
|
|
|
def build_openai_chat_request(
|
|
self,
|
|
model: str,
|
|
prompt: str,
|
|
system_prompt: str,
|
|
api_base_url: str,
|
|
api_key: str,
|
|
image_b64_list: List[str],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
|
|
url = f"{normalize_base_url(api_base_url, 'https://api.openai.com/v1')}/chat/completions"
|
|
headers = {"Content-Type": "application/json"}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
messages: List[Dict[str, Any]] = []
|
|
if system_prompt.strip():
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
|
|
content: List[Dict[str, Any]] = [{"type": "text", "text": prompt}]
|
|
for image_b64 in image_b64_list:
|
|
content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}})
|
|
messages.append({"role": "user", "content": content})
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
}
|
|
return url, headers, payload
|
|
|
|
def build_openai_responses_request(
|
|
self,
|
|
model: str,
|
|
prompt: str,
|
|
system_prompt: str,
|
|
api_base_url: str,
|
|
api_key: str,
|
|
image_b64_list: List[str],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
|
|
url = f"{normalize_base_url(api_base_url, 'https://api.openai.com/v1')}/responses"
|
|
headers = {"Content-Type": "application/json"}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
content: List[Dict[str, Any]] = [{"type": "input_text", "text": prompt}]
|
|
for image_b64 in image_b64_list:
|
|
content.append({"type": "input_image", "image_url": f"data:image/png;base64,{image_b64}"})
|
|
|
|
payload: Dict[str, Any] = {
|
|
"model": model,
|
|
"input": [{"role": "user", "content": content}],
|
|
"max_output_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
}
|
|
if system_prompt.strip():
|
|
payload["instructions"] = system_prompt
|
|
return url, headers, payload
|
|
|
|
def build_anthropic_messages_request(
|
|
self,
|
|
model: str,
|
|
prompt: str,
|
|
system_prompt: str,
|
|
api_base_url: str,
|
|
api_key: str,
|
|
image_b64_list: List[str],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
anthropic_version: str,
|
|
) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
|
|
url = f"{normalize_base_url(api_base_url, 'https://api.anthropic.com/v1')}/messages"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"anthropic-version": anthropic_version.strip() or "2023-06-01",
|
|
}
|
|
if api_key:
|
|
headers["x-api-key"] = api_key
|
|
|
|
content: List[Dict[str, Any]] = [{"type": "text", "text": prompt}]
|
|
for image_b64 in image_b64_list:
|
|
content.append(
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": "image/png",
|
|
"data": image_b64,
|
|
},
|
|
}
|
|
)
|
|
|
|
payload: Dict[str, Any] = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": content}],
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
}
|
|
if system_prompt.strip():
|
|
payload["system"] = system_prompt
|
|
return url, headers, payload
|
|
|
|
def build_gemini_generate_content_request(
|
|
self,
|
|
model: str,
|
|
prompt: str,
|
|
system_prompt: str,
|
|
api_base_url: str,
|
|
api_key: str,
|
|
image_b64_list: List[str],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
|
|
base_url = normalize_base_url(api_base_url, "https://generativelanguage.googleapis.com/v1beta")
|
|
if "{model}" in base_url:
|
|
url = base_url.replace("{model}", model)
|
|
elif base_url.endswith(":generateContent"):
|
|
url = base_url
|
|
else:
|
|
url = f"{base_url}/models/{model}:generateContent"
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
if api_key:
|
|
headers["x-goog-api-key"] = api_key
|
|
|
|
parts: List[Dict[str, Any]] = [{"text": prompt}]
|
|
for image_b64 in image_b64_list:
|
|
parts.append({"inline_data": {"mime_type": "image/png", "data": image_b64}})
|
|
|
|
payload: Dict[str, Any] = {
|
|
"contents": [{"role": "user", "parts": parts}],
|
|
"generationConfig": {
|
|
"maxOutputTokens": max_tokens,
|
|
"temperature": temperature,
|
|
},
|
|
}
|
|
if system_prompt.strip():
|
|
payload["systemInstruction"] = {"parts": [{"text": system_prompt}]}
|
|
return url, headers, payload
|
|
|
|
def parse_response(self, api_protocol: str, response_json: Dict[str, Any]) -> str:
|
|
if api_protocol == PROTOCOL_OPENAI_CHAT:
|
|
return first_text_from_openai_chat(response_json)
|
|
if api_protocol == PROTOCOL_OPENAI_RESPONSES:
|
|
return first_text_from_openai_responses(response_json)
|
|
if api_protocol == PROTOCOL_ANTHROPIC_MESSAGES:
|
|
return first_text_from_anthropic(response_json)
|
|
return first_text_from_gemini(response_json)
|