boos
305 lines
12 KiB
Python
305 lines
12 KiB
Python
"""
|
||
BoosAPI ComfyUI Custom Nodes
|
||
|
||
让 ComfyUI workflow 调用 BoosAPI 的 AI 模型（图像生成、LLM 等）。
|
||
配置：通过环境变量设置 BoosAPI 地址，或使用默认 http://127.0.0.1:4000/v1
|
||
|
||
环境变量：
|
||
BOOSAPI_BASE — BoosAPI 地址（默认 http://127.0.0.1:4000）
|
||
BOOSAPI_KEY — API Key（默认 sk-boos4721）
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import requests
|
||
|
||
# ── Config ───────────────────────────────────────────────────────────────
|
||
BOOSAPI_BASE_RAW = os.environ.get("BOOSAPI_BASE", "http://127.0.0.1:4000")
|
||
BOOSAPI_KEY = os.environ.get("BOOSAPI_KEY", "sk-boos4721")
|
||
HEADERS = {
|
||
"Authorization": f"Bearer {BOOSAPI_KEY}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
FALLBACK_CHAT_MODELS = ["gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex", "gpt-5.2"]
|
||
FALLBACK_IMAGE_MODELS = ["gpt-image-2", "gpt-image-1", "dall-e-3", "dall-e-2"]
|
||
FALLBACK_ALL_MODELS = FALLBACK_IMAGE_MODELS + FALLBACK_CHAT_MODELS
|
||
|
||
|
||
def normalize_boosapi_root(base):
|
||
base = str(base or "http://127.0.0.1:4000").strip().rstrip("/")
|
||
if base.endswith("/v1"):
|
||
base = base[:-3].rstrip("/")
|
||
return base or "http://127.0.0.1:4000"
|
||
|
||
|
||
BOOSAPI_ROOT = normalize_boosapi_root(BOOSAPI_BASE_RAW)
|
||
BOOSAPI_V1 = f"{BOOSAPI_ROOT}/v1"
|
||
|
||
|
||
def raise_for_status_with_body(resp):
|
||
try:
|
||
resp.raise_for_status()
|
||
except requests.HTTPError as e:
|
||
body = (resp.text or "").strip()
|
||
if body:
|
||
print(f"[BoosAPI] HTTP {resp.status_code} response: {body[:2000]}")
|
||
raise requests.HTTPError(f"{e}; response body: {body[:1000]}", response=resp)
|
||
raise
|
||
|
||
|
||
def _ordered_unique(values):
|
||
seen = set()
|
||
result = []
|
||
for value in values:
|
||
value = str(value or "").strip()
|
||
if not value or value in seen:
|
||
continue
|
||
seen.add(value)
|
||
result.append(value)
|
||
return result
|
||
|
||
|
||
def fetch_boosapi_models():
|
||
"""Fetch models exposed by BoosAPI's own /v1/models route."""
|
||
try:
|
||
resp = requests.get(f"{BOOSAPI_V1}/models", headers=HEADERS, timeout=8)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
models = []
|
||
for item in data.get("data", []):
|
||
if isinstance(item, dict):
|
||
models.append(item.get("id") or item.get("name"))
|
||
elif isinstance(item, str):
|
||
models.append(item)
|
||
return _ordered_unique(models)
|
||
except Exception as e:
|
||
print(f"[BoosAPI] model list fallback: {e}")
|
||
return None
|
||
|
||
|
||
def model_choices(kind="all"):
|
||
models = fetch_boosapi_models()
|
||
if models is None:
|
||
if kind == "image":
|
||
return FALLBACK_IMAGE_MODELS
|
||
if kind == "chat":
|
||
return FALLBACK_CHAT_MODELS
|
||
return FALLBACK_ALL_MODELS
|
||
if kind == "image":
|
||
image_like = [
|
||
model for model in models
|
||
if any(token in model.lower() for token in ("image", "dall", "flux", "sd", "midjourney"))
|
||
]
|
||
return _ordered_unique(image_like + models) or FALLBACK_IMAGE_MODELS
|
||
if kind == "chat":
|
||
chat_like = [
|
||
model for model in models
|
||
if not any(token in model.lower() for token in ("image", "dall", "flux", "sd", "midjourney"))
|
||
]
|
||
return _ordered_unique(chat_like + models) or FALLBACK_CHAT_MODELS
|
||
return _ordered_unique(models) or FALLBACK_ALL_MODELS
|
||
|
||
|
||
def default_model(kind="all"):
|
||
choices = model_choices(kind)
|
||
return choices[0] if choices else FALLBACK_ALL_MODELS[0]
|
||
|
||
# ── Node Mappings ────────────────────────────────────────────────────────
|
||
NODE_CLASS_MAPPINGS = {}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
# Node: BoosAPIImageGen — 图像生成 (OpenAI DALL-E / 兼容格式)
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
|
||
class BoosAPIImageGen:
|
||
"""调用 BoosAPI /v1/images/generations 生成图像"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
||
"model": (model_choices("image"), {"default": default_model("image")}),
|
||
"size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
|
||
"quality": (["standard", "hd"], {"default": "standard"}),
|
||
"n": ("INT", {"default": 1, "min": 1, "max": 4}),
|
||
},
|
||
"optional": {
|
||
"style": (["vivid", "natural"], {"default": "vivid"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "STRING")
|
||
RETURN_NAMES = ("image", "metadata")
|
||
FUNCTION = "generate"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "BoosAPI"
|
||
|
||
def generate(self, prompt, model, size, quality, n, style="vivid"):
|
||
url = f"{BOOSAPI_V1}/images/generations"
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": n,
|
||
"size": size,
|
||
"quality": quality,
|
||
"style": style,
|
||
"response_format": "b64_json",
|
||
}
|
||
try:
|
||
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
|
||
raise_for_status_with_body(resp)
|
||
data = resp.json()
|
||
images = data.get("data", [])
|
||
# Return first image as b64
|
||
from io import BytesIO
|
||
import base64
|
||
import torch
|
||
from PIL import Image
|
||
|
||
b64_data = images[0]["b64_json"]
|
||
img_bytes = base64.b64decode(b64_data)
|
||
pil_img = Image.open(BytesIO(img_bytes)).convert("RGB")
|
||
img_tensor = torch.from_numpy(
|
||
__import__("numpy").array(pil_img).astype("float32") / 255.0
|
||
).unsqueeze(0)
|
||
|
||
metadata = json.dumps({
|
||
"model": model,
|
||
"size": size,
|
||
"revised_prompt": data.get("data", [{}])[0].get("revised_prompt", ""),
|
||
})
|
||
return (img_tensor, metadata)
|
||
except Exception as e:
|
||
print(f"[BoosAPI] ImageGen failed: {e}")
|
||
raise
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
# Node: BoosAPILLM — LLM Chat (流式文本生成)
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
|
||
class BoosAPILLM:
|
||
"""调用 BoosAPI /v1/chat/completions 生成文本"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
||
"model": (model_choices("chat"), {"default": default_model("chat")}),
|
||
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0}),
|
||
"max_tokens": ("INT", {"default": 2048, "min": 1, "max": 16384}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("text",)
|
||
FUNCTION = "chat"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "BoosAPI"
|
||
|
||
def chat(self, prompt, model, system_prompt, temperature, max_tokens):
|
||
url = f"{BOOSAPI_V1}/chat/completions"
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
messages.append({"role": "user", "content": prompt})
|
||
|
||
payload = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": False,
|
||
}
|
||
try:
|
||
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
|
||
raise_for_status_with_body(resp)
|
||
data = resp.json()
|
||
text = data["choices"][0]["message"]["content"]
|
||
return {"ui": {"text": (text,)}, "result": (text,)}
|
||
except Exception as e:
|
||
print(f"[BoosAPI] LLM chat failed: {e}")
|
||
raise
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
# Node: BoosAPITextToImage — 文生图 (ComfyUI 原生格式)
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
|
||
class BoosAPITextToImage:
|
||
"""输入文本提示词，通过 BoosAPI 生成图像，输出 ComfyUI IMAGE 格式"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"positive": ("STRING", {"multiline": True, "default": ""}),
|
||
"model": (model_choices("image"), {"default": default_model("image")}),
|
||
},
|
||
"optional": {
|
||
"negative": ("STRING", {"multiline": True, "default": ""}),
|
||
"size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "generate"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "BoosAPI"
|
||
|
||
def generate(self, positive, model, negative="", size="1024x1024"):
|
||
# Build prompt with negative guidance
|
||
prompt = positive
|
||
if negative:
|
||
prompt = f"{positive}\n避免: {negative}"
|
||
|
||
url = f"{BOOSAPI_V1}/images/generations"
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": size,
|
||
"response_format": "b64_json",
|
||
}
|
||
try:
|
||
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
|
||
raise_for_status_with_body(resp)
|
||
data = resp.json()
|
||
|
||
import base64
|
||
from io import BytesIO
|
||
import torch
|
||
from PIL import Image
|
||
import numpy as np
|
||
|
||
b64 = data["data"][0]["b64_json"]
|
||
img = Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")
|
||
tensor = torch.from_numpy(np.array(img).astype("float32") / 255.0).unsqueeze(0)
|
||
return (tensor,)
|
||
except Exception as e:
|
||
print(f"[BoosAPI] T2I failed: {e}")
|
||
raise
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
# Register Nodes
|
||
# ═════════════════════════════════════════════════════════════════════════
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"BoosAPIImageGen": BoosAPIImageGen,
|
||
"BoosAPILLM": BoosAPILLM,
|
||
"BoosAPITextToImage": BoosAPITextToImage,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"BoosAPIImageGen": "BoosAPI 图像生成",
|
||
"BoosAPILLM": "BoosAPI LLM 对话",
|
||
"BoosAPITextToImage": "BoosAPI 文生图",
|
||
}
|
||
|
||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|