376 lines
11 KiB
Python
376 lines
11 KiB
Python
import base64
|
|
import io
|
|
from typing import Dict, Any, List, Tuple, Union
|
|
|
|
import torch
|
|
from PIL import Image
|
|
import requests
|
|
from requests.adapters import HTTPAdapter
|
|
from requests.packages.urllib3.util.retry import Retry
|
|
|
|
|
|
class Image2Prompt:
|
|
"""
|
|
ComfyUI 自定义节点：将图片转换为文本提示
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"APIUrl": ("STRING", {"default": "", "multiline": False}),
|
|
"APIKey": ("STRING", {"default": "", "multiline": False}),
|
|
"RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}),
|
|
"Neutral": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "NudeJS"
|
|
|
|
def __init__(self):
|
|
self.session = requests.Session()
|
|
# 配置重试策略
|
|
retry_strategy = Retry(
|
|
total=3,
|
|
backoff_factor=1,
|
|
status_forcelist=[429, 500, 502, 503, 504],
|
|
)
|
|
adapter = HTTPAdapter(max_retries=retry_strategy)
|
|
self.session.mount("http://", adapter)
|
|
self.session.mount("https://", adapter)
|
|
|
|
def execute(
|
|
self,
|
|
images: torch.Tensor,
|
|
APIUrl: str,
|
|
APIKey: str,
|
|
RetryCount: int,
|
|
Neutral: bool,
|
|
) -> Tuple[Union[str, None]]:
|
|
"""
|
|
执行图片转提示词的主逻辑
|
|
|
|
Args:
|
|
images: 输入的图片张量
|
|
APIUrl: API 接口地址
|
|
APIKey: API 密钥
|
|
RetryCount: 重试次数
|
|
Neutral: 是否使用中性文本
|
|
|
|
Returns:
|
|
转换后的文本提示
|
|
"""
|
|
if len(images) == 0:
|
|
return (None,)
|
|
|
|
# 获取第一张图片（ComfyUI 通常批量处理，这里取第一张）
|
|
img_tensor = images[0]
|
|
|
|
# 转换为 PIL Image
|
|
pil_image = self._tensor_to_pil(img_tensor)
|
|
|
|
# 检查并缩放图片
|
|
resized_image = self._resize_image(pil_image)
|
|
|
|
# 转换为 base64
|
|
img_base64 = self._image_to_base64(resized_image)
|
|
|
|
# 发送请求获取提示词
|
|
prompt_text = self._call_api(
|
|
APIUrl=APIUrl,
|
|
APIKey=APIKey,
|
|
imgBase64=img_base64,
|
|
retry_count=RetryCount,
|
|
neutral=Neutral,
|
|
)
|
|
|
|
return (prompt_text,)
|
|
|
|
def _tensor_to_pil(self, img_tensor: torch.Tensor) -> Image.Image:
|
|
"""
|
|
将 ComfyUI 的张量转换为 PIL Image
|
|
|
|
Args:
|
|
img_tensor: ComfyUI 中的图片张量 (B, H, W, C) 或 (H, W, C)
|
|
|
|
Returns:
|
|
PIL Image
|
|
"""
|
|
# 确保张量在 CPU 上
|
|
if img_tensor.is_cuda:
|
|
img_tensor = img_tensor.cpu()
|
|
|
|
# 处理不同的张量形状
|
|
if len(img_tensor.shape) == 4:
|
|
# 批量图片，取第一张
|
|
img_tensor = img_tensor[0]
|
|
elif len(img_tensor.shape) == 3:
|
|
# 单张图片
|
|
pass
|
|
else:
|
|
raise ValueError(f"Unexpected tensor shape: {img_tensor.shape}")
|
|
|
|
# 转换为 PIL Image
|
|
# ComfyUI 的图片张量通常值范围是 [0, 1]，形状是 [H, W, C]
|
|
if img_tensor.shape[-1] in [3, 4]: # RGB 或 RGBA
|
|
array = (img_tensor.numpy() * 255).astype("uint8")
|
|
pil_image = Image.fromarray(array, mode="RGB")
|
|
else:
|
|
raise ValueError(f"Unexpected number of channels: {img_tensor.shape[-1]}")
|
|
|
|
return pil_image
|
|
|
|
def _resize_image(self, image: Image.Image) -> Image.Image:
|
|
"""
|
|
缩放图片，确保宽度不超过 1000，高度不超过 800
|
|
|
|
Args:
|
|
image: 输入的 PIL 图片
|
|
|
|
Returns:
|
|
缩放后的图片
|
|
"""
|
|
width, height = image.size
|
|
new_width = width
|
|
new_height = height
|
|
|
|
# 检查宽度
|
|
if width > 1000:
|
|
new_width = 1000
|
|
new_height = int((height * 1000) / width)
|
|
|
|
# 检查高度
|
|
if new_height > 800:
|
|
new_height = 800
|
|
new_width = int((new_width * 800) / new_height)
|
|
|
|
# 如果需要缩放，则调整大小
|
|
if new_width != width or new_height != height:
|
|
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
|
|
return image
|
|
|
|
def _image_to_base64(self, image: Image.Image) -> str:
|
|
"""
|
|
将 PIL 图片转换为 base64 字符串（不包含 data:image/png;base64, 前缀）
|
|
|
|
Args:
|
|
image: 输入的 PIL 图片
|
|
|
|
Returns:
|
|
base64 编码的字符串
|
|
"""
|
|
# 转换为 RGB 模式（如果需要）
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
|
|
# 保存到内存中的字节流
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
img_bytes = buffer.getvalue()
|
|
|
|
# 编码为 base64
|
|
img_base64 = base64.b64encode(img_bytes).decode("utf-8")
|
|
|
|
return img_base64
|
|
|
|
def _call_api(
|
|
self,
|
|
APIUrl: str,
|
|
APIKey: str,
|
|
imgBase64: str,
|
|
retry_count: int,
|
|
neutral: bool,
|
|
) -> Union[str, None]:
|
|
"""
|
|
调用 API 接口获取提示词
|
|
|
|
Args:
|
|
APIUrl: API 地址
|
|
APIKey: API 密钥
|
|
imgBase64: base64 编码的图片
|
|
retry_count: 最大重试次数
|
|
neutral: 是否使用中性文本
|
|
|
|
Returns:
|
|
提示词文本，失败返回 None
|
|
"""
|
|
url = f"{APIUrl}/generate"
|
|
headers = {"authorization": APIKey}
|
|
data = { "type": "img2img", "base64": imgBase64 }
|
|
|
|
max_retries = max(0, retry_count)
|
|
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
response = self.session.post(url, json=data, headers=headers, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
|
|
if result.get("code") == 200:
|
|
if neutral:
|
|
return result.get("neutralText")
|
|
else:
|
|
return result.get("promptText")
|
|
else:
|
|
# API 返回错误码，记录日志
|
|
print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}")
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Attempt {attempt + 1} failed: {str(e)}")
|
|
|
|
except Exception as e:
|
|
print(f"Unexpected error on attempt {attempt + 1}: {str(e)}")
|
|
|
|
# 如果不是最后一次尝试，等待一下再重试
|
|
if attempt < max_retries:
|
|
import time
|
|
time.sleep(1)
|
|
|
|
# 所有重试都失败了
|
|
print(f"Failed after {max_retries} retries")
|
|
return None
|
|
|
|
|
|
class Text2Prompt:
|
|
"""
|
|
ComfyUI 自定义节点：将文本描述转换为图片提示
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"imgDesc": ("STRING", {"default": "", "multiline": True}),
|
|
"APIUrl": ("STRING", {"default": "", "multiline": False}),
|
|
"APIKey": ("STRING", {"default": "", "multiline": False}),
|
|
"RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}),
|
|
"Neutral": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "NudeJS"
|
|
|
|
def __init__(self):
|
|
self.session = requests.Session()
|
|
# 配置重试策略
|
|
retry_strategy = Retry(
|
|
total=3,
|
|
backoff_factor=1,
|
|
status_forcelist=[429, 500, 502, 503, 504],
|
|
)
|
|
adapter = HTTPAdapter(max_retries=retry_strategy)
|
|
self.session.mount("http://", adapter)
|
|
self.session.mount("https://", adapter)
|
|
|
|
def execute(
|
|
self,
|
|
imgDesc: str,
|
|
APIUrl: str,
|
|
APIKey: str,
|
|
RetryCount: int,
|
|
Neutral: bool,
|
|
) -> Tuple[Union[str, None]]:
|
|
"""
|
|
执行文本转提示词的主逻辑
|
|
|
|
Args:
|
|
imgDesc: 图片的文本描述
|
|
APIUrl: API 接口地址
|
|
APIKey: API 密钥
|
|
RetryCount: 重试次数
|
|
Neutral: 是否使用中性文本
|
|
|
|
Returns:
|
|
转换后的文本提示
|
|
"""
|
|
|
|
# 发送请求获取提示词
|
|
prompt_text = self._call_api(
|
|
APIUrl=APIUrl,
|
|
APIKey=APIKey,
|
|
imgDesc=imgDesc,
|
|
retry_count=RetryCount,
|
|
neutral=Neutral,
|
|
)
|
|
|
|
return (prompt_text,)
|
|
|
|
def _call_api(
|
|
self,
|
|
APIUrl: str,
|
|
APIKey: str,
|
|
imgDesc: str,
|
|
retry_count: int,
|
|
neutral: bool,
|
|
) -> Union[str, None]:
|
|
"""
|
|
调用 API 接口获取提示词
|
|
|
|
Args:
|
|
APIUrl: API 地址
|
|
APIKey: API 密钥
|
|
imgDesc: 图片的文本描述
|
|
retry_count: 最大重试次数
|
|
neutral: 是否使用中性文本
|
|
|
|
Returns:
|
|
提示词文本，失败返回 None
|
|
"""
|
|
url = f"{APIUrl}/generate"
|
|
headers = {"authorization": APIKey}
|
|
data = { "type": "text2img", "desc": imgDesc }
|
|
|
|
max_retries = max(0, retry_count)
|
|
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
response = self.session.post(url, json=data, headers=headers, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
|
|
if result.get("code") == 200:
|
|
if neutral:
|
|
return result.get("neutralText")
|
|
else:
|
|
return result.get("promptText")
|
|
else:
|
|
# API 返回错误码，记录日志
|
|
print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}")
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Attempt {attempt + 1} failed: {str(e)}")
|
|
|
|
except Exception as e:
|
|
print(f"Unexpected error on attempt {attempt + 1}: {str(e)}")
|
|
|
|
# 如果不是最后一次尝试，等待一下再重试
|
|
if attempt < max_retries:
|
|
import time
|
|
time.sleep(1)
|
|
|
|
# 所有重试都失败了
|
|
print(f"Failed after {max_retries} retries")
|
|
return None
|
|
|
|
|
|
# 注册节点
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Image2Prompt": Image2Prompt,
|
|
"Text2Prompt": Text2Prompt,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Image2Prompt": "Image2Prompt",
|
|
"Text2Prompt": "Text2Prompt",
|
|
}
|