feat: 新增启动欢迎通知、流式预览节点及多项功能更新
- 新增启动弹窗通知(绿色主题,支持关闭) - 新增 StreamPreview 流式文本预览节点 - 新增 fileUpload、updateNotifier 前端 JS 模块 - 重构多个 client,统一错误处理 - 删除废弃节点 batch_nano_banana_v2、quan_neng_sheng_tu 等 - 将 .config 纳入版本控制(已清空密钥)
This commit is contained in:
@@ -1,7 +1,3 @@
|
|||||||
# .config 文件包含敏感信息,不提交到版本控制
|
|
||||||
# 用户可通过 setup_api_key.bat 自动创建本地配置
|
|
||||||
.config
|
|
||||||
|
|
||||||
# Python 缓存
|
# Python 缓存
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.py[cod]
|
*.py[cod]
|
||||||
|
|||||||
+22
-11
@@ -12,16 +12,15 @@ Comfyui_o1key - ComfyUI 自定义节点集合
|
|||||||
# 检查更新(仅在启动时检查一次)
|
# 检查更新(仅在启动时检查一次)
|
||||||
try:
|
try:
|
||||||
from .utils.update_checker import check_for_updates, notify_update_available
|
from .utils.update_checker import check_for_updates, notify_update_available
|
||||||
|
|
||||||
if check_for_updates():
|
notify_update_available() # TODO: 测试用,改回 if check_for_updates(): notify_update_available()
|
||||||
notify_update_available()
|
|
||||||
except Exception:
|
except Exception:
|
||||||
# 静默失败,不影响插件加载
|
# 静默失败,不影响插件加载
|
||||||
pass
|
pass
|
||||||
|
|
||||||
import ssl
|
import ssl
|
||||||
|
|
||||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, QuanNengShengTu, BatchQuanNengShengTu, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance
|
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview
|
||||||
|
|
||||||
# 报错弹框友好文案(不修改原节点代码,仅在外层统一处理)
|
# 报错弹框友好文案(不修改原节点代码,仅在外层统一处理)
|
||||||
_MSG_TIMEOUT = "API 请求超时,请稍后重试或检查网络。"
|
_MSG_TIMEOUT = "API 请求超时,请稍后重试或检查网络。"
|
||||||
@@ -55,8 +54,6 @@ def _wrap_generate_for_error_display(cls, attr="generate"):
|
|||||||
|
|
||||||
_wrap_generate_for_error_display(NanoBananaPro)
|
_wrap_generate_for_error_display(NanoBananaPro)
|
||||||
_wrap_generate_for_error_display(BatchNanoBananaPro)
|
_wrap_generate_for_error_display(BatchNanoBananaPro)
|
||||||
_wrap_generate_for_error_display(QuanNengShengTu)
|
|
||||||
_wrap_generate_for_error_display(BatchQuanNengShengTu, "process_batch")
|
|
||||||
|
|
||||||
# ComfyUI 节点注册
|
# ComfyUI 节点注册
|
||||||
NODE_CLASS_MAPPINGS = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
@@ -74,12 +71,12 @@ NODE_CLASS_MAPPINGS = {
|
|||||||
"KlingVideo": KlingVideo,
|
"KlingVideo": KlingVideo,
|
||||||
"KlingFirstLastFrame": KlingFirstLastFrame,
|
"KlingFirstLastFrame": KlingFirstLastFrame,
|
||||||
"KlingMotionControlTest": KlingMotionControlTest,
|
"KlingMotionControlTest": KlingMotionControlTest,
|
||||||
"QuanNengShengTu": QuanNengShengTu,
|
|
||||||
"BatchQuanNengShengTu": BatchQuanNengShengTu,
|
|
||||||
"AspectRatioPreset": AspectRatioPreset,
|
"AspectRatioPreset": AspectRatioPreset,
|
||||||
"MultiResPreview": MultiResPreview,
|
"MultiResPreview": MultiResPreview,
|
||||||
"BatchImagesO1key": BatchImagesO1key,
|
"BatchImagesO1key": BatchImagesO1key,
|
||||||
"Seedance": Seedance,
|
"Seedance": Seedance,
|
||||||
|
"SeedanceMultiModal": SeedanceMultiModal,
|
||||||
|
"StreamPreview": StreamPreview,
|
||||||
}
|
}
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
@@ -90,21 +87,35 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
|||||||
"ImageStitchPro": "图像拼接 Pro",
|
"ImageStitchPro": "图像拼接 Pro",
|
||||||
"SaveCleanImage": "保存图像(防AI识别)",
|
"SaveCleanImage": "保存图像(防AI识别)",
|
||||||
"BatchCleanMetadata": "批量任务(防AI识别)",
|
"BatchCleanMetadata": "批量任务(防AI识别)",
|
||||||
"VideoPreview": "视频预览",
|
"VideoPreview": "预览视频",
|
||||||
"GoogleVeo": "Google Veo - ab",
|
"GoogleVeo": "Google Veo - ab",
|
||||||
"FluxImageEdit": "Flux2 图像编辑",
|
"FluxImageEdit": "Flux2 图像编辑",
|
||||||
"UniversalLLMChat": "全能LLM对话助手",
|
"UniversalLLMChat": "全能LLM对话助手",
|
||||||
"KlingVideo": "文/图生视频 自研模型",
|
"KlingVideo": "文/图生视频 自研模型",
|
||||||
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
|
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
|
||||||
"KlingMotionControlTest": "动作控制 自研模型",
|
"KlingMotionControlTest": "动作控制 自研模型",
|
||||||
"QuanNengShengTu": "全能生图",
|
|
||||||
"BatchQuanNengShengTu": "全能生图(批量)",
|
|
||||||
"AspectRatioPreset": "图片宽高比预设",
|
"AspectRatioPreset": "图片宽高比预设",
|
||||||
"MultiResPreview": "预览图像(v2)",
|
"MultiResPreview": "预览图像(v2)",
|
||||||
"BatchImagesO1key": "加载图像(批量)",
|
"BatchImagesO1key": "加载图像(批量)",
|
||||||
"Seedance": "Seedance 视频生成",
|
"Seedance": "Seedance 视频生成",
|
||||||
|
"SeedanceMultiModal": "Seedance 多模态参考生视频",
|
||||||
|
"StreamPreview": "流式文本预览",
|
||||||
}
|
}
|
||||||
|
|
||||||
WEB_DIRECTORY = "./web"
|
WEB_DIRECTORY = "./web"
|
||||||
|
|
||||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
|
||||||
|
|
||||||
|
# 注册 /o1key/input_dir 接口,供前端文件上传按钮获取 input 目录绝对路径
|
||||||
|
try:
|
||||||
|
from aiohttp import web
|
||||||
|
from server import PromptServer
|
||||||
|
import folder_paths
|
||||||
|
|
||||||
|
@PromptServer.instance.routes.get("/o1key/input_dir")
|
||||||
|
async def get_input_dir(request):
|
||||||
|
import os
|
||||||
|
path = os.path.abspath(folder_paths.get_input_directory())
|
||||||
|
return web.json_response({"path": path})
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|||||||
+100
-28
@@ -79,6 +79,15 @@ class BaseAPIClient(ABC):
|
|||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
def _make_session(self) -> aiohttp.ClientSession:
|
||||||
|
"""
|
||||||
|
创建统一的 aiohttp ClientSession,全局禁用 SSL 验证。
|
||||||
|
所有需要独立创建 session 的地方都应调用此方法,
|
||||||
|
避免因客户端系统缺少根证书导致 SSLCertVerificationError。
|
||||||
|
"""
|
||||||
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
return aiohttp.ClientSession(connector=connector)
|
||||||
|
|
||||||
def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]:
|
def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]:
|
||||||
"""
|
"""
|
||||||
获取请求头
|
获取请求头
|
||||||
@@ -142,67 +151,132 @@ class BaseAPIClient(ABC):
|
|||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
发送异步 HTTP 请求(带详细计时)
|
发送异步 HTTP 请求(带详细计时)
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
endpoint: API 端点
|
endpoint: API 端点
|
||||||
request_body: 请求体
|
request_body: 请求体
|
||||||
session: aiohttp 会话(可选)
|
session: aiohttp 会话(可选)
|
||||||
use_bearer_token: 是否使用 Bearer Token 认证
|
use_bearer_token: 是否使用 Bearer Token 认证
|
||||||
timeout: 超时时间(秒)- 已废弃,由服务器端控制
|
timeout: 超时时间(秒),默认 900 秒
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
响应 JSON
|
响应 JSON
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
RuntimeError: 请求失败时
|
RuntimeError: 请求失败时
|
||||||
|
InterruptProcessingException: 用户点击终止按钮时
|
||||||
"""
|
"""
|
||||||
import time
|
import time
|
||||||
|
|
||||||
|
# 尝试导入 ComfyUI 中断机制
|
||||||
|
try:
|
||||||
|
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||||
|
_interrupt_available = True
|
||||||
|
except ImportError:
|
||||||
|
_interrupt_available = False
|
||||||
|
|
||||||
url = f"{self.base_url}{endpoint}"
|
url = f"{self.base_url}{endpoint}"
|
||||||
headers = self.get_headers(use_bearer_token)
|
headers = self.get_headers(use_bearer_token)
|
||||||
|
|
||||||
# 检查请求大小
|
# 检查请求大小
|
||||||
self.check_request_size(request_body)
|
self.check_request_size(request_body)
|
||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
# 设置请求超时:连接超时 30s,读取超时 900s(防止服务器出图后卡住)
|
||||||
# 连接计时
|
_timeout_seconds = timeout if timeout is not None else 900
|
||||||
|
_aiohttp_timeout = aiohttp.ClientTimeout(
|
||||||
|
total=_timeout_seconds,
|
||||||
|
connect=30,
|
||||||
|
sock_read=_timeout_seconds
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _do_request():
|
||||||
connect_start = time.time()
|
connect_start = time.time()
|
||||||
|
async with session.post(url, json=request_body, headers=headers, timeout=_aiohttp_timeout) as response:
|
||||||
async with session.post(url, json=request_body, headers=headers) as response:
|
|
||||||
connect_time = time.time() - connect_start
|
connect_time = time.time() - connect_start
|
||||||
|
|
||||||
if response.status != 200:
|
if response.status != 200:
|
||||||
error_text = await response.text()
|
error_text = await response.text()
|
||||||
raise RuntimeError(error_text)
|
raise RuntimeError(error_text)
|
||||||
|
|
||||||
# 接收响应体
|
|
||||||
wait_start = time.time()
|
wait_start = time.time()
|
||||||
response_data = await response.json()
|
response_data = await response.json()
|
||||||
download_time = time.time() - wait_start
|
download_time = time.time() - wait_start
|
||||||
|
|
||||||
# 附加计时信息到响应数据(供上层使用)
|
|
||||||
response_size = len(str(response_data))
|
response_size = len(str(response_data))
|
||||||
if not isinstance(response_data, dict):
|
if not isinstance(response_data, dict):
|
||||||
response_data = {"data": response_data}
|
response_data = {"data": response_data}
|
||||||
|
|
||||||
# 将计时信息存储在响应的元数据中
|
|
||||||
response_data["_timing"] = {
|
response_data["_timing"] = {
|
||||||
"connect_time": connect_time,
|
"connect_time": connect_time,
|
||||||
"download_time": download_time,
|
"download_time": download_time,
|
||||||
"response_size": response_size
|
"response_size": response_size
|
||||||
}
|
}
|
||||||
|
|
||||||
return response_data
|
return response_data
|
||||||
|
|
||||||
|
async def _poll_interrupt():
|
||||||
|
"""每 0.5s 轮询一次中断标志"""
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(0.5)
|
||||||
|
if processing_interrupted():
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
if _interrupt_available:
|
||||||
|
request_task = asyncio.ensure_future(_do_request())
|
||||||
|
interrupt_task = asyncio.ensure_future(_poll_interrupt())
|
||||||
|
|
||||||
|
done, pending = await asyncio.wait(
|
||||||
|
[request_task, interrupt_task],
|
||||||
|
return_when=asyncio.FIRST_COMPLETED
|
||||||
|
)
|
||||||
|
|
||||||
|
# 取消未完成的任务
|
||||||
|
for t in pending:
|
||||||
|
t.cancel()
|
||||||
|
try:
|
||||||
|
await t
|
||||||
|
except (asyncio.CancelledError, Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 判断是哪个先完成
|
||||||
|
if interrupt_task in done and request_task not in done:
|
||||||
|
raise InterruptProcessingException()
|
||||||
|
|
||||||
|
# 请求完成,取出结果(可能含异常)
|
||||||
|
return request_task.result()
|
||||||
|
else:
|
||||||
|
return await _do_request()
|
||||||
|
|
||||||
|
except InterruptProcessingException:
|
||||||
|
raise
|
||||||
|
|
||||||
|
except aiohttp.ServerTimeoutError as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"请求超时!等待服务器响应超过 {_timeout_seconds} 秒。\n"
|
||||||
|
f"服务器可能仍在生成图片,请稍后重试,或检查网络连接。"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
except aiohttp.ClientConnectorError as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"无法连接到服务器:{str(e)}\n"
|
||||||
|
f"请检查网络连接是否正常。"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
except asyncio.TimeoutError as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"请求超时!等待服务器响应超过 {_timeout_seconds} 秒。\n"
|
||||||
|
f"服务器可能仍在生成图片,请稍后重试,或检查网络连接。"
|
||||||
|
) from e
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
if close_session:
|
if close_session:
|
||||||
await session.close()
|
await session.close()
|
||||||
|
|
||||||
async def request_get_async(
|
async def request_get_async(
|
||||||
self,
|
self,
|
||||||
endpoint: str,
|
endpoint: str,
|
||||||
@@ -230,9 +304,9 @@ class BaseAPIClient(ABC):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
async with session.get(url, headers=headers) as response:
|
async with session.get(url, headers=headers) as response:
|
||||||
if response.status != 200:
|
if response.status != 200:
|
||||||
@@ -327,9 +401,7 @@ class BaseAPIClient(ABC):
|
|||||||
total = len(requests)
|
total = len(requests)
|
||||||
|
|
||||||
# 创建无限制的连接器
|
# 创建无限制的连接器
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
async with self._make_session() as session:
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
|
||||||
tasks = []
|
tasks = []
|
||||||
|
|
||||||
for req in requests:
|
for req in requests:
|
||||||
|
|||||||
+27
-67
@@ -141,7 +141,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
resolution: str = "2K",
|
resolution: str = "2K",
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
candidate_count: int = 1,
|
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
@@ -154,7 +153,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
resolution: 分辨率
|
resolution: 分辨率
|
||||||
enable_grounding: 是否启用 Google Search Grounding
|
enable_grounding: 是否启用 Google Search Grounding
|
||||||
enable_image_search: 是否同时启用 Google Image Search(仅 Gemini 3.1 Flash 支持)
|
enable_image_search: 是否同时启用 Google Image Search(仅 Gemini 3.1 Flash 支持)
|
||||||
candidate_count: 单次请求返回的候选图数量,默认 1
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
请求体字典
|
请求体字典
|
||||||
@@ -184,7 +182,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"generationConfig": {
|
"generationConfig": {
|
||||||
"candidateCount": candidate_count,
|
|
||||||
"responseModalities": ["TEXT", "IMAGE"],
|
"responseModalities": ["TEXT", "IMAGE"],
|
||||||
"imageConfig": {
|
"imageConfig": {
|
||||||
"aspectRatio": aspect_ratio,
|
"aspectRatio": aspect_ratio,
|
||||||
@@ -303,7 +300,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
# 需要关闭 session 的标记
|
# 需要关闭 session 的标记
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -438,7 +435,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request: bool = False,
|
debug_request: bool = False,
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
candidate_count: int = 1
|
|
||||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||||
"""
|
"""
|
||||||
单次异步生成请求(极简单行日志)
|
单次异步生成请求(极简单行日志)
|
||||||
@@ -465,7 +461,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
total_start = time.time()
|
total_start = time.time()
|
||||||
|
|
||||||
# 任务前缀
|
# 任务前缀
|
||||||
task_prefix = f"[{task_index}/{total_tasks}]" if task_index is not None and total_tasks else ""
|
task_prefix = f"[{task_index}/{total_tasks}] " if task_index is not None and total_tasks else ""
|
||||||
|
|
||||||
# ========== 1. 构建请求 ==========
|
# ========== 1. 构建请求 ==========
|
||||||
build_start = time.time()
|
build_start = time.time()
|
||||||
@@ -477,7 +473,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
resolution=resolution,
|
resolution=resolution,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
candidate_count=candidate_count
|
|
||||||
)
|
)
|
||||||
build_time = time.time() - build_start
|
build_time = time.time() - build_start
|
||||||
|
|
||||||
@@ -499,7 +494,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
safe_request = _truncate_base64_req(request_body)
|
safe_request = _truncate_base64_req(request_body)
|
||||||
print(
|
print(
|
||||||
f"\n{'='*60}\n"
|
f"\n{'='*60}\n"
|
||||||
f"[请求体日志] 任务 {task_prefix or '?'} 发送请求体:\n"
|
f"[请求体日志] {task_prefix}发送请求体:\n"
|
||||||
f"端点: {endpoint}\n"
|
f"端点: {endpoint}\n"
|
||||||
f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n"
|
f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n"
|
||||||
f"{'='*60}\n"
|
f"{'='*60}\n"
|
||||||
@@ -541,7 +536,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
safe_response = _truncate_base64(response)
|
safe_response = _truncate_base64(response)
|
||||||
print(
|
print(
|
||||||
f"\n{'='*60}\n"
|
f"\n{'='*60}\n"
|
||||||
f"[调试日志] 任务 {task_prefix or '?'} 完整 API 响应:\n"
|
f"[调试日志] {task_prefix}完整 API 响应:\n"
|
||||||
f"{_json.dumps(safe_response, ensure_ascii=False, indent=2)}\n"
|
f"{_json.dumps(safe_response, ensure_ascii=False, indent=2)}\n"
|
||||||
f"{'='*60}\n"
|
f"{'='*60}\n"
|
||||||
)
|
)
|
||||||
@@ -554,7 +549,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
parse_time = time.time() - parse_start
|
parse_time = time.time() - parse_start
|
||||||
error_first_line = str(e).split('\n')[0]
|
error_first_line = str(e).split('\n')[0]
|
||||||
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗")
|
print(f"{task_prefix}请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗")
|
||||||
raise
|
raise
|
||||||
|
|
||||||
parse_time = time.time() - parse_start
|
parse_time = time.time() - parse_start
|
||||||
@@ -578,7 +573,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
download_info = f"{img_size_str}"
|
download_info = f"{img_size_str}"
|
||||||
|
|
||||||
# 单行输出
|
# 单行输出
|
||||||
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓")
|
print(f"{task_prefix}请求 {size_str} → API {request_time:.1f}s → {download_info} ✓")
|
||||||
|
|
||||||
# 返回结果和计时信息
|
# 返回结果和计时信息
|
||||||
total_time = time.time() - total_start
|
total_time = time.time() - total_start
|
||||||
@@ -605,7 +600,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request: bool = False,
|
debug_request: bool = False,
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
candidate_count: int = 1
|
|
||||||
) -> List[Image.Image]:
|
) -> List[Image.Image]:
|
||||||
"""
|
"""
|
||||||
批量全并发生成 - 改进版:支持分批处理和内存管理
|
批量全并发生成 - 改进版:支持分批处理和内存管理
|
||||||
@@ -622,7 +616,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request: 是否打印发送的请求体
|
debug_request: 是否打印发送的请求体
|
||||||
enable_grounding: 是否启用 Google Search Grounding
|
enable_grounding: 是否启用 Google Search Grounding
|
||||||
enable_image_search: 是否同时启用 Google Image Search
|
enable_image_search: 是否同时启用 Google Image Search
|
||||||
candidate_count: 单次请求返回的候选图数量
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
生成的图像列表
|
生成的图像列表
|
||||||
@@ -635,28 +628,23 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
success_count = 0
|
success_count = 0
|
||||||
fail_count = 0
|
fail_count = 0
|
||||||
first_error = None # 保存第一个错误
|
first_error = None # 保存第一个错误
|
||||||
|
|
||||||
# 分批处理配置
|
# 分批处理配置
|
||||||
max_concurrent = 10 # 最大并发数
|
max_concurrent = 10 # 最大并发数
|
||||||
save_batch_size = 10 # 分批保存大小
|
save_batch_size = 10 # 分批保存大小
|
||||||
|
|
||||||
# 计算需要多少批次
|
# 计算需要多少批次
|
||||||
num_batches = (batch_size + max_concurrent - 1) // max_concurrent
|
num_batches = (batch_size + max_concurrent - 1) // max_concurrent
|
||||||
|
|
||||||
print(f"GeminiClient: 批量生成 {batch_size} 张图片,并发数: {max_concurrent},分 {num_batches} 批执行")
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
# 分批执行
|
# 分批执行
|
||||||
for batch_idx in range(num_batches):
|
for batch_idx in range(num_batches):
|
||||||
batch_start = batch_idx * max_concurrent
|
batch_start = batch_idx * max_concurrent
|
||||||
batch_end = min(batch_start + max_concurrent, batch_size)
|
batch_end = min(batch_start + max_concurrent, batch_size)
|
||||||
batch_size_current = batch_end - batch_start
|
batch_size_current = batch_end - batch_start
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...")
|
|
||||||
|
|
||||||
# 创建当前批次的任务
|
# 创建当前批次的任务
|
||||||
tasks = []
|
tasks = []
|
||||||
for i in range(batch_size_current):
|
for i in range(batch_size_current):
|
||||||
@@ -675,7 +663,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request=debug_request,
|
debug_request=debug_request,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
candidate_count=candidate_count
|
|
||||||
),
|
),
|
||||||
name=f"task_{task_index}"
|
name=f"task_{task_index}"
|
||||||
)
|
)
|
||||||
@@ -698,14 +685,11 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
all_images.append(img)
|
all_images.append(img)
|
||||||
|
|
||||||
success_count += 1
|
success_count += 1
|
||||||
|
|
||||||
# 通知进度
|
# 通知进度
|
||||||
if progress_callback:
|
if progress_callback:
|
||||||
progress_callback(completed, batch_size, True, None)
|
progress_callback(completed, batch_size, True, None)
|
||||||
|
|
||||||
# 每成功生成一张图片就打印日志
|
|
||||||
print(f"GeminiClient: 任务 {completed}/{batch_size} 成功生成图片 ✓")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
fail_count += 1
|
fail_count += 1
|
||||||
# 保存第一个错误(用于后续抛出)
|
# 保存第一个错误(用于后续抛出)
|
||||||
@@ -719,25 +703,20 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
|
|
||||||
# 当前批次完成后,立即清理内存
|
# 当前批次完成后,立即清理内存
|
||||||
if batch_images:
|
if batch_images:
|
||||||
print(f"GeminiClient: 第 {batch_idx + 1} 批完成,生成 {len(batch_images)} 张图片")
|
|
||||||
|
|
||||||
# 强制垃圾回收,释放内存
|
|
||||||
import gc
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
# 短暂暂停,让系统处理内存
|
|
||||||
await asyncio.sleep(0.1)
|
await asyncio.sleep(0.1)
|
||||||
|
|
||||||
# 清空当前批次图片引用,帮助垃圾回收
|
# 清空当前批次图片引用,帮助垃圾回收
|
||||||
batch_images = []
|
batch_images = []
|
||||||
|
|
||||||
# 最终结果检查
|
# 最终结果检查
|
||||||
if not all_images:
|
if not all_images:
|
||||||
# 如果有保存的原始错误,直接抛出原始错误
|
# 如果有保存的原始错误,直接抛出原始错误
|
||||||
if first_error:
|
if first_error:
|
||||||
raise first_error
|
raise first_error
|
||||||
raise RuntimeError(f"批量生成失败,{fail_count} 个请求全部失败")
|
raise RuntimeError(f"批量生成失败,{fail_count} 个请求全部失败")
|
||||||
|
|
||||||
print(f"GeminiClient: 批量生成完成,成功 {success_count}/{batch_size},失败 {fail_count}")
|
print(f"GeminiClient: 批量生成完成,成功 {success_count}/{batch_size},失败 {fail_count}")
|
||||||
return all_images
|
return all_images
|
||||||
|
|
||||||
@@ -754,7 +733,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request: bool = False,
|
debug_request: bool = False,
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
candidate_count: int = 1
|
|
||||||
) -> List[Image.Image]:
|
) -> List[Image.Image]:
|
||||||
"""
|
"""
|
||||||
同步生成接口(用于 ComfyUI)
|
同步生成接口(用于 ComfyUI)
|
||||||
@@ -771,7 +749,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request: 是否打印发送的请求体
|
debug_request: 是否打印发送的请求体
|
||||||
enable_grounding: 是否启用 Google Search Grounding
|
enable_grounding: 是否启用 Google Search Grounding
|
||||||
enable_image_search: 是否同时启用 Google Image Search
|
enable_image_search: 是否同时启用 Google Image Search
|
||||||
candidate_count: 单次请求返回的候选图数量
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
生成的图像列表
|
生成的图像列表
|
||||||
@@ -788,7 +765,6 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
debug_request=debug_request,
|
debug_request=debug_request,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
candidate_count=candidate_count
|
|
||||||
)
|
)
|
||||||
|
|
||||||
return self.run_async_in_thread(coro)
|
return self.run_async_in_thread(coro)
|
||||||
@@ -837,13 +813,11 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
fail_count = 0
|
fail_count = 0
|
||||||
first_error = None # 保存第一个错误
|
first_error = None # 保存第一个错误
|
||||||
total_tasks = len(prompts) * images_per_prompt
|
total_tasks = len(prompts) * images_per_prompt
|
||||||
|
|
||||||
# 分批处理配置
|
# 分批处理配置
|
||||||
max_concurrent = 10 # 最大并发数
|
max_concurrent = 10 # 最大并发数
|
||||||
|
|
||||||
print(f"GeminiClient: 多提示词批量生成,共 {total_tasks} 个任务,{len(prompts)} 个提示词,每个 {images_per_prompt} 张")
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
# 创建所有任务
|
# 创建所有任务
|
||||||
@@ -874,15 +848,12 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
# 分批处理:每10个任务为一组
|
# 分批处理:每10个任务为一组
|
||||||
batch_size = max_concurrent
|
batch_size = max_concurrent
|
||||||
num_batches = (total_tasks + batch_size - 1) // batch_size
|
num_batches = (total_tasks + batch_size - 1) // batch_size
|
||||||
|
|
||||||
for batch_idx in range(num_batches):
|
for batch_idx in range(num_batches):
|
||||||
batch_start = batch_idx * batch_size
|
batch_start = batch_idx * batch_size
|
||||||
batch_end = min(batch_start + batch_size, total_tasks)
|
batch_end = min(batch_start + batch_size, total_tasks)
|
||||||
batch_tasks = tasks[batch_start:batch_end]
|
batch_tasks = tasks[batch_start:batch_end]
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...")
|
|
||||||
|
|
||||||
# 收集当前批次的结果
|
# 收集当前批次的结果
|
||||||
batch_images = []
|
batch_images = []
|
||||||
|
|
||||||
@@ -898,36 +869,26 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
all_images.append(img)
|
all_images.append(img)
|
||||||
|
|
||||||
success_count += 1
|
success_count += 1
|
||||||
|
|
||||||
# 通知进度
|
# 通知进度
|
||||||
if progress_callback:
|
if progress_callback:
|
||||||
progress_callback(completed, total_tasks, True, None)
|
progress_callback(completed, total_tasks, True, None)
|
||||||
|
|
||||||
# 每成功生成一张图片就打印日志
|
|
||||||
print(f"GeminiClient: 任务 {completed}/{total_tasks} 成功生成图片 ✓")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
fail_count += 1
|
fail_count += 1
|
||||||
# 保存第一个错误(用于后续抛出)
|
# 保存第一个错误(用于后续抛出)
|
||||||
if first_error is None:
|
if first_error is None:
|
||||||
first_error = e
|
first_error = e
|
||||||
error_msg = str(e)
|
error_msg = str(e)
|
||||||
|
|
||||||
# 传递完整的错误信息(用于排查问题)
|
# 传递完整的错误信息(用于排查问题)
|
||||||
if progress_callback:
|
if progress_callback:
|
||||||
progress_callback(completed, total_tasks, False, error_msg)
|
progress_callback(completed, total_tasks, False, error_msg)
|
||||||
|
|
||||||
print(f"GeminiClient: 任务 {completed}/{total_tasks} 失败 ✗")
|
|
||||||
|
|
||||||
# 当前批次完成后,立即清理内存
|
# 当前批次完成后,立即清理内存
|
||||||
if batch_images:
|
if batch_images:
|
||||||
print(f"GeminiClient: 第 {batch_idx + 1} 批完成,生成 {len(batch_images)} 张图片")
|
|
||||||
|
|
||||||
# 强制垃圾回收,释放内存
|
|
||||||
import gc
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
# 短暂暂停,让系统处理内存
|
|
||||||
await asyncio.sleep(0.1)
|
await asyncio.sleep(0.1)
|
||||||
|
|
||||||
# 清空当前批次图片引用,帮助垃圾回收
|
# 清空当前批次图片引用,帮助垃圾回收
|
||||||
@@ -938,8 +899,7 @@ class GeminiAPIClient(BaseAPIClient):
|
|||||||
if first_error:
|
if first_error:
|
||||||
raise first_error
|
raise first_error
|
||||||
raise RuntimeError(f"批量生成失败,{fail_count} 个请求全部失败")
|
raise RuntimeError(f"批量生成失败,{fail_count} 个请求全部失败")
|
||||||
|
|
||||||
print(f"GeminiClient: 多提示词批量生成完成,成功 {success_count}/{total_tasks},失败 {fail_count}")
|
|
||||||
return all_images
|
return all_images
|
||||||
|
|
||||||
def generate_multi_prompts_sync(
|
def generate_multi_prompts_sync(
|
||||||
|
|||||||
@@ -137,7 +137,7 @@ class KlingClient:
|
|||||||
on_progress: Optional[Callable[[int], None]] = None,
|
on_progress: Optional[Callable[[int], None]] = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""提交 → 轮询 → 下载,返回本地文件路径"""
|
"""提交 → 轮询 → 下载,返回本地文件路径"""
|
||||||
connector = aiohttp.TCPConnector(force_close=True)
|
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
if on_stage:
|
if on_stage:
|
||||||
on_stage("submitting")
|
on_stage("submitting")
|
||||||
@@ -196,7 +196,7 @@ class KlingClient:
|
|||||||
"Content-Type": "application/json"}
|
"Content-Type": "application/json"}
|
||||||
interval = self.POLL_INITIAL_INTERVAL
|
interval = self.POLL_INITIAL_INTERVAL
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(force_close=True)
|
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
|
|
||||||
# 1. 提交
|
# 1. 提交
|
||||||
|
|||||||
@@ -309,7 +309,7 @@ class OpenAIAPIClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -677,7 +677,7 @@ class OpenAIAPIClient(BaseAPIClient):
|
|||||||
|
|
||||||
print(f"OpenAIClient: 批量生成 {batch_size} 张,并发数: {max_concurrent},分 {num_batches} 批")
|
print(f"OpenAIClient: 批量生成 {batch_size} 张,并发数: {max_concurrent},分 {num_batches} 批")
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
for batch_idx in range(num_batches):
|
for batch_idx in range(num_batches):
|
||||||
|
|||||||
@@ -169,7 +169,7 @@ class SeedanceClient:
|
|||||||
on_progress: Optional[Callable[[int], None]] = None,
|
on_progress: Optional[Callable[[int], None]] = None,
|
||||||
) -> tuple:
|
) -> tuple:
|
||||||
"""提交 → 轮询 → 下载,返回 (本地视频路径, 末帧图片URL或None)"""
|
"""提交 → 轮询 → 下载,返回 (本地视频路径, 末帧图片URL或None)"""
|
||||||
connector = aiohttp.TCPConnector(force_close=True)
|
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
|
|
||||||
# 提交
|
# 提交
|
||||||
|
|||||||
@@ -132,7 +132,7 @@ class SoraClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -174,7 +174,7 @@ class SoraClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
interval = self.POLL_INITIAL_INTERVAL
|
interval = self.POLL_INITIAL_INTERVAL
|
||||||
@@ -242,7 +242,7 @@ class SoraClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -299,7 +299,7 @@ class SoraClient(BaseAPIClient):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
async def _run():
|
async def _run():
|
||||||
connector = aiohttp.TCPConnector(limit=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0)
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
# 1. 提交任务
|
# 1. 提交任务
|
||||||
if on_stage:
|
if on_stage:
|
||||||
@@ -408,7 +408,7 @@ class SoraClient(BaseAPIClient):
|
|||||||
成功生成的视频路径列表
|
成功生成的视频路径列表
|
||||||
"""
|
"""
|
||||||
batch_size = len(save_paths)
|
batch_size = len(save_paths)
|
||||||
connector = aiohttp.TCPConnector(limit=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0)
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
tasks = [
|
tasks = [
|
||||||
|
|||||||
@@ -160,7 +160,7 @@ class VeoClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -202,7 +202,7 @@ class VeoClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
interval = self.POLL_INITIAL_INTERVAL
|
interval = self.POLL_INITIAL_INTERVAL
|
||||||
@@ -265,7 +265,7 @@ class VeoClient(BaseAPIClient):
|
|||||||
|
|
||||||
close_session = False
|
close_session = False
|
||||||
if session is None:
|
if session is None:
|
||||||
session = aiohttp.ClientSession()
|
session = self._make_session()
|
||||||
close_session = True
|
close_session = True
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -318,7 +318,7 @@ class VeoClient(BaseAPIClient):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
async def _run():
|
async def _run():
|
||||||
connector = aiohttp.TCPConnector(limit=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0)
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
# 1. 提交任务
|
# 1. 提交任务
|
||||||
if on_stage:
|
if on_stage:
|
||||||
@@ -383,7 +383,7 @@ class VeoClient(BaseAPIClient):
|
|||||||
"""
|
"""
|
||||||
async def _run():
|
async def _run():
|
||||||
batch_size = len(save_paths)
|
batch_size = len(save_paths)
|
||||||
connector = aiohttp.TCPConnector(limit=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0)
|
||||||
|
|
||||||
async def generate_one(save_path: str):
|
async def generate_one(save_path: str):
|
||||||
return await self._generate_one_video_async(
|
return await self._generate_one_video_async(
|
||||||
|
|||||||
+3
-6
@@ -3,6 +3,7 @@
|
|||||||
包含所有 ComfyUI 自定义节点的实现
|
包含所有 ComfyUI 自定义节点的实现
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from .stream_preview import StreamPreview
|
||||||
from .nano_banana_pro import NanoBananaPro
|
from .nano_banana_pro import NanoBananaPro
|
||||||
from .batch_nano_banana_pro import BatchNanoBananaPro
|
from .batch_nano_banana_pro import BatchNanoBananaPro
|
||||||
from .google_gemini import GoogleGemini
|
from .google_gemini import GoogleGemini
|
||||||
@@ -14,12 +15,8 @@ from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest
|
|||||||
from .veo_video import GoogleVeo
|
from .veo_video import GoogleVeo
|
||||||
from .flux_edit import FluxImageEdit
|
from .flux_edit import FluxImageEdit
|
||||||
from .universal_llm import UniversalLLMChat
|
from .universal_llm import UniversalLLMChat
|
||||||
from .quan_neng_sheng_tu import QuanNengShengTu
|
|
||||||
from .batch_quan_neng_sheng_tu import BatchQuanNengShengTu
|
|
||||||
from .multi_res_preview import MultiResPreview
|
from .multi_res_preview import MultiResPreview
|
||||||
from .batch_images_o1key import BatchImagesO1key
|
from .batch_images_o1key import BatchImagesO1key
|
||||||
from .nano_banana_v2 import NanaBananaV2
|
from .seedance_video import Seedance, SeedanceMultiModal
|
||||||
from .batch_nano_banana_v2 import BatchNanaBananaV2
|
|
||||||
from .seedance_video import Seedance
|
|
||||||
|
|
||||||
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'QuanNengShengTu', 'BatchQuanNengShengTu', 'MultiResPreview', 'BatchImagesO1key', 'NanaBananaV2', 'BatchNanaBananaV2', 'Seedance']
|
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview']
|
||||||
|
|||||||
@@ -236,18 +236,6 @@ class BatchNanoBananaPro:
|
|||||||
"分辨率": (all_resolutions, {
|
"分辨率": (all_resolutions, {
|
||||||
"default": "2K"
|
"default": "2K"
|
||||||
}),
|
}),
|
||||||
"像素缩放": ("BOOLEAN", {
|
|
||||||
"default": False,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
}),
|
|
||||||
"分辨率像素": ("FLOAT", {
|
|
||||||
"default": 1.0,
|
|
||||||
"min": 0.1,
|
|
||||||
"max": 100.0,
|
|
||||||
"step": 0.1,
|
|
||||||
"display": "number"
|
|
||||||
}),
|
|
||||||
"谷歌搜索(联网)": (["关闭", "打开"], {
|
"谷歌搜索(联网)": (["关闭", "打开"], {
|
||||||
"default": "关闭"
|
"default": "关闭"
|
||||||
}),
|
}),
|
||||||
@@ -298,11 +286,6 @@ class BatchNanoBananaPro:
|
|||||||
"保存路径": ("STRING", {
|
"保存路径": ("STRING", {
|
||||||
"default": "",
|
"default": "",
|
||||||
"multiline": False
|
"multiline": False
|
||||||
}),
|
|
||||||
"跳过错误": ("BOOLEAN", {
|
|
||||||
"default": False,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
})
|
})
|
||||||
},
|
},
|
||||||
"optional": optional_inputs
|
"optional": optional_inputs
|
||||||
@@ -324,8 +307,6 @@ class BatchNanoBananaPro:
|
|||||||
folder2: Optional[str],
|
folder2: Optional[str],
|
||||||
folder3: Optional[str],
|
folder3: Optional[str],
|
||||||
folder4: Optional[str],
|
folder4: Optional[str],
|
||||||
enable_scaling: bool,
|
|
||||||
target_megapixels: float,
|
|
||||||
folder5: Optional[str] = None,
|
folder5: Optional[str] = None,
|
||||||
folder6: Optional[str] = None,
|
folder6: Optional[str] = None,
|
||||||
folder7: Optional[str] = None,
|
folder7: Optional[str] = None,
|
||||||
@@ -334,48 +315,25 @@ class BatchNanoBananaPro:
|
|||||||
) -> List[List[ImageInfo]]:
|
) -> List[List[ImageInfo]]:
|
||||||
"""
|
"""
|
||||||
加载所有文件夹中的图片
|
加载所有文件夹中的图片
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
folder1-9: 文件夹路径
|
folder1-9: 文件夹路径
|
||||||
enable_scaling: 是否启用像素缩放
|
|
||||||
target_megapixels: 目标像素数(百万像素)
|
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
图片列表的列表
|
图片列表的列表
|
||||||
"""
|
"""
|
||||||
folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9]
|
folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9]
|
||||||
all_images = []
|
all_images = []
|
||||||
|
|
||||||
for i, folder in enumerate(folders, 1):
|
for i, folder in enumerate(folders, 1):
|
||||||
if folder and folder.strip():
|
if folder and folder.strip():
|
||||||
try:
|
try:
|
||||||
images = load_images_from_folder(folder)
|
images = load_images_from_folder(folder)
|
||||||
if images:
|
if images:
|
||||||
# 应用像素缩放
|
|
||||||
if enable_scaling:
|
|
||||||
scaled_images = []
|
|
||||||
for img_info in images:
|
|
||||||
scaled_img = self.resize_to_megapixels(
|
|
||||||
img_info.image,
|
|
||||||
target_megapixels
|
|
||||||
)
|
|
||||||
# 创建新的 ImageInfo,保留其他元数据
|
|
||||||
scaled_info = ImageInfo(
|
|
||||||
image=scaled_img,
|
|
||||||
filename=img_info.filename,
|
|
||||||
extension=img_info.extension,
|
|
||||||
source_path=img_info.source_path
|
|
||||||
)
|
|
||||||
scaled_images.append(scaled_info)
|
|
||||||
images = scaled_images
|
|
||||||
|
|
||||||
all_images.append(images)
|
all_images.append(images)
|
||||||
else:
|
|
||||||
# 空文件夹,静默跳过
|
|
||||||
pass
|
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}")
|
print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}")
|
||||||
|
|
||||||
return all_images
|
return all_images
|
||||||
|
|
||||||
def _create_pairs(
|
def _create_pairs(
|
||||||
@@ -596,8 +554,7 @@ class BatchNanoBananaPro:
|
|||||||
|
|
||||||
total_tasks = len(pairs)
|
total_tasks = len(pairs)
|
||||||
|
|
||||||
# 保持并发数为10不变(按用户要求)
|
max_concurrent = 50
|
||||||
max_concurrent = 10
|
|
||||||
|
|
||||||
# 分批保存的批次大小(与并发数一致)
|
# 分批保存的批次大小(与并发数一致)
|
||||||
save_batch_size = 10
|
save_batch_size = 10
|
||||||
@@ -627,7 +584,7 @@ class BatchNanoBananaPro:
|
|||||||
if num_batches > 1:
|
if num_batches > 1:
|
||||||
print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent},将分 {num_batches} 批执行")
|
print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent},将分 {num_batches} 批执行")
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
# 分批处理:每批最多10个任务
|
# 分批处理:每批最多10个任务
|
||||||
@@ -761,14 +718,11 @@ class BatchNanoBananaPro:
|
|||||||
文件夹7: str,
|
文件夹7: str,
|
||||||
文件夹8: str,
|
文件夹8: str,
|
||||||
文件夹9: str,
|
文件夹9: str,
|
||||||
像素缩放: bool,
|
|
||||||
分辨率像素: float,
|
|
||||||
seed: int,
|
seed: int,
|
||||||
图片配对模式: str,
|
图片配对模式: str,
|
||||||
模型: str,
|
模型: str,
|
||||||
宽高比: str,
|
宽高比: str,
|
||||||
分辨率: str,
|
分辨率: str,
|
||||||
跳过错误: bool = False,
|
|
||||||
保存路径: str = "",
|
保存路径: str = "",
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> Tuple[torch.Tensor]:
|
) -> Tuple[torch.Tensor]:
|
||||||
@@ -778,8 +732,6 @@ class BatchNanoBananaPro:
|
|||||||
Args:
|
Args:
|
||||||
prompt: 提示词
|
prompt: 提示词
|
||||||
文件夹1-9: 图片文件夹路径
|
文件夹1-9: 图片文件夹路径
|
||||||
像素缩放: 是否启用像素缩放
|
|
||||||
分辨率像素: 目标像素数(百万像素)
|
|
||||||
seed: 随机种子
|
seed: 随机种子
|
||||||
保存路径: 输出保存路径
|
保存路径: 输出保存路径
|
||||||
图片配对模式: 1:1 或 1*N
|
图片配对模式: 1:1 或 1*N
|
||||||
@@ -840,7 +792,6 @@ class BatchNanoBananaPro:
|
|||||||
print("BatchNanoBananaPro: 开始加载图片...")
|
print("BatchNanoBananaPro: 开始加载图片...")
|
||||||
image_lists = self._load_folders(
|
image_lists = self._load_folders(
|
||||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||||
像素缩放, 分辨率像素,
|
|
||||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -856,10 +807,6 @@ class BatchNanoBananaPro:
|
|||||||
if key in kwargs and kwargs[key] is not None:
|
if key in kwargs and kwargs[key] is not None:
|
||||||
pil_images = tensor_to_pil(kwargs[key])
|
pil_images = tensor_to_pil(kwargs[key])
|
||||||
for j, img in enumerate(pil_images):
|
for j, img in enumerate(pil_images):
|
||||||
# 如果启用像素缩放,也对参考图进行缩放
|
|
||||||
if 像素缩放:
|
|
||||||
img = self.resize_to_megapixels(img, 分辨率像素)
|
|
||||||
|
|
||||||
manual_images.append(
|
manual_images.append(
|
||||||
ImageInfo(
|
ImageInfo(
|
||||||
image=img,
|
image=img,
|
||||||
@@ -976,9 +923,9 @@ class BatchNanoBananaPro:
|
|||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||||
future = executor.submit(run_async_in_thread)
|
future = executor.submit(run_async_in_thread)
|
||||||
try:
|
try:
|
||||||
results = future.result(timeout=3600) # 1小时超时
|
results = future.result(timeout=900) # 900秒超时
|
||||||
except TimeoutError:
|
except TimeoutError:
|
||||||
print("BatchNanoBananaPro: 任务执行超时(1小时)")
|
print("BatchNanoBananaPro: 任务执行超时(900秒)")
|
||||||
raise RuntimeError("任务执行超时,请减少任务数量或检查网络连接")
|
raise RuntimeError("任务执行超时,请减少任务数量或检查网络连接")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
# 即使失败,也尝试返回部分结果
|
# 即使失败,也尝试返回部分结果
|
||||||
@@ -1078,24 +1025,12 @@ class BatchNanoBananaPro:
|
|||||||
if str(e) == "未授权!":
|
if str(e) == "未授权!":
|
||||||
print("请联系作者授权后方可使用!")
|
print("请联系作者授权后方可使用!")
|
||||||
raise ValueError("未授权!") from None
|
raise ValueError("未授权!") from None
|
||||||
if 跳过错误:
|
|
||||||
print("BatchNanoBananaPro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise ValueError(str(e)) from None
|
raise ValueError(str(e)) from None
|
||||||
|
|
||||||
except RuntimeError as e:
|
except RuntimeError as e:
|
||||||
if 跳过错误:
|
|
||||||
print("BatchNanoBananaPro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise RuntimeError(str(e)) from None
|
raise RuntimeError(str(e)) from None
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
if 跳过错误:
|
|
||||||
print("BatchNanoBananaPro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise type(e)(str(e)) from None
|
raise type(e)(str(e)) from None
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
|
|||||||
@@ -1,781 +0,0 @@
|
|||||||
"""
|
|
||||||
批量 Nano Banana v2 节点
|
|
||||||
BatchNanoBananaPro 的完全复刻,唯一改动:
|
|
||||||
|
|
||||||
将原来 9 个独立「参考图1~9」输入端
|
|
||||||
改为 1 个「参考图」输入端(可选),配合「加载图像(批量)」节点使用。
|
|
||||||
|
|
||||||
「加载图像(批量)」输出 is_output_list=True(list[Tensor]),
|
|
||||||
本节点声明 INPUT_IS_LIST = True 来整体接收该列表,
|
|
||||||
然后在 process_batch() 开头对所有参数统一解包,其余业务逻辑与原节点完全一致。
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
import gc
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import random
|
|
||||||
import asyncio
|
|
||||||
import aiohttp
|
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
|
||||||
from typing import Optional, Tuple, List
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
|
||||||
from ..utils.file_utils import (
|
|
||||||
ImageInfo,
|
|
||||||
load_images_from_folder,
|
|
||||||
pair_images_by_name,
|
|
||||||
pair_images_cartesian,
|
|
||||||
generate_timestamp_filename,
|
|
||||||
save_image,
|
|
||||||
)
|
|
||||||
from ..clients.gemini_client import GeminiAPIClient
|
|
||||||
from ..models_config import (
|
|
||||||
get_enabled_models,
|
|
||||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
|
||||||
get_model_supported_resolutions, get_all_supported_resolutions
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
from comfy.utils import ProgressBar
|
|
||||||
PROGRESS_BAR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
PROGRESS_BAR_AVAILABLE = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
import folder_paths
|
|
||||||
FOLDER_PATHS_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
import psutil
|
|
||||||
MEMORY_MONITOR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
MEMORY_MONITOR_AVAILABLE = False
|
|
||||||
|
|
||||||
DEBUG_LOG_ENABLED = False
|
|
||||||
REQUEST_LOG_ENABLED = False
|
|
||||||
|
|
||||||
_NODE = "BatchNanoBananaV2"
|
|
||||||
|
|
||||||
|
|
||||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
|
||||||
"""
|
|
||||||
将 PIL Image 列表转换为 ComfyUI tensor,安全处理多张不同尺寸的情况。
|
|
||||||
|
|
||||||
ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C],要求 batch 内所有图尺寸相同。
|
|
||||||
当 API 返回多张不同分辨率的图时(主图 + 附图),直接 stack 会崩溃。
|
|
||||||
|
|
||||||
策略:
|
|
||||||
- 所有图均已按原始分辨率保存到磁盘(调用此函数前已完成)
|
|
||||||
- 以第一张图的尺寸为基准,只将尺寸相同的图纳入 tensor 输出
|
|
||||||
- 尺寸不同的图跳过(不 resize、不丢弃磁盘文件),并打印日志提示
|
|
||||||
- 若没有任何图与第一张尺寸相同(极罕见),则只输出第一张
|
|
||||||
"""
|
|
||||||
if not images:
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return pil_to_tensor([placeholder])
|
|
||||||
|
|
||||||
base_size = images[0].size # PIL size = (W, H)
|
|
||||||
matched = [img for img in images if img.size == base_size]
|
|
||||||
skipped = [img for img in images if img.size != base_size]
|
|
||||||
|
|
||||||
if skipped:
|
|
||||||
sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped)
|
|
||||||
print(
|
|
||||||
f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str}),"
|
|
||||||
f"已按原始分辨率保存到磁盘,tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 "
|
|
||||||
f"({base_size[0]}×{base_size[1]})"
|
|
||||||
)
|
|
||||||
|
|
||||||
return pil_to_tensor(matched if matched else [images[0]])
|
|
||||||
|
|
||||||
|
|
||||||
class BatchNanaBananaV2:
|
|
||||||
"""
|
|
||||||
批量 Nano Banana v2
|
|
||||||
|
|
||||||
与 BatchNanoBananaPro 完全一致,参考图输入方式不同:
|
|
||||||
- 原版:9 个独立可选端口(参考图1~9)
|
|
||||||
- v2:1 个可选端口「参考图」,配合「加载图像(批量)」可传入任意数量图片
|
|
||||||
"""
|
|
||||||
|
|
||||||
ASPECT_RATIOS = [
|
|
||||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
|
||||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
|
||||||
"1:4", "4:1", "1:8", "8:1"
|
|
||||||
]
|
|
||||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
|
||||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
self.client = None
|
|
||||||
|
|
||||||
def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image:
|
|
||||||
current_pixels = image.width * image.height
|
|
||||||
target_pixels = int(target_megapixels * 1_000_000)
|
|
||||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
|
||||||
return image
|
|
||||||
scale = (target_pixels / current_pixels) ** 0.5
|
|
||||||
new_width = max(1, int(image.width * scale))
|
|
||||||
new_height = max(1, int(image.height * scale))
|
|
||||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
enabled_models = get_enabled_models()
|
|
||||||
if not enabled_models:
|
|
||||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
|
||||||
|
|
||||||
all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS
|
|
||||||
all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS
|
|
||||||
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt": ("STRING", {"default": "一个中国女子的OOTD", "multiline": True}),
|
|
||||||
"模型": (enabled_models, {"default": enabled_models[0]}),
|
|
||||||
"宽高比": (all_aspect_ratios, {"default": "1:1"}),
|
|
||||||
"分辨率": (all_resolutions, {"default": "2K"}),
|
|
||||||
"像素缩放": ("BOOLEAN", {"default": False, "label_on": "打开", "label_off": "关闭"}),
|
|
||||||
"分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}),
|
|
||||||
"谷歌搜索(联网)": (["关闭", "打开"], {"default": "关闭"}),
|
|
||||||
"图片搜索(联网)": (["关闭", "打开"], {"default": "关闭"}),
|
|
||||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
||||||
"文件夹1": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹2": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹3": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹4": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹5": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹6": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹7": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹8": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"文件夹9": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
"保存路径": ("STRING", {"default": "", "multiline": False}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
# 单个参考图端口,接受普通 IMAGE 或「加载图像(批量)」输出的列表
|
|
||||||
"参考图": ("IMAGE",),
|
|
||||||
"图片配对模式": (cls.PAIRING_MODES, {"default": "不配对"}),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("IMAGE",)
|
|
||||||
RETURN_NAMES = ("输出图像",)
|
|
||||||
FUNCTION = "process_batch"
|
|
||||||
CATEGORY = "image/batch"
|
|
||||||
|
|
||||||
# 声明 INPUT_IS_LIST,使 ComfyUI 将「加载图像(批量)」的 list[Tensor]
|
|
||||||
# 整体传入而非逐张迭代执行,同时其余所有参数也会被包进 list,需解包。
|
|
||||||
INPUT_IS_LIST = True
|
|
||||||
|
|
||||||
# ------------------------------------------------------------------ #
|
|
||||||
# 以下方法与 BatchNanoBananaPro 完全相同
|
|
||||||
# ------------------------------------------------------------------ #
|
|
||||||
|
|
||||||
def _load_folders(
|
|
||||||
self,
|
|
||||||
folder1, folder2, folder3, folder4,
|
|
||||||
enable_scaling, target_megapixels,
|
|
||||||
folder5=None, folder6=None, folder7=None, folder8=None, folder9=None,
|
|
||||||
) -> List[List[ImageInfo]]:
|
|
||||||
folders = [folder1, folder2, folder3, folder4,
|
|
||||||
folder5, folder6, folder7, folder8, folder9]
|
|
||||||
all_images = []
|
|
||||||
for i, folder in enumerate(folders, 1):
|
|
||||||
if folder and folder.strip():
|
|
||||||
try:
|
|
||||||
images = load_images_from_folder(folder)
|
|
||||||
if images:
|
|
||||||
if enable_scaling:
|
|
||||||
scaled = []
|
|
||||||
for info in images:
|
|
||||||
scaled_img = self.resize_to_megapixels(info.image, target_megapixels)
|
|
||||||
scaled.append(ImageInfo(
|
|
||||||
image=scaled_img,
|
|
||||||
filename=info.filename,
|
|
||||||
extension=info.extension,
|
|
||||||
source_path=info.source_path
|
|
||||||
))
|
|
||||||
images = scaled
|
|
||||||
all_images.append(images)
|
|
||||||
except ValueError as e:
|
|
||||||
print(f"{_NODE}: 文件夹{i} 加载失败 - {e}")
|
|
||||||
return all_images
|
|
||||||
|
|
||||||
def _create_pairs(
|
|
||||||
self,
|
|
||||||
image_lists: List[List[ImageInfo]],
|
|
||||||
pairing_mode: str,
|
|
||||||
manual_images: Optional[List[ImageInfo]] = None
|
|
||||||
) -> List[Tuple[ImageInfo, ...]]:
|
|
||||||
if pairing_mode == "不配对":
|
|
||||||
if len(image_lists) > 1:
|
|
||||||
raise ValueError("「不配对」模式只支持单个文件夹,请清空其他文件夹路径")
|
|
||||||
if image_lists and manual_images:
|
|
||||||
return [(img,) + tuple(manual_images) for img in image_lists[0]]
|
|
||||||
elif image_lists:
|
|
||||||
return [(img,) for img in image_lists[0]]
|
|
||||||
else:
|
|
||||||
return []
|
|
||||||
|
|
||||||
if not image_lists:
|
|
||||||
return []
|
|
||||||
|
|
||||||
if len(image_lists) == 1:
|
|
||||||
base_pairs = [(img,) for img in image_lists[0]]
|
|
||||||
elif pairing_mode == "按相同图片命名":
|
|
||||||
base_pairs = list(pair_images_by_name(*image_lists))
|
|
||||||
else:
|
|
||||||
base_pairs = list(pair_images_cartesian(*image_lists))
|
|
||||||
|
|
||||||
if manual_images:
|
|
||||||
manual_tuple = tuple(manual_images)
|
|
||||||
base_pairs = [pair + manual_tuple for pair in base_pairs]
|
|
||||||
|
|
||||||
return base_pairs
|
|
||||||
|
|
||||||
async def _generate_single_task(
|
|
||||||
self,
|
|
||||||
client: GeminiAPIClient,
|
|
||||||
session: aiohttp.ClientSession,
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images: List[ImageInfo],
|
|
||||||
output_folder: str,
|
|
||||||
task_index: int,
|
|
||||||
enable_grounding: bool = True,
|
|
||||||
enable_image_search: bool = False,
|
|
||||||
base_filename: str = None,
|
|
||||||
) -> dict:
|
|
||||||
result = {
|
|
||||||
"task_index": task_index,
|
|
||||||
"prompt": prompt,
|
|
||||||
"success": False,
|
|
||||||
"generated_count": 0,
|
|
||||||
"saved_files": [],
|
|
||||||
"output_images": [],
|
|
||||||
"error": None
|
|
||||||
}
|
|
||||||
try:
|
|
||||||
input_pil_images = [info.image for info in images]
|
|
||||||
generated_images = []
|
|
||||||
try:
|
|
||||||
gen_result = await client.generate_single_async(
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=input_pil_images,
|
|
||||||
session=session,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
)
|
|
||||||
if gen_result:
|
|
||||||
images_list, timing_info = gen_result
|
|
||||||
generated_images.extend(images_list)
|
|
||||||
except Exception as e:
|
|
||||||
import traceback
|
|
||||||
error_msg = str(e)
|
|
||||||
error_traceback = traceback.format_exc()
|
|
||||||
print(f"=" * 80)
|
|
||||||
print(f"🔍 【原始报错信息展示】")
|
|
||||||
print(f"=" * 80)
|
|
||||||
print(f"任务编号: {task_index + 1}")
|
|
||||||
print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
|
||||||
print(f"模型: {model}")
|
|
||||||
print(f"分辨率: {resolution}")
|
|
||||||
print(f"宽高比: {aspect_ratio}")
|
|
||||||
print(f"-" * 80)
|
|
||||||
print(f"错误信息: {error_msg}")
|
|
||||||
print(f"-" * 80)
|
|
||||||
print(f"完整堆栈追踪:")
|
|
||||||
print(error_traceback)
|
|
||||||
print(f"=" * 80)
|
|
||||||
result["error"] = error_msg
|
|
||||||
|
|
||||||
for i, gen_img in enumerate(generated_images):
|
|
||||||
if base_filename:
|
|
||||||
base_name = base_filename
|
|
||||||
counter = 0
|
|
||||||
while True:
|
|
||||||
filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png"
|
|
||||||
output_path = os.path.join(output_folder, filename)
|
|
||||||
if not os.path.exists(output_path):
|
|
||||||
break
|
|
||||||
counter += 1
|
|
||||||
else:
|
|
||||||
output_path = generate_timestamp_filename(
|
|
||||||
output_folder=output_folder, extension=".png"
|
|
||||||
)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
result["saved_files"].append(output_path)
|
|
||||||
gen_img = None
|
|
||||||
|
|
||||||
if len(generated_images) > 0:
|
|
||||||
result["success"] = True
|
|
||||||
result["generated_count"] = len(generated_images)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
result["error"] = str(e)
|
|
||||||
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def _process_batch_async(
|
|
||||||
self,
|
|
||||||
pairs: List[Tuple[ImageInfo, ...]],
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
output_folder: str,
|
|
||||||
pbar=None,
|
|
||||||
prompts_per_task: Optional[List[str]] = None,
|
|
||||||
enable_grounding: bool = True,
|
|
||||||
enable_image_search: bool = False,
|
|
||||||
) -> List[dict]:
|
|
||||||
if self.client is None:
|
|
||||||
self.client = GeminiAPIClient()
|
|
||||||
|
|
||||||
total_tasks = len(pairs)
|
|
||||||
max_concurrent = 10
|
|
||||||
|
|
||||||
print(f"{_NODE}: 检测到 {total_tasks} 个任务")
|
|
||||||
|
|
||||||
all_results = []
|
|
||||||
completed = 0
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
|
||||||
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
|
||||||
process = psutil.Process()
|
|
||||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB")
|
|
||||||
|
|
||||||
show_milestone = total_tasks >= 50
|
|
||||||
milestones = [0.2, 0.4, 0.6, 0.8, 1.0]
|
|
||||||
milestone_index = 0
|
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"{_NODE}: 任务数 {total_tasks} 超过并发上限 {max_concurrent},将分 {num_batches} 批执行")
|
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
|
||||||
for batch_idx in range(num_batches):
|
|
||||||
start_idx = batch_idx * max_concurrent
|
|
||||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
|
||||||
batch_pairs = pairs[start_idx:end_idx]
|
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"{_NODE}: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
|
||||||
|
|
||||||
tasks = []
|
|
||||||
for i, pair in enumerate(batch_pairs):
|
|
||||||
task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt
|
|
||||||
base_filename = None
|
|
||||||
if pair and len(pair) > 0:
|
|
||||||
first_image = pair[0]
|
|
||||||
if hasattr(first_image, 'filename'):
|
|
||||||
base_filename = first_image.filename
|
|
||||||
|
|
||||||
task = asyncio.create_task(
|
|
||||||
self._generate_single_task(
|
|
||||||
client=self.client,
|
|
||||||
session=session,
|
|
||||||
prompt=task_prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=list(pair),
|
|
||||||
output_folder=output_folder,
|
|
||||||
task_index=start_idx + i,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
base_filename=base_filename,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
tasks.append(task)
|
|
||||||
|
|
||||||
batch_results = []
|
|
||||||
for coro in asyncio.as_completed(tasks):
|
|
||||||
result_data = None
|
|
||||||
try:
|
|
||||||
result = await coro
|
|
||||||
if isinstance(result, Exception):
|
|
||||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []}
|
|
||||||
batch_results.append(result_data)
|
|
||||||
else:
|
|
||||||
result_data = result
|
|
||||||
batch_results.append(result)
|
|
||||||
except Exception as e:
|
|
||||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []}
|
|
||||||
batch_results.append(result_data)
|
|
||||||
|
|
||||||
completed += 1
|
|
||||||
if result_data and result_data.get("success", False):
|
|
||||||
success_count += 1
|
|
||||||
print(f"{_NODE}: 任务 {completed}/{total_tasks} 成功 ✓")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
|
||||||
print(f"{_NODE}: 任务 {completed}/{total_tasks} 失败 ✗")
|
|
||||||
print(f"=" * 80)
|
|
||||||
print(f"🔍 【原始报错信息展示】")
|
|
||||||
print(f"=" * 80)
|
|
||||||
print(f"任务编号: {completed}/{total_tasks}")
|
|
||||||
print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
|
||||||
print(f"-" * 80)
|
|
||||||
print(f"错误详情:")
|
|
||||||
print(error_msg)
|
|
||||||
print(f"=" * 80)
|
|
||||||
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
|
|
||||||
if show_milestone and milestone_index < len(milestones):
|
|
||||||
progress = completed / total_tasks
|
|
||||||
if progress >= milestones[milestone_index]:
|
|
||||||
percentage = int(milestones[milestone_index] * 100)
|
|
||||||
print(f"{_NODE}: >>> 进度 {percentage}% <<<")
|
|
||||||
milestone_index += 1
|
|
||||||
|
|
||||||
all_results.extend(batch_results)
|
|
||||||
print(f"{_NODE}: 第 {batch_idx + 1} 批完成,开始分批保存...")
|
|
||||||
|
|
||||||
batch_success = sum(1 for r in batch_results if r.get("success", False))
|
|
||||||
batch_fail = len(batch_results) - batch_success
|
|
||||||
batch_generated = sum(r.get("generated_count", 0) for r in batch_results)
|
|
||||||
print(f"{_NODE}: 本批结果 - 成功: {batch_success}/{len(batch_results)},生成: {batch_generated} 张")
|
|
||||||
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
|
||||||
current_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
memory_increase = current_memory - initial_memory
|
|
||||||
print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
|
||||||
if current_memory > 2000:
|
|
||||||
print(f"⚠️ {_NODE}: 内存使用过高!但图片已分批保存,即使崩溃也不会丢失已完成的任务")
|
|
||||||
|
|
||||||
await asyncio.sleep(0.5)
|
|
||||||
|
|
||||||
return all_results
|
|
||||||
|
|
||||||
def process_batch(
|
|
||||||
self,
|
|
||||||
prompt,
|
|
||||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
|
||||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9,
|
|
||||||
像素缩放,
|
|
||||||
分辨率像素,
|
|
||||||
seed,
|
|
||||||
模型,
|
|
||||||
宽高比,
|
|
||||||
分辨率,
|
|
||||||
保存路径,
|
|
||||||
**kwargs
|
|
||||||
) -> Tuple[torch.Tensor]:
|
|
||||||
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# INPUT_IS_LIST=True 时,所有参数均为 list,先统一解包为标量
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
def _unpack(v):
|
|
||||||
return v[0] if isinstance(v, list) else v
|
|
||||||
|
|
||||||
prompt = _unpack(prompt)
|
|
||||||
文件夹1 = _unpack(文件夹1)
|
|
||||||
文件夹2 = _unpack(文件夹2)
|
|
||||||
文件夹3 = _unpack(文件夹3)
|
|
||||||
文件夹4 = _unpack(文件夹4)
|
|
||||||
文件夹5 = _unpack(文件夹5)
|
|
||||||
文件夹6 = _unpack(文件夹6)
|
|
||||||
文件夹7 = _unpack(文件夹7)
|
|
||||||
文件夹8 = _unpack(文件夹8)
|
|
||||||
文件夹9 = _unpack(文件夹9)
|
|
||||||
像素缩放 = _unpack(像素缩放)
|
|
||||||
分辨率像素 = _unpack(分辨率像素)
|
|
||||||
seed = _unpack(seed)
|
|
||||||
模型 = _unpack(模型)
|
|
||||||
宽高比 = _unpack(宽高比)
|
|
||||||
分辨率 = _unpack(分辨率)
|
|
||||||
保存路径 = _unpack(保存路径)
|
|
||||||
|
|
||||||
# 含全角括号的参数名无法作为形参,从 kwargs 中提取
|
|
||||||
enable_grounding: bool = (_unpack(kwargs.pop("谷歌搜索(联网)", "关闭"))) == "打开"
|
|
||||||
enable_image_search: bool = (_unpack(kwargs.pop("图片搜索(联网)", "关闭"))) == "打开"
|
|
||||||
|
|
||||||
# 图片配对模式(可选参数)
|
|
||||||
图片配对模式 = _unpack(kwargs.pop("图片配对模式", "不配对"))
|
|
||||||
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# 收集参考图:兼容两种来源
|
|
||||||
# 1. 「加载图像(批量)」→ is_output_list=True → list[Tensor]
|
|
||||||
# INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor],需展平
|
|
||||||
# 2. 普通 IMAGE 端口(单 tensor 或 batch tensor)→ list 中只有 1 个元素
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
ref_raw = kwargs.pop("参考图", None)
|
|
||||||
manual_images: List[ImageInfo] = []
|
|
||||||
|
|
||||||
if ref_raw is not None:
|
|
||||||
items = ref_raw if isinstance(ref_raw, list) else [ref_raw]
|
|
||||||
idx = 0
|
|
||||||
for item in items:
|
|
||||||
if item is None:
|
|
||||||
continue
|
|
||||||
if isinstance(item, list):
|
|
||||||
sub_tensors = item
|
|
||||||
elif isinstance(item, torch.Tensor):
|
|
||||||
sub_tensors = [item]
|
|
||||||
else:
|
|
||||||
continue
|
|
||||||
for tensor in sub_tensors:
|
|
||||||
if tensor is None or not isinstance(tensor, torch.Tensor):
|
|
||||||
continue
|
|
||||||
pil_images = tensor_to_pil(tensor)
|
|
||||||
for j, img in enumerate(pil_images):
|
|
||||||
if 像素缩放:
|
|
||||||
img = self.resize_to_megapixels(img, 分辨率像素)
|
|
||||||
manual_images.append(ImageInfo(
|
|
||||||
image=img,
|
|
||||||
filename=f"manual_{idx}_{j}",
|
|
||||||
extension=".png",
|
|
||||||
source_path=""
|
|
||||||
))
|
|
||||||
idx += 1
|
|
||||||
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# 以下逻辑与 BatchNanoBananaPro.process_batch() 完全一致
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
start_time = time.time()
|
|
||||||
|
|
||||||
try:
|
|
||||||
random.seed(seed)
|
|
||||||
np.random.seed(seed % (2 ** 32))
|
|
||||||
|
|
||||||
has_any_folder = any(
|
|
||||||
f and f.strip()
|
|
||||||
for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
|
||||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9]
|
|
||||||
)
|
|
||||||
if not has_any_folder:
|
|
||||||
raise ValueError("请至少填写一个文件夹路径,该节点专为批量文件夹处理设计")
|
|
||||||
|
|
||||||
supported_resolutions = get_model_supported_resolutions(模型)
|
|
||||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
|
||||||
raise ValueError(
|
|
||||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的分辨率:{', '.join(supported_resolutions)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
|
||||||
if supported_ratios and 宽高比 not in supported_ratios:
|
|
||||||
raise ValueError(
|
|
||||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = [
|
|
||||||
"nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"
|
|
||||||
]
|
|
||||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
|
||||||
raise ValueError(
|
|
||||||
f"模型 \"{模型}\" 不支持【图片搜索(联网)】功能!"
|
|
||||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
|
||||||
)
|
|
||||||
|
|
||||||
print(f"{_NODE}: 开始加载图片...")
|
|
||||||
image_lists = self._load_folders(
|
|
||||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
|
||||||
像素缩放, 分辨率像素,
|
|
||||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
|
||||||
)
|
|
||||||
|
|
||||||
total_folder_images = sum(len(lst) for lst in image_lists)
|
|
||||||
if total_folder_images == 0:
|
|
||||||
raise ValueError("文件夹中未找到任何图片,请检查文件夹路径是否正确")
|
|
||||||
|
|
||||||
pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None)
|
|
||||||
|
|
||||||
if not pairs:
|
|
||||||
raise ValueError("配对结果为空,请检查输入")
|
|
||||||
|
|
||||||
batch_prompts = parse_batch_prompts(prompt)
|
|
||||||
prompts_per_task = None
|
|
||||||
if batch_prompts:
|
|
||||||
expanded_pairs = []
|
|
||||||
expanded_prompts = []
|
|
||||||
for pair in pairs:
|
|
||||||
for bp in batch_prompts:
|
|
||||||
expanded_pairs.append(pair)
|
|
||||||
expanded_prompts.append(bp)
|
|
||||||
pairs = expanded_pairs
|
|
||||||
prompts_per_task = expanded_prompts
|
|
||||||
|
|
||||||
total_tasks = len(pairs)
|
|
||||||
|
|
||||||
grounding_str = ""
|
|
||||||
if enable_image_search:
|
|
||||||
grounding_str = " | 谷歌图片搜索接地"
|
|
||||||
elif enable_grounding:
|
|
||||||
grounding_str = " | 谷歌搜索接地"
|
|
||||||
|
|
||||||
if batch_prompts:
|
|
||||||
print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}")
|
|
||||||
else:
|
|
||||||
print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}")
|
|
||||||
|
|
||||||
pbar = None
|
|
||||||
if PROGRESS_BAR_AVAILABLE:
|
|
||||||
pbar = ProgressBar(total_tasks)
|
|
||||||
|
|
||||||
has_save_path = bool(保存路径 and 保存路径.strip())
|
|
||||||
if not has_save_path:
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
保存路径 = folder_paths.get_output_directory()
|
|
||||||
has_save_path = True
|
|
||||||
print(f"{_NODE}: 未设置保存路径,将使用 ComfyUI 默认 output 目录: {保存路径}")
|
|
||||||
else:
|
|
||||||
print(f"{_NODE}: 未设置保存路径,图片将输出到节点")
|
|
||||||
|
|
||||||
if has_save_path:
|
|
||||||
try:
|
|
||||||
os.makedirs(保存路径, exist_ok=True)
|
|
||||||
test_file = os.path.join(保存路径, ".write_test")
|
|
||||||
with open(test_file, 'w') as f:
|
|
||||||
f.write("test")
|
|
||||||
os.remove(test_file)
|
|
||||||
print(f"{_NODE}: 保存路径验证通过: {保存路径}")
|
|
||||||
except Exception as e:
|
|
||||||
raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}")
|
|
||||||
|
|
||||||
if self.client is None:
|
|
||||||
try:
|
|
||||||
self.client = GeminiAPIClient()
|
|
||||||
except ValueError as e:
|
|
||||||
raise ValueError(f"初始化 API 客户端失败: {str(e)}")
|
|
||||||
|
|
||||||
def run_async_in_thread():
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
asyncio.set_event_loop(loop)
|
|
||||||
try:
|
|
||||||
return loop.run_until_complete(
|
|
||||||
self._process_batch_async(
|
|
||||||
pairs=pairs,
|
|
||||||
prompt=prompt,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
output_folder=保存路径,
|
|
||||||
pbar=pbar,
|
|
||||||
prompts_per_task=prompts_per_task,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"{_NODE}: 异步任务执行异常: {str(e)}")
|
|
||||||
raise
|
|
||||||
finally:
|
|
||||||
loop.close()
|
|
||||||
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
||||||
future = executor.submit(run_async_in_thread)
|
|
||||||
try:
|
|
||||||
results = future.result(timeout=3600)
|
|
||||||
except TimeoutError:
|
|
||||||
print(f"{_NODE}: 任务执行超时(1小时)")
|
|
||||||
raise RuntimeError("任务执行超时,请减少任务数量或检查网络连接")
|
|
||||||
except Exception as e:
|
|
||||||
print(f"{_NODE}: 任务执行失败: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
|
||||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
avg_time = elapsed / success_count if success_count > 0 else 0
|
|
||||||
avg_time_str = f"{avg_time:.1f}s/张" if success_count > 0 else "N/A"
|
|
||||||
|
|
||||||
print("=" * 60)
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 生成 {total_generated} 张 | 平均 {avg_time_str}")
|
|
||||||
if has_save_path:
|
|
||||||
print(f"保存路径: {保存路径}")
|
|
||||||
else:
|
|
||||||
print("保存路径: 未设置(仅输出到节点)")
|
|
||||||
|
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
|
||||||
if failed_results:
|
|
||||||
print(f"-" * 60)
|
|
||||||
print(f"❌ 失败任务汇总: {len(failed_results)} 个")
|
|
||||||
print(f"-" * 60)
|
|
||||||
for idx, failed in enumerate(failed_results[:3], 1):
|
|
||||||
task_num = failed.get('task_index', '?') + 1
|
|
||||||
error_msg = failed.get('error', '未知错误')
|
|
||||||
print(f"\n【失败任务 #{task_num}】")
|
|
||||||
print(f"错误信息: {error_msg}")
|
|
||||||
if len(failed_results) > 3:
|
|
||||||
remaining = [str(r.get('task_index', '?') + 1) for r in failed_results[3:]]
|
|
||||||
print(f"\n其他失败任务编号: {', '.join(remaining)}")
|
|
||||||
print(f"-" * 60)
|
|
||||||
|
|
||||||
output_images = []
|
|
||||||
if all_saved_files:
|
|
||||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
|
||||||
try:
|
|
||||||
output_images.append(Image.open(fp))
|
|
||||||
except Exception as e:
|
|
||||||
print(f"{_NODE}: 无法加载图片 {fp} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
total_saved = len(all_saved_files)
|
|
||||||
print(f"{_NODE}: 任务完成!共保存 {total_saved} 张图片到磁盘")
|
|
||||||
if total_saved > 0:
|
|
||||||
print(f"{_NODE}: 最新保存的文件: {all_saved_files[-1]}")
|
|
||||||
|
|
||||||
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
except ValueError as e:
|
|
||||||
if str(e) == "未授权!":
|
|
||||||
print("请联系作者授权后方可使用!")
|
|
||||||
raise ValueError("未授权!") from None
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_msg}")
|
|
||||||
raise ValueError(error_msg) from None
|
|
||||||
|
|
||||||
except RuntimeError as e:
|
|
||||||
error_full = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_full}")
|
|
||||||
raise RuntimeError(error_full) from None
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_msg}")
|
|
||||||
raise type(e)(error_msg) from None
|
|
||||||
|
|
||||||
finally:
|
|
||||||
if self.client is not None:
|
|
||||||
try:
|
|
||||||
balance_data = self.client.query_balance_sync()
|
|
||||||
balance_info = self.client.format_balance_info(balance_data)
|
|
||||||
print(f"{_NODE}: {balance_info}")
|
|
||||||
print("=" * 60)
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
gc.collect()
|
|
||||||
@@ -1,642 +0,0 @@
|
|||||||
"""
|
|
||||||
全能生图(批量)节点
|
|
||||||
ComfyUI 自定义节点,用于批量处理图像生成任务
|
|
||||||
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
|
||||||
"""
|
|
||||||
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import random
|
|
||||||
import asyncio
|
|
||||||
import aiohttp
|
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
|
||||||
from typing import Optional, Tuple, List
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
|
||||||
from ..utils.file_utils import (
|
|
||||||
ImageInfo,
|
|
||||||
load_images_from_folder,
|
|
||||||
pair_images_by_name,
|
|
||||||
pair_images_cartesian,
|
|
||||||
generate_timestamp_filename,
|
|
||||||
save_image,
|
|
||||||
)
|
|
||||||
from ..clients.openai_client import OpenAIAPIClient
|
|
||||||
from ..models_config import (
|
|
||||||
get_enabled_models,
|
|
||||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
|
||||||
get_model_supported_resolutions, get_all_supported_resolutions
|
|
||||||
)
|
|
||||||
|
|
||||||
# 导入 ComfyUI 原生进度条
|
|
||||||
try:
|
|
||||||
from comfy.utils import ProgressBar
|
|
||||||
PROGRESS_BAR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
PROGRESS_BAR_AVAILABLE = False
|
|
||||||
print("⚠️ 全能生图(批量): comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
|
||||||
|
|
||||||
# 导入 ComfyUI 的文件夹路径管理
|
|
||||||
try:
|
|
||||||
import folder_paths
|
|
||||||
FOLDER_PATHS_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
|
||||||
print("⚠️ 全能生图(批量): folder_paths 不可用,将无法使用默认保存路径")
|
|
||||||
|
|
||||||
# 内存监控(可选)
|
|
||||||
try:
|
|
||||||
import psutil
|
|
||||||
MEMORY_MONITOR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
MEMORY_MONITOR_AVAILABLE = False
|
|
||||||
print("⚠️ 全能生图(批量): psutil 不可用,内存监控功能禁用")
|
|
||||||
|
|
||||||
# ============================================================================
|
|
||||||
# 调试日志配置
|
|
||||||
# ============================================================================
|
|
||||||
DEBUG_LOG_ENABLED = False
|
|
||||||
REQUEST_LOG_ENABLED = False
|
|
||||||
# ============================================================================
|
|
||||||
|
|
||||||
|
|
||||||
class BatchQuanNengShengTu:
|
|
||||||
"""
|
|
||||||
全能生图(批量)节点
|
|
||||||
|
|
||||||
功能:
|
|
||||||
- 从多个文件夹加载图片
|
|
||||||
- 支持三种配对模式:
|
|
||||||
* 按相同图片命名 - 索引配对(文件夹之间按位置配对)
|
|
||||||
* 1*N - 笛卡尔积配对(所有可能组合)
|
|
||||||
* 不配对 - 固定参考图模式(文件夹图片依次与所有参考图组合)
|
|
||||||
- 批量调用 API 生成图像
|
|
||||||
- 智能命名保存(保留原始文件名)
|
|
||||||
- 并发控制(默认最大 10)
|
|
||||||
|
|
||||||
注意:
|
|
||||||
- 「不配对」模式只支持单个文件夹
|
|
||||||
- 支持的模型列表从 models_config.py 动态加载
|
|
||||||
"""
|
|
||||||
|
|
||||||
MODELS = None
|
|
||||||
ASPECT_RATIOS = [
|
|
||||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
|
||||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
|
||||||
"1:4", "4:1", "1:8", "8:1"
|
|
||||||
]
|
|
||||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
|
||||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
"""初始化节点"""
|
|
||||||
self.client = None
|
|
||||||
|
|
||||||
def resize_to_megapixels(
|
|
||||||
self,
|
|
||||||
image: Image.Image,
|
|
||||||
target_megapixels: float
|
|
||||||
) -> Image.Image:
|
|
||||||
"""将图像缩放到指定的总像素数,保持纵横比"""
|
|
||||||
current_pixels = image.width * image.height
|
|
||||||
target_pixels = int(target_megapixels * 1_000_000)
|
|
||||||
|
|
||||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
|
||||||
return image
|
|
||||||
|
|
||||||
scale = (target_pixels / current_pixels) ** 0.5
|
|
||||||
new_width = max(1, int(image.width * scale))
|
|
||||||
new_height = max(1, int(image.height * scale))
|
|
||||||
|
|
||||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
"""定义输入参数"""
|
|
||||||
enabled_models = get_enabled_models()
|
|
||||||
enabled_models = [m for m in enabled_models if "限时特价" not in m]
|
|
||||||
|
|
||||||
if not enabled_models:
|
|
||||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
|
||||||
|
|
||||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
|
||||||
if not all_aspect_ratios:
|
|
||||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
|
||||||
|
|
||||||
all_resolutions = get_all_supported_resolutions()
|
|
||||||
if not all_resolutions:
|
|
||||||
all_resolutions = cls.RESOLUTIONS
|
|
||||||
|
|
||||||
optional_inputs = {}
|
|
||||||
for i in range(1, 10):
|
|
||||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
|
||||||
|
|
||||||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
|
||||||
"default": "不配对"
|
|
||||||
})
|
|
||||||
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"提示词": ("STRING", {
|
|
||||||
"default": "一个中国女子的OOTD",
|
|
||||||
"multiline": True
|
|
||||||
}),
|
|
||||||
"模型": (enabled_models, {
|
|
||||||
"default": enabled_models[0]
|
|
||||||
}),
|
|
||||||
"宽高比": (all_aspect_ratios, {
|
|
||||||
"default": "1:1"
|
|
||||||
}),
|
|
||||||
"分辨率": (all_resolutions, {
|
|
||||||
"default": "2K"
|
|
||||||
}),
|
|
||||||
"像素缩放": ("BOOLEAN", {
|
|
||||||
"default": False,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
}),
|
|
||||||
"分辨率像素": ("FLOAT", {
|
|
||||||
"default": 1.0,
|
|
||||||
"min": 0.1,
|
|
||||||
"max": 100.0,
|
|
||||||
"step": 0.1,
|
|
||||||
"display": "number"
|
|
||||||
}),
|
|
||||||
"seed": ("INT", {
|
|
||||||
"default": 0,
|
|
||||||
"min": 0,
|
|
||||||
"max": 0xffffffffffffffff
|
|
||||||
}),
|
|
||||||
"文件夹1": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹2": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹3": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹4": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹5": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹6": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹7": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹8": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"文件夹9": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
}),
|
|
||||||
"保存路径": ("STRING", {
|
|
||||||
"default": "",
|
|
||||||
"multiline": False
|
|
||||||
})
|
|
||||||
},
|
|
||||||
"optional": optional_inputs
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("IMAGE",)
|
|
||||||
RETURN_NAMES = ("输出图像",)
|
|
||||||
FUNCTION = "process_batch"
|
|
||||||
CATEGORY = "image/batch"
|
|
||||||
|
|
||||||
def _load_folders(
|
|
||||||
self,
|
|
||||||
folder1: str,
|
|
||||||
folder2: Optional[str],
|
|
||||||
folder3: Optional[str],
|
|
||||||
folder4: Optional[str],
|
|
||||||
enable_scaling: bool,
|
|
||||||
target_megapixels: float,
|
|
||||||
folder5: Optional[str] = None,
|
|
||||||
folder6: Optional[str] = None,
|
|
||||||
folder7: Optional[str] = None,
|
|
||||||
folder8: Optional[str] = None,
|
|
||||||
folder9: Optional[str] = None,
|
|
||||||
) -> List[List[ImageInfo]]:
|
|
||||||
"""加载所有文件夹中的图片"""
|
|
||||||
folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9]
|
|
||||||
all_images = []
|
|
||||||
|
|
||||||
for i, folder in enumerate(folders, 1):
|
|
||||||
if folder and folder.strip():
|
|
||||||
try:
|
|
||||||
images = load_images_from_folder(folder)
|
|
||||||
if images:
|
|
||||||
if enable_scaling:
|
|
||||||
scaled_images = []
|
|
||||||
for img_info in images:
|
|
||||||
scaled_img = self.resize_to_megapixels(
|
|
||||||
img_info.image,
|
|
||||||
target_megapixels
|
|
||||||
)
|
|
||||||
scaled_info = ImageInfo(
|
|
||||||
image=scaled_img,
|
|
||||||
filename=img_info.filename,
|
|
||||||
extension=img_info.extension,
|
|
||||||
source_path=img_info.source_path
|
|
||||||
)
|
|
||||||
scaled_images.append(scaled_info)
|
|
||||||
images = scaled_images
|
|
||||||
all_images.append(images)
|
|
||||||
except ValueError as e:
|
|
||||||
print(f"全能生图(批量): 文件夹{i} 加载失败 - {e}")
|
|
||||||
|
|
||||||
return all_images
|
|
||||||
|
|
||||||
def _create_pairs(
|
|
||||||
self,
|
|
||||||
image_lists: List[List[ImageInfo]],
|
|
||||||
pairing_mode: str,
|
|
||||||
manual_images: Optional[List[ImageInfo]] = None
|
|
||||||
) -> List[Tuple[ImageInfo, ...]]:
|
|
||||||
"""根据配对模式创建图片组合"""
|
|
||||||
if pairing_mode == "不配对":
|
|
||||||
if len(image_lists) > 1:
|
|
||||||
raise ValueError("「不配对」模式只支持单个文件夹,请清空其他文件夹路径")
|
|
||||||
|
|
||||||
if image_lists and manual_images:
|
|
||||||
folder_images = image_lists[0]
|
|
||||||
pairs = []
|
|
||||||
for img in folder_images:
|
|
||||||
pair = (img,) + tuple(manual_images)
|
|
||||||
pairs.append(pair)
|
|
||||||
return pairs
|
|
||||||
elif image_lists:
|
|
||||||
return [(img,) for img in image_lists[0]]
|
|
||||||
else:
|
|
||||||
return []
|
|
||||||
|
|
||||||
if not image_lists:
|
|
||||||
return []
|
|
||||||
|
|
||||||
if len(image_lists) == 1:
|
|
||||||
base_pairs = [(img,) for img in image_lists[0]]
|
|
||||||
elif pairing_mode == "按相同图片命名":
|
|
||||||
base_pairs = list(pair_images_by_name(*image_lists))
|
|
||||||
else:
|
|
||||||
base_pairs = list(pair_images_cartesian(*image_lists))
|
|
||||||
|
|
||||||
if manual_images:
|
|
||||||
manual_tuple = tuple(manual_images)
|
|
||||||
base_pairs = [pair + manual_tuple for pair in base_pairs]
|
|
||||||
|
|
||||||
return base_pairs
|
|
||||||
|
|
||||||
async def _generate_single_task(
|
|
||||||
self,
|
|
||||||
session: aiohttp.ClientSession,
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images: List[ImageInfo],
|
|
||||||
output_folder: str,
|
|
||||||
task_index: int,
|
|
||||||
base_filename: str = None,
|
|
||||||
) -> dict:
|
|
||||||
"""执行单个生成任务"""
|
|
||||||
result = {
|
|
||||||
"task_index": task_index,
|
|
||||||
"prompt": prompt,
|
|
||||||
"success": False,
|
|
||||||
"generated_count": 0,
|
|
||||||
"saved_files": [],
|
|
||||||
"error": None
|
|
||||||
}
|
|
||||||
|
|
||||||
try:
|
|
||||||
input_pil_images = [info.image for info in images]
|
|
||||||
|
|
||||||
gen_result = await self.client.generate_single_async(
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=input_pil_images,
|
|
||||||
session=session,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=False,
|
|
||||||
enable_image_search=False
|
|
||||||
)
|
|
||||||
|
|
||||||
if gen_result:
|
|
||||||
images_list, _ = gen_result
|
|
||||||
|
|
||||||
import os
|
|
||||||
for gen_img in images_list:
|
|
||||||
if base_filename:
|
|
||||||
base_name = base_filename
|
|
||||||
counter = 0
|
|
||||||
while True:
|
|
||||||
filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png"
|
|
||||||
output_path = os.path.join(output_folder, filename)
|
|
||||||
if not os.path.exists(output_path):
|
|
||||||
break
|
|
||||||
counter += 1
|
|
||||||
else:
|
|
||||||
output_path = generate_timestamp_filename(
|
|
||||||
output_folder=output_folder,
|
|
||||||
extension=".png"
|
|
||||||
)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
result["saved_files"].append(output_path)
|
|
||||||
gen_img = None
|
|
||||||
|
|
||||||
result["success"] = True
|
|
||||||
result["generated_count"] = len(images_list)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
result["error"] = str(e)
|
|
||||||
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def _process_batch_async(
|
|
||||||
self,
|
|
||||||
pairs: List[Tuple[ImageInfo, ...]],
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
output_folder: str,
|
|
||||||
pbar=None,
|
|
||||||
prompts_per_task: Optional[List[str]] = None,
|
|
||||||
) -> List[dict]:
|
|
||||||
"""异步批量处理所有任务"""
|
|
||||||
if self.client is None:
|
|
||||||
self.client = OpenAIAPIClient()
|
|
||||||
|
|
||||||
total_tasks = len(pairs)
|
|
||||||
max_concurrent = 10
|
|
||||||
|
|
||||||
print(f"全能生图(批量): 检测到 {total_tasks} 个任务")
|
|
||||||
|
|
||||||
all_results = []
|
|
||||||
completed = 0
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
|
|
||||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"全能生图(批量): 任务数 {total_tasks} 超过并发上限 {max_concurrent},将分 {num_batches} 批执行")
|
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
|
||||||
for batch_idx in range(num_batches):
|
|
||||||
start_idx = batch_idx * max_concurrent
|
|
||||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
|
||||||
batch_pairs = pairs[start_idx:end_idx]
|
|
||||||
|
|
||||||
if num_batches > 1:
|
|
||||||
print(f"全能生图(批量): 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
|
||||||
|
|
||||||
tasks = []
|
|
||||||
for i, pair in enumerate(batch_pairs):
|
|
||||||
task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt
|
|
||||||
|
|
||||||
base_filename = None
|
|
||||||
if pair and len(pair) > 0:
|
|
||||||
first_image = pair[0]
|
|
||||||
if hasattr(first_image, 'filename'):
|
|
||||||
base_filename = first_image.filename
|
|
||||||
|
|
||||||
task = asyncio.create_task(
|
|
||||||
self._generate_single_task(
|
|
||||||
session=session,
|
|
||||||
prompt=task_prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=list(pair),
|
|
||||||
output_folder=output_folder,
|
|
||||||
task_index=start_idx + i,
|
|
||||||
base_filename=base_filename,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
tasks.append(task)
|
|
||||||
|
|
||||||
batch_results = []
|
|
||||||
for coro in asyncio.as_completed(tasks):
|
|
||||||
result_data = None
|
|
||||||
try:
|
|
||||||
result = await coro
|
|
||||||
if isinstance(result, Exception):
|
|
||||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []}
|
|
||||||
else:
|
|
||||||
result_data = result
|
|
||||||
batch_results.append(result_data)
|
|
||||||
except Exception as e:
|
|
||||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []}
|
|
||||||
batch_results.append(result_data)
|
|
||||||
|
|
||||||
completed += 1
|
|
||||||
|
|
||||||
if result_data and result_data.get("success", False):
|
|
||||||
success_count += 1
|
|
||||||
print(f"全能生图(批量): 任务 {completed}/{total_tasks} 成功 ✓")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
|
||||||
print(f"全能生图(批量): 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}")
|
|
||||||
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
|
|
||||||
all_results.extend(batch_results)
|
|
||||||
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
return all_results
|
|
||||||
|
|
||||||
def process_batch(
|
|
||||||
self,
|
|
||||||
提示词: str,
|
|
||||||
模型: str,
|
|
||||||
宽高比: str,
|
|
||||||
分辨率: str,
|
|
||||||
像素缩放: bool,
|
|
||||||
分辨率像素: float,
|
|
||||||
seed: int,
|
|
||||||
文件夹1: str,
|
|
||||||
文件夹2: str,
|
|
||||||
文件夹3: str,
|
|
||||||
文件夹4: str,
|
|
||||||
文件夹5: str,
|
|
||||||
文件夹6: str,
|
|
||||||
文件夹7: str,
|
|
||||||
文件夹8: str,
|
|
||||||
文件夹9: str,
|
|
||||||
保存路径: str,
|
|
||||||
**kwargs
|
|
||||||
) -> Tuple[torch.Tensor]:
|
|
||||||
"""批量处理图像生成"""
|
|
||||||
start_time = time.time()
|
|
||||||
|
|
||||||
try:
|
|
||||||
random.seed(seed)
|
|
||||||
np.random.seed(seed % (2**32))
|
|
||||||
|
|
||||||
if self.client is None:
|
|
||||||
self.client = OpenAIAPIClient()
|
|
||||||
|
|
||||||
supported_resolutions = get_model_supported_resolutions(模型)
|
|
||||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
|
||||||
raise ValueError(
|
|
||||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的分辨率:{', '.join(supported_resolutions)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
|
||||||
if supported_ratios and 宽高比 not in supported_ratios:
|
|
||||||
raise ValueError(
|
|
||||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
manual_images = []
|
|
||||||
for i in range(1, 10):
|
|
||||||
key = f"参考图{i}"
|
|
||||||
if key in kwargs and kwargs[key] is not None:
|
|
||||||
pil_imgs = tensor_to_pil(kwargs[key])
|
|
||||||
for pil_img in pil_imgs:
|
|
||||||
manual_images.append(ImageInfo(
|
|
||||||
image=pil_img,
|
|
||||||
filename=f"manual_{i}",
|
|
||||||
extension=".png",
|
|
||||||
source_path=""
|
|
||||||
))
|
|
||||||
|
|
||||||
if 像素缩放 and manual_images:
|
|
||||||
scaled_manual = []
|
|
||||||
for img_info in manual_images:
|
|
||||||
scaled_img = self.resize_to_megapixels(img_info.image, 分辨率像素)
|
|
||||||
scaled_manual.append(ImageInfo(
|
|
||||||
image=scaled_img,
|
|
||||||
filename=img_info.filename,
|
|
||||||
extension=img_info.extension,
|
|
||||||
source_path=img_info.source_path
|
|
||||||
))
|
|
||||||
manual_images = scaled_manual
|
|
||||||
|
|
||||||
folder_images = self._load_folders(
|
|
||||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
|
||||||
像素缩放, 分辨率像素,
|
|
||||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
|
||||||
)
|
|
||||||
|
|
||||||
pairing_mode = kwargs.get("图片配对模式", "不配对")
|
|
||||||
pairs = self._create_pairs(folder_images, pairing_mode, manual_images if manual_images else None)
|
|
||||||
|
|
||||||
if not pairs:
|
|
||||||
raise ValueError("没有可处理的图片组合,请检查文件夹路径和参考图输入")
|
|
||||||
|
|
||||||
batch_prompts = parse_batch_prompts(提示词)
|
|
||||||
prompts_per_task = None
|
|
||||||
|
|
||||||
if batch_prompts:
|
|
||||||
if len(batch_prompts) != len(pairs):
|
|
||||||
raise ValueError(
|
|
||||||
f"批量提示词数量 ({len(batch_prompts)}) 与任务数量 ({len(pairs)}) 不匹配!\n"
|
|
||||||
f"请确保提示词数量与图片组合数量一致"
|
|
||||||
)
|
|
||||||
prompts_per_task = batch_prompts
|
|
||||||
print(f"全能生图(批量): 批量提示词模式 - {len(batch_prompts)} 个提示词")
|
|
||||||
|
|
||||||
output_folder = 保存路径.strip() if 保存路径 else ""
|
|
||||||
if not output_folder and FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
|
|
||||||
if not output_folder:
|
|
||||||
raise ValueError("无法确定保存路径,请指定保存路径或确保 folder_paths 可用")
|
|
||||||
|
|
||||||
import os
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
print(f"全能生图(批量): 保存路径 → {output_folder}")
|
|
||||||
|
|
||||||
pbar = None
|
|
||||||
if PROGRESS_BAR_AVAILABLE:
|
|
||||||
pbar = ProgressBar(len(pairs))
|
|
||||||
|
|
||||||
def run_async():
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
asyncio.set_event_loop(loop)
|
|
||||||
try:
|
|
||||||
return loop.run_until_complete(
|
|
||||||
self._process_batch_async(
|
|
||||||
pairs=pairs,
|
|
||||||
prompt=提示词,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
output_folder=output_folder,
|
|
||||||
pbar=pbar,
|
|
||||||
prompts_per_task=prompts_per_task,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
loop.close()
|
|
||||||
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
||||||
future = executor.submit(run_async)
|
|
||||||
results = future.result(timeout=3600)
|
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
all_saved_files = []
|
|
||||||
for r in results:
|
|
||||||
all_saved_files.extend(r.get("saved_files", []))
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
print(f"全能生图(批量): 完成!总耗时 {elapsed:.2f}s | 成功: {success_count}/{len(pairs)} | 失败: {fail_count}")
|
|
||||||
|
|
||||||
output_images = []
|
|
||||||
max_output = 10
|
|
||||||
recent_files = all_saved_files[-min(max_output, len(all_saved_files)):]
|
|
||||||
for file_path in recent_files:
|
|
||||||
try:
|
|
||||||
img = Image.open(file_path)
|
|
||||||
output_images.append(img)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"全能生图(批量): 无法加载 {file_path} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
output_images = [placeholder]
|
|
||||||
|
|
||||||
output_tensor = pil_to_tensor(output_images)
|
|
||||||
print(f"全能生图(批量): 共保存 {len(all_saved_files)} 张图片,节点输出最后 {len(output_images)} 张")
|
|
||||||
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"全能生图(批量): ❌ {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
+104
-125
@@ -1,146 +1,125 @@
|
|||||||
"""
|
"""
|
||||||
LoadFile 节点
|
LoadFile 节点(增强版)
|
||||||
ComfyUI 自定义节点,用于加载文件并转换为 FILE 类型数据
|
支持单文件路径和文件夹路径,输出 FILE_LIST 类型供全能LLM等节点使用
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Tuple
|
from typing import Tuple, List
|
||||||
|
|
||||||
from ..utils.file_types import FileData, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMITS
|
from ..utils.file_types import FileData, FileList, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMIT, TOTAL_FILE_SIZE_LIMIT
|
||||||
|
|
||||||
|
|
||||||
class LoadFile:
|
class LoadFile:
|
||||||
"""
|
"""
|
||||||
LoadFile 节点
|
加载文件节点
|
||||||
|
|
||||||
功能:
|
- 单文件路径:加载指定文件
|
||||||
- 从文件系统加载文件
|
- 文件夹路径:加载文件夹内所有支持的文件(非递归)
|
||||||
- 支持 PDF 和 TXT 文件
|
- 两者可同时使用,结果合并输出
|
||||||
- 转换为 FILE 类型数据(包含 base64 编码内容)
|
- 输出 FILE_LIST 类型,可直接连接到全能LLM对话助手
|
||||||
- 验证文件大小和格式
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(cls):
|
def INPUT_TYPES(cls):
|
||||||
"""
|
|
||||||
定义输入参数
|
|
||||||
"""
|
|
||||||
return {
|
return {
|
||||||
"required": {
|
"required": {},
|
||||||
"文件路径": ("STRING", {
|
"optional": {
|
||||||
|
"单文件路径": ("STRING", {
|
||||||
"default": "",
|
"default": "",
|
||||||
"multiline": False
|
"multiline": False,
|
||||||
})
|
"placeholder": "文件完整路径,多个文件用英文逗号分隔",
|
||||||
}
|
}),
|
||||||
|
"文件夹路径": ("STRING", {
|
||||||
|
"default": "",
|
||||||
|
"multiline": False,
|
||||||
|
"placeholder": "文件夹路径,自动读取其中所有支持的文件",
|
||||||
|
}),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
# 返回值类型
|
RETURN_TYPES = ("FILE_LIST", "STRING")
|
||||||
RETURN_TYPES = ("FILE", "STRING")
|
RETURN_NAMES = ("文件列表", "文件信息")
|
||||||
RETURN_NAMES = ("文件", "文件信息")
|
|
||||||
|
|
||||||
# 执行函数名
|
|
||||||
FUNCTION = "load_file"
|
FUNCTION = "load_file"
|
||||||
|
|
||||||
# 节点分类
|
|
||||||
CATEGORY = "file/input"
|
CATEGORY = "file/input"
|
||||||
|
|
||||||
def load_file(self, 文件路径: str) -> Tuple[FileData, str]:
|
def load_file(self, 单文件路径: str = "", 文件夹路径: str = "") -> Tuple[FileList, str]:
|
||||||
"""
|
collected: List[Path] = []
|
||||||
加载文件并转换为 FILE 类型
|
|
||||||
|
# 1. 单文件路径(逗号分隔,支持多个)
|
||||||
Args:
|
if 单文件路径.strip():
|
||||||
文件路径: 文件的完整路径(支持绝对路径和相对路径)
|
for raw in 单文件路径.split(","):
|
||||||
|
p = Path(raw.strip().strip('"').strip("'"))
|
||||||
Returns:
|
if not p.is_absolute():
|
||||||
(FileData, 文件信息预览)
|
p = Path.cwd() / p
|
||||||
|
if not p.exists():
|
||||||
Raises:
|
raise ValueError(f"文件不存在: {p}")
|
||||||
ValueError: 文件不存在、不支持的文件类型或文件过大
|
if not p.is_file():
|
||||||
"""
|
raise ValueError(f"路径不是文件: {p}")
|
||||||
try:
|
collected.append(p)
|
||||||
# 清理路径(去除空格和引号)
|
|
||||||
file_path = 文件路径.strip().strip('"').strip("'")
|
# 2. 文件夹路径
|
||||||
|
if 文件夹路径.strip():
|
||||||
if not file_path:
|
folder = Path(文件夹路径.strip().strip('"').strip("'"))
|
||||||
raise ValueError("文件路径不能为空")
|
if not folder.is_absolute():
|
||||||
|
folder = Path.cwd() / folder
|
||||||
# 转换为 Path 对象
|
if not folder.exists():
|
||||||
path = Path(file_path)
|
raise ValueError(f"文件夹不存在: {folder}")
|
||||||
|
if not folder.is_dir():
|
||||||
# 如果是相对路径,转换为绝对路径
|
raise ValueError(f"路径不是文件夹: {folder}")
|
||||||
if not path.is_absolute():
|
for p in sorted(folder.iterdir()):
|
||||||
# 相对于当前工作目录
|
if p.is_file() and p.suffix.lower() in DOCUMENT_MIME_TYPES:
|
||||||
path = Path.cwd() / path
|
collected.append(p)
|
||||||
|
if not collected:
|
||||||
# 验证文件是否存在
|
raise ValueError(f"文件夹中没有支持的文件: {folder}")
|
||||||
if not path.exists():
|
|
||||||
raise ValueError(f"文件不存在: {file_path}")
|
if not collected:
|
||||||
|
raise ValueError("请至少提供一个文件路径或文件夹路径")
|
||||||
if not path.is_file():
|
|
||||||
raise ValueError(f"路径不是文件: {file_path}")
|
# 去重(保持顺序)
|
||||||
|
seen = set()
|
||||||
# 获取文件信息
|
unique: List[Path] = []
|
||||||
extension = path.suffix.lower()
|
for p in collected:
|
||||||
filename = path.stem
|
key = str(p.resolve())
|
||||||
file_size = path.stat().st_size
|
if key not in seen:
|
||||||
|
seen.add(key)
|
||||||
# 验证文件类型
|
unique.append(p)
|
||||||
if extension not in DOCUMENT_MIME_TYPES:
|
|
||||||
supported_types = ", ".join(DOCUMENT_MIME_TYPES.keys())
|
# 大小检查 & 读取
|
||||||
|
total_size = 0
|
||||||
|
file_list: FileList = []
|
||||||
|
info_lines = []
|
||||||
|
|
||||||
|
for p in unique:
|
||||||
|
ext = p.suffix.lower()
|
||||||
|
if ext not in DOCUMENT_MIME_TYPES:
|
||||||
|
print(f"LoadFile: 跳过不支持的文件类型 {p.name}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
file_size = p.stat().st_size
|
||||||
|
if file_size > FILE_SIZE_LIMIT:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"不支持的文件类型: {extension}\n"
|
f"文件 {p.name} 大小 {file_size / 1024 / 1024:.1f}MB 超过单文件 50MB 限制"
|
||||||
f"支持的类型: {supported_types}"
|
|
||||||
)
|
)
|
||||||
|
total_size += file_size
|
||||||
# 获取 MIME 类型
|
if total_size > TOTAL_FILE_SIZE_LIMIT:
|
||||||
mime_type = DOCUMENT_MIME_TYPES[extension]
|
raise ValueError(f"所有文件总大小超过 50MB 限制")
|
||||||
|
|
||||||
# 验证文件大小
|
mime = DOCUMENT_MIME_TYPES[ext]
|
||||||
size_limit = FILE_SIZE_LIMITS.get(extension, 20 * 1024 * 1024)
|
with open(p, "rb") as f:
|
||||||
if file_size > size_limit:
|
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||||
raise ValueError(
|
|
||||||
f"文件过大 ({file_size / 1024 / 1024:.2f}MB),"
|
file_list.append(FileData(
|
||||||
f"最大支持 {size_limit / 1024 / 1024:.0f}MB"
|
path=str(p),
|
||||||
)
|
filename=p.stem,
|
||||||
|
extension=ext,
|
||||||
# 读取文件并转换为 base64
|
mime_type=mime,
|
||||||
print(f"LoadFile: 正在加载文件 {filename}{extension}")
|
data=b64,
|
||||||
print(f"LoadFile: 文件大小 = {file_size / 1024:.2f}KB")
|
size=file_size,
|
||||||
|
))
|
||||||
with open(path, "rb") as f:
|
info_lines.append(f" {p.name} ({file_size / 1024:.1f}KB, {mime})")
|
||||||
file_bytes = f.read()
|
print(f"LoadFile: 加载 {p.name} ({file_size / 1024:.1f}KB)")
|
||||||
|
|
||||||
# Base64 编码
|
info = f"共 {len(file_list)} 个文件,总大小 {total_size / 1024:.1f}KB\n" + "\n".join(info_lines)
|
||||||
b64_str = base64.b64encode(file_bytes).decode("utf-8")
|
return (file_list, info)
|
||||||
|
|
||||||
# 创建 FileData 对象
|
|
||||||
file_data = FileData(
|
|
||||||
path=str(path),
|
|
||||||
filename=filename,
|
|
||||||
extension=extension,
|
|
||||||
mime_type=mime_type,
|
|
||||||
data=b64_str,
|
|
||||||
size=file_size
|
|
||||||
)
|
|
||||||
|
|
||||||
# 生成文件信息预览
|
|
||||||
file_info = (
|
|
||||||
f"文件名: {filename}{extension}\n"
|
|
||||||
f"类型: {mime_type}\n"
|
|
||||||
f"大小: {file_size / 1024:.2f}KB\n"
|
|
||||||
f"路径: {path}"
|
|
||||||
)
|
|
||||||
|
|
||||||
print(f"LoadFile: 加载成功")
|
|
||||||
|
|
||||||
return (file_data, file_info)
|
|
||||||
|
|
||||||
except ValueError as e:
|
|
||||||
print(f"LoadFile: 输入错误 - {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"LoadFile: 未知错误 - {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|||||||
+56
-147
@@ -52,7 +52,7 @@ except ImportError:
|
|||||||
# ============================================================================
|
# ============================================================================
|
||||||
# 是否启用调试日志(打印完整的 API 响应内容)
|
# 是否启用调试日志(打印完整的 API 响应内容)
|
||||||
# 设置为 True 以启用调试日志,False 以禁用
|
# 设置为 True 以启用调试日志,False 以禁用
|
||||||
DEBUG_LOG_ENABLED = False
|
DEBUG_LOG_ENABLED = True
|
||||||
# 是否启用请求体日志(打印发送给 API 的请求体,base64 图片数据将自动截断)
|
# 是否启用请求体日志(打印发送给 API 的请求体,base64 图片数据将自动截断)
|
||||||
# 设置为 True 以启用请求体日志,False 以禁用
|
# 设置为 True 以启用请求体日志,False 以禁用
|
||||||
REQUEST_LOG_ENABLED = False
|
REQUEST_LOG_ENABLED = False
|
||||||
@@ -177,18 +177,6 @@ class NanoBananaPro:
|
|||||||
"max": 1000,
|
"max": 1000,
|
||||||
"step": 1
|
"step": 1
|
||||||
}),
|
}),
|
||||||
"像素缩放": ("BOOLEAN", {
|
|
||||||
"default": True,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
}),
|
|
||||||
"分辨率像素": ("FLOAT", {
|
|
||||||
"default": 1.0,
|
|
||||||
"min": 0.1,
|
|
||||||
"max": 100.0,
|
|
||||||
"step": 0.1,
|
|
||||||
"display": "number"
|
|
||||||
}),
|
|
||||||
"谷歌搜索(联网)": (["关闭", "打开"], {
|
"谷歌搜索(联网)": (["关闭", "打开"], {
|
||||||
"default": "关闭"
|
"default": "关闭"
|
||||||
}),
|
}),
|
||||||
@@ -199,11 +187,6 @@ class NanoBananaPro:
|
|||||||
"default": 0,
|
"default": 0,
|
||||||
"min": 0,
|
"min": 0,
|
||||||
"max": 0xffffffffffffffff
|
"max": 0xffffffffffffffff
|
||||||
}),
|
|
||||||
"跳过错误": ("BOOLEAN", {
|
|
||||||
"default": False,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
})
|
})
|
||||||
},
|
},
|
||||||
"optional": optional_inputs
|
"optional": optional_inputs
|
||||||
@@ -309,14 +292,16 @@ class NanoBananaPro:
|
|||||||
global_task_index: int,
|
global_task_index: int,
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
|
save_to_disk: bool = True,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
"""执行单个生成任务,生成后立即保存到磁盘"""
|
"""执行单个生成任务"""
|
||||||
result = {
|
result = {
|
||||||
"global_task_index": global_task_index,
|
"global_task_index": global_task_index,
|
||||||
"prompt": prompt,
|
"prompt": prompt,
|
||||||
"success": False,
|
"success": False,
|
||||||
"generated_count": 0,
|
"generated_count": 0,
|
||||||
"saved_files": [],
|
"saved_files": [],
|
||||||
|
"output_images": [],
|
||||||
"error": None
|
"error": None
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -335,20 +320,23 @@ class NanoBananaPro:
|
|||||||
)
|
)
|
||||||
if gen_result:
|
if gen_result:
|
||||||
images_list, _ = gen_result
|
images_list, _ = gen_result
|
||||||
for gen_img in images_list:
|
if save_to_disk:
|
||||||
output_path = generate_timestamp_filename(
|
for gen_img in images_list:
|
||||||
output_folder=output_folder,
|
output_path = generate_timestamp_filename(
|
||||||
extension=".png"
|
output_folder=output_folder,
|
||||||
)
|
extension=".png"
|
||||||
save_image(gen_img, output_path)
|
)
|
||||||
result["saved_files"].append(output_path)
|
save_image(gen_img, output_path)
|
||||||
gen_img = None # 释放内存
|
result["saved_files"].append(output_path)
|
||||||
|
gen_img = None
|
||||||
|
else:
|
||||||
|
result["output_images"] = images_list
|
||||||
|
|
||||||
result["success"] = True
|
result["success"] = True
|
||||||
result["generated_count"] = len(images_list)
|
result["generated_count"] = len(images_list)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
result["error"] = str(e)
|
result["error"] = str(e)
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
async def _process_batch_async(
|
async def _process_batch_async(
|
||||||
@@ -363,8 +351,9 @@ class NanoBananaPro:
|
|||||||
pbar=None,
|
pbar=None,
|
||||||
enable_grounding: bool = False,
|
enable_grounding: bool = False,
|
||||||
enable_image_search: bool = False,
|
enable_image_search: bool = False,
|
||||||
|
save_to_disk: bool = True,
|
||||||
) -> List[dict]:
|
) -> List[dict]:
|
||||||
"""异步批量处理:每个提示词独立调用 API,生成后立即写磁盘"""
|
"""异步批量处理:每个提示词独立调用 API"""
|
||||||
# 构建任务列表:(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
# 构建任务列表:(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
||||||
tasks_def = []
|
tasks_def = []
|
||||||
for p_idx, prompt in enumerate(prompts):
|
for p_idx, prompt in enumerate(prompts):
|
||||||
@@ -373,9 +362,8 @@ class NanoBananaPro:
|
|||||||
|
|
||||||
total_tasks = len(tasks_def)
|
total_tasks = len(tasks_def)
|
||||||
num_prompts = len(prompts)
|
num_prompts = len(prompts)
|
||||||
print(f"Nano Banana Pro: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
|
||||||
|
max_concurrent = 50
|
||||||
max_concurrent = 10
|
|
||||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
@@ -383,7 +371,7 @@ class NanoBananaPro:
|
|||||||
success_count = 0
|
success_count = 0
|
||||||
fail_count = 0
|
fail_count = 0
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
async with aiohttp.ClientSession(connector=connector) as session:
|
||||||
for batch_idx in range(num_batches):
|
for batch_idx in range(num_batches):
|
||||||
@@ -405,6 +393,7 @@ class NanoBananaPro:
|
|||||||
global_task_index=i,
|
global_task_index=i,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
|
save_to_disk=save_to_disk,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
tasks.append(task)
|
tasks.append(task)
|
||||||
@@ -454,23 +443,18 @@ class NanoBananaPro:
|
|||||||
宽高比: str,
|
宽高比: str,
|
||||||
分辨率: str,
|
分辨率: str,
|
||||||
生图数量: int,
|
生图数量: int,
|
||||||
像素缩放: bool,
|
|
||||||
分辨率像素: float,
|
|
||||||
seed: int,
|
seed: int,
|
||||||
跳过错误: bool = False,
|
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> Tuple[torch.Tensor]:
|
) -> Tuple[torch.Tensor]:
|
||||||
"""
|
"""
|
||||||
生成图像
|
生成图像
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
prompt: 提示词
|
prompt: 提示词
|
||||||
模型: 模型名称
|
模型: 模型名称
|
||||||
宽高比: 宽高比
|
宽高比: 宽高比
|
||||||
分辨率: 分辨率
|
分辨率: 分辨率
|
||||||
生图数量: 批次大小
|
生图数量: 批次大小
|
||||||
像素缩放: 是否启用像素缩放
|
|
||||||
分辨率像素: 目标像素数(百万像素)
|
|
||||||
seed: 随机种子
|
seed: 随机种子
|
||||||
**kwargs: 搜索开关(谷歌搜索(联网)/ 图片搜索(联网))及动态参考图输入 (参考图1-9)
|
**kwargs: 搜索开关(谷歌搜索(联网)/ 图片搜索(联网))及动态参考图输入 (参考图1-9)
|
||||||
注:两个搜索参数名含全角括号,不能作为 Python 形参,从 kwargs 中提取
|
注:两个搜索参数名含全角括号,不能作为 Python 形参,从 kwargs 中提取
|
||||||
@@ -550,15 +534,7 @@ class NanoBananaPro:
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"输入图像数量 {len(input_images)} 超过限制 14 张,请减少输入图像数量"
|
f"输入图像数量 {len(input_images)} 超过限制 14 张,请减少输入图像数量"
|
||||||
)
|
)
|
||||||
|
|
||||||
# 应用像素缩放(如果启用)
|
|
||||||
if input_images and 像素缩放:
|
|
||||||
scaled_images = []
|
|
||||||
for img in input_images:
|
|
||||||
scaled = self.resize_to_megapixels(img, 分辨率像素)
|
|
||||||
scaled_images.append(scaled)
|
|
||||||
input_images = scaled_images
|
|
||||||
|
|
||||||
# 解析批量提示词
|
# 解析批量提示词
|
||||||
batch_prompts = parse_batch_prompts(prompt)
|
batch_prompts = parse_batch_prompts(prompt)
|
||||||
|
|
||||||
@@ -602,14 +578,13 @@ class NanoBananaPro:
|
|||||||
nonlocal success_count, fail_count
|
nonlocal success_count, fail_count
|
||||||
if success:
|
if success:
|
||||||
success_count += 1
|
success_count += 1
|
||||||
print(f"Nano Banana Pro: 任务 {current}/{total} 成功 ✓")
|
|
||||||
else:
|
else:
|
||||||
fail_count += 1
|
fail_count += 1
|
||||||
|
|
||||||
# 更新 ComfyUI 原生进度条
|
# 更新 ComfyUI 原生进度条
|
||||||
if pbar is not None:
|
if pbar is not None:
|
||||||
pbar.update(1)
|
pbar.update(1)
|
||||||
|
|
||||||
# 内存监控(每完成10个任务检查一次)
|
# 内存监控(每完成10个任务检查一次)
|
||||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
||||||
import gc
|
import gc
|
||||||
@@ -627,21 +602,10 @@ class NanoBananaPro:
|
|||||||
num_prompts = len(batch_prompts)
|
num_prompts = len(batch_prompts)
|
||||||
total_images = num_prompts * 生图数量
|
total_images = num_prompts * 生图数量
|
||||||
|
|
||||||
# ===== 批量提示词模式:异步并发+磁盘保存 =====
|
# ===== 批量提示词模式:异步并发,内存输出 =====
|
||||||
if pbar is not None:
|
if pbar is not None:
|
||||||
pbar = ProgressBar(total_images)
|
pbar = ProgressBar(total_images)
|
||||||
|
|
||||||
# 确定保存路径
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
|
|
||||||
import os
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
def run_async_in_thread():
|
def run_async_in_thread():
|
||||||
loop = asyncio.new_event_loop()
|
loop = asyncio.new_event_loop()
|
||||||
asyncio.set_event_loop(loop)
|
asyncio.set_event_loop(loop)
|
||||||
@@ -654,10 +618,11 @@ class NanoBananaPro:
|
|||||||
aspect_ratio=宽高比,
|
aspect_ratio=宽高比,
|
||||||
images_per_prompt=生图数量,
|
images_per_prompt=生图数量,
|
||||||
input_images=input_images,
|
input_images=input_images,
|
||||||
output_folder=output_folder,
|
output_folder="",
|
||||||
pbar=pbar,
|
pbar=pbar,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
|
save_to_disk=False,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
@@ -666,23 +631,20 @@ class NanoBananaPro:
|
|||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||||
future = executor.submit(run_async_in_thread)
|
future = executor.submit(run_async_in_thread)
|
||||||
try:
|
try:
|
||||||
results = future.result(timeout=3600)
|
results = future.result(timeout=900)
|
||||||
except TimeoutError:
|
except TimeoutError:
|
||||||
raise RuntimeError("任务执行超时(1小时),请减少提示词数量或检查网络连接")
|
raise RuntimeError("任务执行超时(900秒),请减少提示词数量或检查网络连接")
|
||||||
|
|
||||||
# 统计结果
|
# 统计结果
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
success_count = sum(1 for r in results if r.get("success", False))
|
||||||
fail_count = len(results) - success_count
|
fail_count = len(results) - success_count
|
||||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||||
all_saved_files = []
|
|
||||||
for r in results:
|
|
||||||
all_saved_files.extend(r.get("saved_files", []))
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
elapsed = time.time() - start_time
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||||
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}")
|
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}")
|
||||||
|
|
||||||
# 失败详情
|
# 失败详情
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
failed_results = [r for r in results if not r.get("success", False)]
|
||||||
if failed_results:
|
if failed_results:
|
||||||
@@ -691,24 +653,17 @@ class NanoBananaPro:
|
|||||||
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
||||||
error_msg = fr.get("error", "未知错误")
|
error_msg = fr.get("error", "未知错误")
|
||||||
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
||||||
|
|
||||||
# 从磁盘加载最后 10 张图片
|
# 收集内存中的图像
|
||||||
output_images = []
|
output_images = []
|
||||||
max_output_images = 10
|
for r in results:
|
||||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
output_images.extend(r.get("output_images", []))
|
||||||
for file_path in recent_files:
|
|
||||||
try:
|
|
||||||
img = Image.open(file_path)
|
|
||||||
output_images.append(img)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Nano Banana Pro: 无法加载 {file_path} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
if not output_images:
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||||
output_images = [placeholder]
|
output_images = [placeholder]
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||||
print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
|
|
||||||
import gc
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
@@ -716,7 +671,7 @@ class NanoBananaPro:
|
|||||||
else:
|
else:
|
||||||
# 单提示词模式
|
# 单提示词模式
|
||||||
if 生图数量 == 1:
|
if 生图数量 == 1:
|
||||||
# 单张:同步生成 + 保存到磁盘 + 输出 tensor
|
# 单张:同步生成,输出 tensor
|
||||||
generated_images = self.client.generate_sync(
|
generated_images = self.client.generate_sync(
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
model=模型,
|
model=模型,
|
||||||
@@ -730,35 +685,12 @@ class NanoBananaPro:
|
|||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
)
|
)
|
||||||
# 单张:保存到磁盘
|
|
||||||
import os
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
for gen_img in generated_images:
|
|
||||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
else:
|
else:
|
||||||
# 多张:异步并发 + 磁盘保存(与批量提示词逻辑一致)
|
# 多张:异步并发,内存输出
|
||||||
print(f"Nano Banana Pro: 单提示词×{生图数量}张 → 异步并发模式")
|
|
||||||
|
|
||||||
if pbar is not None:
|
if pbar is not None:
|
||||||
pbar = ProgressBar(生图数量)
|
pbar = ProgressBar(生图数量)
|
||||||
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
|
|
||||||
import os
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
def run_async_in_thread():
|
def run_async_in_thread():
|
||||||
loop = asyncio.new_event_loop()
|
loop = asyncio.new_event_loop()
|
||||||
asyncio.set_event_loop(loop)
|
asyncio.set_event_loop(loop)
|
||||||
@@ -771,10 +703,11 @@ class NanoBananaPro:
|
|||||||
aspect_ratio=宽高比,
|
aspect_ratio=宽高比,
|
||||||
images_per_prompt=生图数量,
|
images_per_prompt=生图数量,
|
||||||
input_images=input_images,
|
input_images=input_images,
|
||||||
output_folder=output_folder,
|
output_folder="",
|
||||||
pbar=pbar,
|
pbar=pbar,
|
||||||
enable_grounding=enable_grounding,
|
enable_grounding=enable_grounding,
|
||||||
enable_image_search=enable_image_search,
|
enable_image_search=enable_image_search,
|
||||||
|
save_to_disk=False,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
@@ -783,16 +716,13 @@ class NanoBananaPro:
|
|||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||||
future = executor.submit(run_async_in_thread)
|
future = executor.submit(run_async_in_thread)
|
||||||
try:
|
try:
|
||||||
results = future.result(timeout=3600)
|
results = future.result(timeout=900)
|
||||||
except TimeoutError:
|
except TimeoutError:
|
||||||
raise RuntimeError("任务执行超时(1小时),请减少生图数量或检查网络连接")
|
raise RuntimeError("任务执行超时(900秒),请减少生图数量或检查网络连接")
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
success_count = sum(1 for r in results if r.get("success", False))
|
||||||
fail_count = len(results) - success_count
|
fail_count = len(results) - success_count
|
||||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||||
all_saved_files = []
|
|
||||||
for r in results:
|
|
||||||
all_saved_files.extend(r.get("saved_files", []))
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
elapsed = time.time() - start_time
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||||
@@ -806,24 +736,16 @@ class NanoBananaPro:
|
|||||||
error_msg = fr.get("error", "未知错误")
|
error_msg = fr.get("error", "未知错误")
|
||||||
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}")
|
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}")
|
||||||
|
|
||||||
# 从磁盘加载最后 10 张图片
|
# 收集内存中的图像
|
||||||
output_images = []
|
output_images = []
|
||||||
max_output_images = 10
|
for r in results:
|
||||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
output_images.extend(r.get("output_images", []))
|
||||||
for file_path in recent_files:
|
|
||||||
try:
|
|
||||||
img = Image.open(file_path)
|
|
||||||
output_images.append(img)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Nano Banana Pro: 无法加载 {file_path} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
if not output_images:
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||||
output_images = [placeholder]
|
output_images = [placeholder]
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||||
print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
# 不生成 prompts_map.txt(单提示词无需映射)
|
|
||||||
|
|
||||||
import gc
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
@@ -851,9 +773,9 @@ class NanoBananaPro:
|
|||||||
|
|
||||||
# 打印最终汇总
|
# 打印最终汇总
|
||||||
if fail_count > 0:
|
if fail_count > 0:
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
print(f"完成!总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
||||||
else:
|
else:
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {len(generated_images)}张")
|
print(f"完成!总耗时 {time_str} | 成功 {len(generated_images)}张")
|
||||||
|
|
||||||
# 最终内存清理
|
# 最终内存清理
|
||||||
import gc
|
import gc
|
||||||
@@ -861,7 +783,7 @@ class NanoBananaPro:
|
|||||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||||
final_memory = process.memory_info().rss / 1024 / 1024
|
final_memory = process.memory_info().rss / 1024 / 1024
|
||||||
print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB")
|
print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB")
|
||||||
|
|
||||||
return (output_tensor,)
|
return (output_tensor,)
|
||||||
|
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
@@ -869,24 +791,12 @@ class NanoBananaPro:
|
|||||||
if str(e) == "未授权!":
|
if str(e) == "未授权!":
|
||||||
print("请联系作者授权后方可使用!")
|
print("请联系作者授权后方可使用!")
|
||||||
raise ValueError("未授权!") from None
|
raise ValueError("未授权!") from None
|
||||||
if 跳过错误:
|
|
||||||
print("Nano Banana Pro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise ValueError(str(e)) from None
|
raise ValueError(str(e)) from None
|
||||||
|
|
||||||
except RuntimeError as e:
|
except RuntimeError as e:
|
||||||
if 跳过错误:
|
|
||||||
print("Nano Banana Pro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise RuntimeError(str(e)) from None
|
raise RuntimeError(str(e)) from None
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
if 跳过错误:
|
|
||||||
print("Nano Banana Pro: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise type(e)(str(e)) from None
|
raise type(e)(str(e)) from None
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
@@ -898,8 +808,7 @@ class NanoBananaPro:
|
|||||||
print(f"Nano Banana Pro: {balance_info}")
|
print(f"Nano Banana Pro: {balance_info}")
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
# 最终内存清理
|
# 最终内存清理
|
||||||
import gc
|
import gc
|
||||||
gc.collect()
|
gc.collect()
|
||||||
print(f"Nano Banana Pro: 最终内存清理完成")
|
|
||||||
@@ -1,656 +0,0 @@
|
|||||||
"""
|
|
||||||
Nano Banana v2 节点
|
|
||||||
NanoBananaPro 的完全复刻,唯一改动:
|
|
||||||
|
|
||||||
将原来 9 个独立「参考图1~9」输入端
|
|
||||||
改为 1 个「参考图」输入端(可选),配合「加载图像(批量)」节点使用。
|
|
||||||
|
|
||||||
「加载图像(批量)」输出 is_output_list=True(list[Tensor]),
|
|
||||||
本节点声明 INPUT_IS_LIST = True 来整体接收该列表,
|
|
||||||
然后在 generate() 开头对所有参数统一解包,其余业务逻辑与原节点完全一致。
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
import gc
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import random
|
|
||||||
import asyncio
|
|
||||||
import aiohttp
|
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
|
||||||
from typing import Optional, Tuple, List
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
|
||||||
from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
|
||||||
from ..clients.gemini_client import GeminiAPIClient
|
|
||||||
from ..models_config import (
|
|
||||||
get_enabled_models, get_model_description,
|
|
||||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
|
||||||
get_model_supported_resolutions, get_all_supported_resolutions
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
import folder_paths
|
|
||||||
FOLDER_PATHS_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
from comfy.utils import ProgressBar
|
|
||||||
PROGRESS_BAR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
PROGRESS_BAR_AVAILABLE = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
import psutil
|
|
||||||
MEMORY_MONITOR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
MEMORY_MONITOR_AVAILABLE = False
|
|
||||||
|
|
||||||
DEBUG_LOG_ENABLED = False
|
|
||||||
REQUEST_LOG_ENABLED = False
|
|
||||||
|
|
||||||
_NODE = "Nano Banana v2"
|
|
||||||
|
|
||||||
|
|
||||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
|
||||||
"""
|
|
||||||
将 PIL Image 列表转换为 ComfyUI tensor,安全处理多张不同尺寸的情况。
|
|
||||||
|
|
||||||
ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C],要求 batch 内所有图尺寸相同。
|
|
||||||
当 API 返回多张不同分辨率的图时(主图 + 附图),直接 stack 会崩溃。
|
|
||||||
|
|
||||||
策略:
|
|
||||||
- 所有图均已按原始分辨率保存到磁盘(调用此函数前已完成)
|
|
||||||
- 以第一张图的尺寸为基准,只将尺寸相同的图纳入 tensor 输出
|
|
||||||
- 尺寸不同的图跳过(不 resize、不丢弃磁盘文件),并打印日志提示
|
|
||||||
- 若没有任何图与第一张尺寸相同(极罕见),则只输出第一张
|
|
||||||
"""
|
|
||||||
if not images:
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return pil_to_tensor([placeholder])
|
|
||||||
|
|
||||||
base_size = images[0].size # PIL size = (W, H)
|
|
||||||
matched = [img for img in images if img.size == base_size]
|
|
||||||
skipped = [img for img in images if img.size != base_size]
|
|
||||||
|
|
||||||
if skipped:
|
|
||||||
sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped)
|
|
||||||
print(
|
|
||||||
f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str}),"
|
|
||||||
f"已按原始分辨率保存到磁盘,tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 "
|
|
||||||
f"({base_size[0]}×{base_size[1]})"
|
|
||||||
)
|
|
||||||
|
|
||||||
return pil_to_tensor(matched if matched else [images[0]])
|
|
||||||
|
|
||||||
|
|
||||||
class NanaBananaV2:
|
|
||||||
"""
|
|
||||||
Nano Banana v2
|
|
||||||
|
|
||||||
与 NanoBananaPro 完全一致,参考图输入方式不同:
|
|
||||||
- 原版:9 个独立可选端口(参考图1~9)
|
|
||||||
- v2:1 个可选端口「参考图」,配合「加载图像(批量)」可传入任意数量图片
|
|
||||||
"""
|
|
||||||
|
|
||||||
ASPECT_RATIOS = [
|
|
||||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
|
||||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
|
||||||
"1:4", "4:1", "1:8", "8:1"
|
|
||||||
]
|
|
||||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
self.client = None
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
enabled_models = get_enabled_models()
|
|
||||||
if not enabled_models:
|
|
||||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
|
||||||
|
|
||||||
all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS
|
|
||||||
all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS
|
|
||||||
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt": ("STRING", {
|
|
||||||
"default": "一个中国女子的OOTD",
|
|
||||||
"multiline": True
|
|
||||||
}),
|
|
||||||
"模型": (enabled_models, {"default": enabled_models[0]}),
|
|
||||||
"宽高比": (all_aspect_ratios, {"default": "1:1"}),
|
|
||||||
"分辨率": (all_resolutions, {"default": "2K"}),
|
|
||||||
"生图数量": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
|
||||||
"像素缩放": ("BOOLEAN", {"default": True, "label_on": "打开", "label_off": "关闭"}),
|
|
||||||
"分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}),
|
|
||||||
"谷歌搜索(联网)": (["关闭", "打开"], {"default": "关闭"}),
|
|
||||||
"图片搜索(联网)": (["关闭", "打开"], {"default": "关闭"}),
|
|
||||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
# 单个参考图端口,接受普通 IMAGE 或「加载图像(批量)」输出的列表
|
|
||||||
"参考图": ("IMAGE",),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("IMAGE",)
|
|
||||||
RETURN_NAMES = ("输出图像",)
|
|
||||||
FUNCTION = "generate"
|
|
||||||
CATEGORY = "image/generation"
|
|
||||||
|
|
||||||
# 声明 INPUT_IS_LIST,使 ComfyUI 将「加载图像(批量)」的 list[Tensor]
|
|
||||||
# 整体传入而非逐张迭代执行,同时其余所有参数也会被包进 list,需解包。
|
|
||||||
INPUT_IS_LIST = True
|
|
||||||
|
|
||||||
# ------------------------------------------------------------------ #
|
|
||||||
# 以下方法与 NanoBananaPro 完全相同,仅 generate() 开头增加了解包逻辑
|
|
||||||
# ------------------------------------------------------------------ #
|
|
||||||
|
|
||||||
def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image:
|
|
||||||
current_pixels = image.width * image.height
|
|
||||||
target_pixels = int(target_megapixels * 1_000_000)
|
|
||||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
|
||||||
return image
|
|
||||||
scale = (target_pixels / current_pixels) ** 0.5
|
|
||||||
new_width = max(1, int(image.width * scale))
|
|
||||||
new_height = max(1, int(image.height * scale))
|
|
||||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
||||||
|
|
||||||
async def _generate_single_task(
|
|
||||||
self,
|
|
||||||
session: aiohttp.ClientSession,
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images: List[Image.Image],
|
|
||||||
output_folder: str,
|
|
||||||
global_task_index: int,
|
|
||||||
enable_grounding: bool = False,
|
|
||||||
enable_image_search: bool = False,
|
|
||||||
) -> dict:
|
|
||||||
result = {
|
|
||||||
"global_task_index": global_task_index,
|
|
||||||
"prompt": prompt,
|
|
||||||
"success": False,
|
|
||||||
"generated_count": 0,
|
|
||||||
"saved_files": [],
|
|
||||||
"error": None
|
|
||||||
}
|
|
||||||
try:
|
|
||||||
gen_result = await self.client.generate_single_async(
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=images if images else None,
|
|
||||||
session=session,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
)
|
|
||||||
if gen_result:
|
|
||||||
images_list, _ = gen_result
|
|
||||||
for gen_img in images_list:
|
|
||||||
output_path = generate_timestamp_filename(
|
|
||||||
output_folder=output_folder,
|
|
||||||
extension=".png"
|
|
||||||
)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
result["saved_files"].append(output_path)
|
|
||||||
gen_img = None
|
|
||||||
result["success"] = True
|
|
||||||
result["generated_count"] = len(images_list)
|
|
||||||
except Exception as e:
|
|
||||||
result["error"] = str(e)
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def _process_batch_async(
|
|
||||||
self,
|
|
||||||
prompts: List[str],
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images_per_prompt: int,
|
|
||||||
input_images: List[Image.Image],
|
|
||||||
output_folder: str,
|
|
||||||
pbar=None,
|
|
||||||
enable_grounding: bool = False,
|
|
||||||
enable_image_search: bool = False,
|
|
||||||
) -> List[dict]:
|
|
||||||
tasks_def = []
|
|
||||||
for p_idx, prompt in enumerate(prompts):
|
|
||||||
for sub_idx in range(images_per_prompt):
|
|
||||||
tasks_def.append((p_idx, sub_idx, prompt))
|
|
||||||
|
|
||||||
total_tasks = len(tasks_def)
|
|
||||||
num_prompts = len(prompts)
|
|
||||||
print(f"{_NODE}: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
|
||||||
|
|
||||||
max_concurrent = 10
|
|
||||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
|
||||||
all_results = []
|
|
||||||
completed = 0
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
|
||||||
for batch_idx in range(num_batches):
|
|
||||||
start_idx = batch_idx * max_concurrent
|
|
||||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
|
||||||
tasks = []
|
|
||||||
for i in range(start_idx, end_idx):
|
|
||||||
_, _, prompt = tasks_def[i]
|
|
||||||
task = asyncio.create_task(
|
|
||||||
self._generate_single_task(
|
|
||||||
session=session,
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
global_task_index=i,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
tasks.append(task)
|
|
||||||
|
|
||||||
batch_results = []
|
|
||||||
for coro in asyncio.as_completed(tasks):
|
|
||||||
result_data = None
|
|
||||||
try:
|
|
||||||
result = await coro
|
|
||||||
if isinstance(result, Exception):
|
|
||||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
|
||||||
else:
|
|
||||||
result_data = result
|
|
||||||
batch_results.append(result_data)
|
|
||||||
except Exception as e:
|
|
||||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""}
|
|
||||||
batch_results.append(result_data)
|
|
||||||
|
|
||||||
completed += 1
|
|
||||||
prompt_snippet = (result_data.get("prompt", "") or "")[:30]
|
|
||||||
if result_data and result_data.get("success", False):
|
|
||||||
success_count += 1
|
|
||||||
count = result_data.get("generated_count", 1)
|
|
||||||
print(f"{_NODE}: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
|
||||||
print(f"{_NODE}: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
|
||||||
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
|
|
||||||
all_results.extend(batch_results)
|
|
||||||
gc.collect()
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
return all_results
|
|
||||||
|
|
||||||
def generate(
|
|
||||||
self,
|
|
||||||
prompt,
|
|
||||||
模型,
|
|
||||||
宽高比,
|
|
||||||
分辨率,
|
|
||||||
生图数量,
|
|
||||||
像素缩放,
|
|
||||||
分辨率像素,
|
|
||||||
**kwargs
|
|
||||||
) -> Tuple[torch.Tensor]:
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# INPUT_IS_LIST=True 时,所有参数均为 list,先统一解包为标量
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
prompt = prompt[0] if isinstance(prompt, list) else prompt
|
|
||||||
模型 = 模型[0] if isinstance(模型, list) else 模型
|
|
||||||
宽高比 = 宽高比[0] if isinstance(宽高比, list) else 宽高比
|
|
||||||
分辨率 = 分辨率[0] if isinstance(分辨率, list) else 分辨率
|
|
||||||
生图数量 = 生图数量[0] if isinstance(生图数量, list) else 生图数量
|
|
||||||
像素缩放 = 像素缩放[0] if isinstance(像素缩放, list) else 像素缩放
|
|
||||||
分辨率像素 = 分辨率像素[0] if isinstance(分辨率像素, list) else 分辨率像素
|
|
||||||
|
|
||||||
# seed 也在 kwargs 里(含全角括号的参数名无法作为形参)
|
|
||||||
seed_raw = kwargs.pop("seed", [0])
|
|
||||||
seed: int = seed_raw[0] if isinstance(seed_raw, list) else seed_raw
|
|
||||||
|
|
||||||
# 搜索开关同理
|
|
||||||
grounding_raw = kwargs.pop("谷歌搜索(联网)", ["关闭"])
|
|
||||||
image_search_raw = kwargs.pop("图片搜索(联网)", ["关闭"])
|
|
||||||
enable_grounding: bool = (grounding_raw[0] if isinstance(grounding_raw, list) else grounding_raw) == "打开"
|
|
||||||
enable_image_search: bool = (image_search_raw[0] if isinstance(image_search_raw, list) else image_search_raw) == "打开"
|
|
||||||
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# 收集参考图:兼容两种来源
|
|
||||||
# 1. 「加载图像(批量)」→ is_output_list=True → list[Tensor]
|
|
||||||
# INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor],需展平
|
|
||||||
# 2. 普通 IMAGE 端口(单 tensor 或 batch tensor)→ list 中只有 1 个元素
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
ref_raw = kwargs.pop("参考图", None)
|
|
||||||
input_images: List[Image.Image] = []
|
|
||||||
|
|
||||||
if ref_raw is not None:
|
|
||||||
# INPUT_IS_LIST 下,可选端口若连接则为 list;元素可能是 Tensor 或 list[Tensor]
|
|
||||||
items = ref_raw if isinstance(ref_raw, list) else [ref_raw]
|
|
||||||
for item in items:
|
|
||||||
if item is None:
|
|
||||||
continue
|
|
||||||
if isinstance(item, list):
|
|
||||||
# 来自 is_output_list 的嵌套 list,继续展平
|
|
||||||
for sub in item:
|
|
||||||
if sub is not None and isinstance(sub, torch.Tensor):
|
|
||||||
input_images.extend(tensor_to_pil(sub))
|
|
||||||
elif isinstance(item, torch.Tensor):
|
|
||||||
input_images.extend(tensor_to_pil(item))
|
|
||||||
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
# 以下逻辑与 NanoBananaPro.generate() 完全一致
|
|
||||||
# ----------------------------------------------------------------
|
|
||||||
start_time = time.time()
|
|
||||||
|
|
||||||
pbar = None
|
|
||||||
if PROGRESS_BAR_AVAILABLE:
|
|
||||||
pbar = ProgressBar(生图数量)
|
|
||||||
|
|
||||||
try:
|
|
||||||
random.seed(seed)
|
|
||||||
np.random.seed(seed % (2 ** 32))
|
|
||||||
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
|
||||||
process = psutil.Process()
|
|
||||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB")
|
|
||||||
|
|
||||||
if self.client is None:
|
|
||||||
try:
|
|
||||||
self.client = GeminiAPIClient()
|
|
||||||
except ValueError as e:
|
|
||||||
raise ValueError(f"初始化失败: {str(e)}")
|
|
||||||
|
|
||||||
# 校验分辨率
|
|
||||||
supported_resolutions = get_model_supported_resolutions(模型)
|
|
||||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
|
||||||
raise ValueError(
|
|
||||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的分辨率:{', '.join(supported_resolutions)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 校验宽高比
|
|
||||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
|
||||||
if supported_ratios and 宽高比 not in supported_ratios:
|
|
||||||
raise ValueError(
|
|
||||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 校验图片搜索与模型兼容性
|
|
||||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = [
|
|
||||||
"nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"
|
|
||||||
]
|
|
||||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
|
||||||
raise ValueError(
|
|
||||||
f"模型 \"{模型}\" 不支持【图片搜索(联网)】功能!"
|
|
||||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 验证输入图像数量上限
|
|
||||||
if len(input_images) > 14:
|
|
||||||
raise ValueError(
|
|
||||||
f"输入图像数量 {len(input_images)} 超过限制 14 张,请减少输入图像数量"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 像素缩放
|
|
||||||
if input_images and 像素缩放:
|
|
||||||
input_images = [self.resize_to_megapixels(img, 分辨率像素) for img in input_images]
|
|
||||||
|
|
||||||
# 解析批量提示词
|
|
||||||
batch_prompts = parse_batch_prompts(prompt)
|
|
||||||
|
|
||||||
# 打印概览
|
|
||||||
grounding_str = ""
|
|
||||||
if enable_image_search:
|
|
||||||
grounding_str = " | 谷歌图片搜索接地"
|
|
||||||
elif enable_grounding:
|
|
||||||
grounding_str = " | 谷歌搜索接地"
|
|
||||||
|
|
||||||
if batch_prompts:
|
|
||||||
num_prompts = len(batch_prompts)
|
|
||||||
total_images = num_prompts * 生图数量
|
|
||||||
mode_str = f"批量提示词模式 ({num_prompts}个提示词)"
|
|
||||||
if input_images:
|
|
||||||
mode_str += f" (输入{len(input_images)}张)"
|
|
||||||
print(f"{_NODE}: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}")
|
|
||||||
if total_images > 100:
|
|
||||||
print(f"⚠️ {_NODE}: 警告!批量生成 {total_images} 张图片,内存占用可能较高")
|
|
||||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
|
||||||
else:
|
|
||||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
|
||||||
print(f"{_NODE}: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}")
|
|
||||||
if 生图数量 > 100:
|
|
||||||
print(f"⚠️ {_NODE}: 警告!批量生成 {生图数量} 张图片,内存占用可能较高")
|
|
||||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
|
||||||
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
|
|
||||||
def progress_callback(current, total, success, error_msg=None):
|
|
||||||
nonlocal success_count, fail_count
|
|
||||||
if success:
|
|
||||||
success_count += 1
|
|
||||||
print(f"{_NODE}: 任务 {current}/{total} 成功 ✓")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
if error_msg:
|
|
||||||
print(f"{_NODE}: 任务 {current}/{total} 失败 ✗")
|
|
||||||
print(f"原始错误详情:\n{error_msg}")
|
|
||||||
else:
|
|
||||||
print(f"{_NODE}: 任务 {current}/{total} 失败 ✗")
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
|
||||||
gc.collect()
|
|
||||||
current_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
memory_increase = current_memory - initial_memory
|
|
||||||
print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
|
||||||
if current_memory > 2000:
|
|
||||||
print(f"⚠️ {_NODE}: 内存使用过高!建议减少生图数量或分批执行")
|
|
||||||
|
|
||||||
def _get_output_folder():
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
folder = folder_paths.get_output_directory()
|
|
||||||
return folder
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
|
|
||||||
def run_async_in_thread(coro_fn):
|
|
||||||
def _run():
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
asyncio.set_event_loop(loop)
|
|
||||||
try:
|
|
||||||
return loop.run_until_complete(coro_fn())
|
|
||||||
finally:
|
|
||||||
loop.close()
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
||||||
future = executor.submit(_run)
|
|
||||||
try:
|
|
||||||
return future.result(timeout=3600)
|
|
||||||
except TimeoutError:
|
|
||||||
raise RuntimeError("任务执行超时(1小时),请减少数量或检查网络连接")
|
|
||||||
|
|
||||||
# ── 批量提示词模式 ──────────────────────────────────────
|
|
||||||
if batch_prompts:
|
|
||||||
num_prompts = len(batch_prompts)
|
|
||||||
total_images = num_prompts * 生图数量
|
|
||||||
if pbar is not None:
|
|
||||||
pbar = ProgressBar(total_images)
|
|
||||||
|
|
||||||
output_folder = _get_output_folder()
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
results = run_async_in_thread(lambda: self._process_batch_async(
|
|
||||||
prompts=batch_prompts,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
images_per_prompt=生图数量,
|
|
||||||
input_images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
pbar=pbar,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
))
|
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}")
|
|
||||||
|
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
|
||||||
for fr in failed_results:
|
|
||||||
idx = fr.get("global_task_index", -1) + 1
|
|
||||||
snippet = (fr.get("prompt", "") or "")[:30]
|
|
||||||
print(f" 失败 #{idx}: {snippet}{'...' if len(snippet) >= 30 else ''} → {fr.get('error', '未知错误')}")
|
|
||||||
|
|
||||||
output_images = []
|
|
||||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
|
||||||
try:
|
|
||||||
output_images.append(Image.open(fp))
|
|
||||||
except Exception as e:
|
|
||||||
print(f"{_NODE}: 无法加载 {fp} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
|
||||||
print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
# ── 单提示词模式 ────────────────────────────────────────
|
|
||||||
if 生图数量 == 1:
|
|
||||||
generated_images = self.client.generate_sync(
|
|
||||||
prompt=prompt,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
batch_size=1,
|
|
||||||
images=input_images,
|
|
||||||
progress_callback=progress_callback,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
)
|
|
||||||
output_folder = _get_output_folder()
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
for gen_img in generated_images:
|
|
||||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
else:
|
|
||||||
print(f"{_NODE}: 单提示词×{生图数量}张 → 异步并发模式")
|
|
||||||
if pbar is not None:
|
|
||||||
pbar = ProgressBar(生图数量)
|
|
||||||
|
|
||||||
output_folder = _get_output_folder()
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
results = run_async_in_thread(lambda: self._process_batch_async(
|
|
||||||
prompts=[prompt],
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
images_per_prompt=生图数量,
|
|
||||||
input_images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
pbar=pbar,
|
|
||||||
enable_grounding=enable_grounding,
|
|
||||||
enable_image_search=enable_image_search,
|
|
||||||
))
|
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}")
|
|
||||||
|
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
|
||||||
for fr in failed_results:
|
|
||||||
idx = fr.get("global_task_index", -1) + 1
|
|
||||||
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {fr.get('error', '未知错误')}")
|
|
||||||
|
|
||||||
output_images = []
|
|
||||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
|
||||||
try:
|
|
||||||
output_images.append(Image.open(fp))
|
|
||||||
except Exception as e:
|
|
||||||
print(f"{_NODE}: 无法加载 {fp} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
|
||||||
print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
# 单张同步模式的输出路径(生图数量==1 走到这里)
|
|
||||||
max_output_images = 20
|
|
||||||
if len(generated_images) > max_output_images:
|
|
||||||
print(f"{_NODE}: 生成 {len(generated_images)} 张图片,限制输出前 {max_output_images} 张到ComfyUI")
|
|
||||||
output_images = generated_images[:max_output_images]
|
|
||||||
else:
|
|
||||||
output_images = generated_images
|
|
||||||
|
|
||||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
if fail_count > 0:
|
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
|
||||||
else:
|
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {len(generated_images)}张")
|
|
||||||
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
except ValueError as e:
|
|
||||||
if str(e) == "未授权!":
|
|
||||||
print("请联系作者授权后方可使用!")
|
|
||||||
raise ValueError("未授权!") from None
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_msg}")
|
|
||||||
raise ValueError(error_msg) from None
|
|
||||||
|
|
||||||
except RuntimeError as e:
|
|
||||||
error_full = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_full}")
|
|
||||||
raise RuntimeError(error_full) from None
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"{_NODE}: ❌ {error_msg}")
|
|
||||||
raise type(e)(error_msg) from None
|
|
||||||
|
|
||||||
finally:
|
|
||||||
if self.client is not None:
|
|
||||||
try:
|
|
||||||
balance_data = self.client.query_balance_sync()
|
|
||||||
balance_info = self.client.format_balance_info(balance_data)
|
|
||||||
print(f"{_NODE}: {balance_info}")
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
gc.collect()
|
|
||||||
@@ -1,844 +0,0 @@
|
|||||||
"""
|
|
||||||
全能生图 节点
|
|
||||||
ComfyUI 自定义节点,用于调用 Gemini 模型生成图像
|
|
||||||
"""
|
|
||||||
|
|
||||||
import time
|
|
||||||
import math
|
|
||||||
import random
|
|
||||||
import asyncio
|
|
||||||
import aiohttp
|
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
|
||||||
from typing import Optional, Tuple, List
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
|
||||||
from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
|
||||||
from ..clients.openai_client import OpenAIAPIClient
|
|
||||||
from ..models_config import (
|
|
||||||
get_enabled_models, get_model_description,
|
|
||||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
|
||||||
get_model_supported_resolutions, get_all_supported_resolutions
|
|
||||||
)
|
|
||||||
|
|
||||||
# 检查 folder_paths 是否可用
|
|
||||||
try:
|
|
||||||
import folder_paths
|
|
||||||
FOLDER_PATHS_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
|
||||||
|
|
||||||
# 导入 ComfyUI 原生进度条
|
|
||||||
try:
|
|
||||||
from comfy.utils import ProgressBar
|
|
||||||
PROGRESS_BAR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
PROGRESS_BAR_AVAILABLE = False
|
|
||||||
print("⚠️ 全能生图: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
|
||||||
|
|
||||||
# 内存监控(可选)
|
|
||||||
try:
|
|
||||||
import psutil
|
|
||||||
MEMORY_MONITOR_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
MEMORY_MONITOR_AVAILABLE = False
|
|
||||||
print("⚠️ 全能生图: psutil 不可用,内存监控功能禁用")
|
|
||||||
|
|
||||||
# ============================================================================
|
|
||||||
# 调试日志配置
|
|
||||||
# ============================================================================
|
|
||||||
# 是否启用调试日志(打印完整的 API 响应内容)
|
|
||||||
# 设置为 True 以启用调试日志,False 以禁用
|
|
||||||
DEBUG_LOG_ENABLED = False
|
|
||||||
# 是否启用请求体日志(打印发送给 API 的请求体,base64 图片数据将自动截断)
|
|
||||||
# 设置为 True 以启用请求体日志,False 以禁用
|
|
||||||
REQUEST_LOG_ENABLED = False
|
|
||||||
# ============================================================================
|
|
||||||
|
|
||||||
|
|
||||||
class QuanNengShengTu:
|
|
||||||
"""
|
|
||||||
全能生图 节点
|
|
||||||
|
|
||||||
功能:
|
|
||||||
- 文生图:基于提示词生成图像
|
|
||||||
- 图生图:基于输入图像和提示词生成新图像
|
|
||||||
- 批量生成:支持并发生成多张图像
|
|
||||||
|
|
||||||
注意:
|
|
||||||
- 支持的模型列表从 models_config.py 动态加载
|
|
||||||
- 要添加/禁用模型,请编辑 models_config.py 文件
|
|
||||||
"""
|
|
||||||
|
|
||||||
# 支持的模型列表(从配置文件动态加载)
|
|
||||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
|
||||||
|
|
||||||
# 支持的宽高比列表(全量:所有启用模型的并集,动态加载)
|
|
||||||
ASPECT_RATIOS = [
|
|
||||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
|
||||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
|
||||||
"1:4", "4:1", "1:8", "8:1"
|
|
||||||
]
|
|
||||||
|
|
||||||
# 支持的分辨率列表(全量兜底,实际由 get_all_supported_resolutions() 动态生成)
|
|
||||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
"""初始化节点"""
|
|
||||||
self.client = None
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
"""
|
|
||||||
定义输入参数
|
|
||||||
|
|
||||||
ComfyUI 节点规范:
|
|
||||||
- required: 必选参数
|
|
||||||
- optional: 可选参数
|
|
||||||
"""
|
|
||||||
# 从配置文件动态获取启用的模型列表
|
|
||||||
enabled_models = get_enabled_models()
|
|
||||||
|
|
||||||
# 过滤掉包含"限时特价"的模型
|
|
||||||
enabled_models = [m for m in enabled_models if "限时特价" not in m]
|
|
||||||
|
|
||||||
# 如果没有启用的模型,使用空列表(会导致节点不可用,提示用户配置)
|
|
||||||
if not enabled_models:
|
|
||||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
|
||||||
|
|
||||||
# 动态获取所有启用模型支持的宽高比(去重合并)
|
|
||||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
|
||||||
if not all_aspect_ratios:
|
|
||||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
|
||||||
|
|
||||||
# 动态获取所有启用模型支持的分辨率(去重合并)
|
|
||||||
all_resolutions = get_all_supported_resolutions()
|
|
||||||
if not all_resolutions:
|
|
||||||
all_resolutions = cls.RESOLUTIONS
|
|
||||||
|
|
||||||
# 创建9个独立的图像输入
|
|
||||||
optional_inputs = {}
|
|
||||||
for i in range(1, 10): # 1-9
|
|
||||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
|
||||||
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"提示词": ("STRING", {
|
|
||||||
"default": "一个中国女子的OOTD",
|
|
||||||
"multiline": True
|
|
||||||
}),
|
|
||||||
"模型": (enabled_models, {
|
|
||||||
"default": enabled_models[0]
|
|
||||||
}),
|
|
||||||
"宽高比": (all_aspect_ratios, {
|
|
||||||
"default": "1:1"
|
|
||||||
}),
|
|
||||||
"分辨率": (all_resolutions, {
|
|
||||||
"default": "2K"
|
|
||||||
}),
|
|
||||||
"生图数量": ("INT", {
|
|
||||||
"default": 1,
|
|
||||||
"min": 1,
|
|
||||||
"max": 1000,
|
|
||||||
"step": 1
|
|
||||||
}),
|
|
||||||
"像素缩放": ("BOOLEAN", {
|
|
||||||
"default": True,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
}),
|
|
||||||
"分辨率像素": ("FLOAT", {
|
|
||||||
"default": 1.0,
|
|
||||||
"min": 0.1,
|
|
||||||
"max": 100.0,
|
|
||||||
"step": 0.1,
|
|
||||||
"display": "number"
|
|
||||||
}),
|
|
||||||
"seed": ("INT", {
|
|
||||||
"default": 0,
|
|
||||||
"min": 0,
|
|
||||||
"max": 0xffffffffffffffff
|
|
||||||
}),
|
|
||||||
"跳过错误": ("BOOLEAN", {
|
|
||||||
"default": False,
|
|
||||||
"label_on": "打开",
|
|
||||||
"label_off": "关闭"
|
|
||||||
})
|
|
||||||
},
|
|
||||||
"optional": optional_inputs
|
|
||||||
}
|
|
||||||
|
|
||||||
# 返回值类型
|
|
||||||
RETURN_TYPES = ("IMAGE",)
|
|
||||||
RETURN_NAMES = ("输出图像",)
|
|
||||||
|
|
||||||
# 导入 ComfyUI 的文件夹路径管理
|
|
||||||
try:
|
|
||||||
import folder_paths
|
|
||||||
FOLDER_PATHS_AVAILABLE = True
|
|
||||||
except ImportError:
|
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
|
||||||
|
|
||||||
# 执行函数名
|
|
||||||
FUNCTION = "generate"
|
|
||||||
|
|
||||||
# 节点分类
|
|
||||||
CATEGORY = "image/generation"
|
|
||||||
|
|
||||||
def resize_to_megapixels(
|
|
||||||
self,
|
|
||||||
image: Image.Image,
|
|
||||||
target_megapixels: float
|
|
||||||
) -> Image.Image:
|
|
||||||
"""
|
|
||||||
将图像缩放到指定的总像素数,保持纵横比
|
|
||||||
|
|
||||||
Args:
|
|
||||||
image: PIL Image 对象
|
|
||||||
target_megapixels: 目标像素数(百万像素)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
缩放后的 PIL Image
|
|
||||||
"""
|
|
||||||
# 计算当前像素数
|
|
||||||
current_pixels = image.width * image.height
|
|
||||||
target_pixels = int(target_megapixels * 1_000_000)
|
|
||||||
|
|
||||||
# 如果当前像素数已经接近目标,则不缩放
|
|
||||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
|
||||||
return image
|
|
||||||
|
|
||||||
# 计算缩放比例
|
|
||||||
scale = (target_pixels / current_pixels) ** 0.5
|
|
||||||
|
|
||||||
# 计算新尺寸
|
|
||||||
new_width = int(image.width * scale)
|
|
||||||
new_height = int(image.height * scale)
|
|
||||||
|
|
||||||
# 确保至少为1像素
|
|
||||||
new_width = max(1, new_width)
|
|
||||||
new_height = max(1, new_height)
|
|
||||||
|
|
||||||
# 使用 Lanczos 重采样
|
|
||||||
resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
||||||
|
|
||||||
return resized_image
|
|
||||||
|
|
||||||
def validate_inputs(
|
|
||||||
self,
|
|
||||||
images: Optional[torch.Tensor],
|
|
||||||
batch_size: int
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
验证输入参数
|
|
||||||
|
|
||||||
Args:
|
|
||||||
images: 输入图像张量(可选)
|
|
||||||
batch_size: 批次大小
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: 如果输入参数不合法
|
|
||||||
"""
|
|
||||||
# 检查图像数量
|
|
||||||
if images is not None:
|
|
||||||
num_images = images.shape[0]
|
|
||||||
if num_images > 14:
|
|
||||||
raise ValueError(
|
|
||||||
f"输入图像数量 {num_images} 超过限制 14 张,请减少输入图像数量"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 检查批次大小
|
|
||||||
if batch_size < 1 or batch_size > 1000:
|
|
||||||
raise ValueError(
|
|
||||||
f"批次大小 {batch_size} 超出范围 [1, 1000]"
|
|
||||||
)
|
|
||||||
|
|
||||||
async def _generate_single_task(
|
|
||||||
self,
|
|
||||||
session: aiohttp.ClientSession,
|
|
||||||
prompt: str,
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images: List[Image.Image],
|
|
||||||
output_folder: str,
|
|
||||||
global_task_index: int,
|
|
||||||
) -> dict:
|
|
||||||
"""执行单个生成任务,生成后立即保存到磁盘"""
|
|
||||||
result = {
|
|
||||||
"global_task_index": global_task_index,
|
|
||||||
"prompt": prompt,
|
|
||||||
"success": False,
|
|
||||||
"generated_count": 0,
|
|
||||||
"saved_files": [],
|
|
||||||
"error": None
|
|
||||||
}
|
|
||||||
|
|
||||||
try:
|
|
||||||
gen_result = await self.client.generate_single_async(
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=images if images else None,
|
|
||||||
session=session,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=False,
|
|
||||||
enable_image_search=False
|
|
||||||
)
|
|
||||||
if gen_result:
|
|
||||||
images_list, _ = gen_result
|
|
||||||
for gen_img in images_list:
|
|
||||||
output_path = generate_timestamp_filename(
|
|
||||||
output_folder=output_folder,
|
|
||||||
extension=".png"
|
|
||||||
)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
result["saved_files"].append(output_path)
|
|
||||||
gen_img = None # 释放内存
|
|
||||||
|
|
||||||
result["success"] = True
|
|
||||||
result["generated_count"] = len(images_list)
|
|
||||||
except Exception as e:
|
|
||||||
result["error"] = str(e)
|
|
||||||
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def _process_batch_async(
|
|
||||||
self,
|
|
||||||
prompts: List[str],
|
|
||||||
model: str,
|
|
||||||
resolution: str,
|
|
||||||
aspect_ratio: str,
|
|
||||||
images_per_prompt: int,
|
|
||||||
input_images: List[Image.Image],
|
|
||||||
output_folder: str,
|
|
||||||
pbar=None,
|
|
||||||
) -> List[dict]:
|
|
||||||
"""异步批量处理:每个提示词独立调用 API,生成后立即写磁盘"""
|
|
||||||
# 构建任务列表:(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
|
||||||
tasks_def = []
|
|
||||||
for p_idx, prompt in enumerate(prompts):
|
|
||||||
for sub_idx in range(images_per_prompt):
|
|
||||||
tasks_def.append((p_idx, sub_idx, prompt))
|
|
||||||
|
|
||||||
total_tasks = len(tasks_def)
|
|
||||||
num_prompts = len(prompts)
|
|
||||||
print(f"全能生图: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
|
||||||
|
|
||||||
max_concurrent = 10
|
|
||||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
|
||||||
|
|
||||||
all_results = []
|
|
||||||
completed = 0
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
|
|
||||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
|
||||||
|
|
||||||
async with aiohttp.ClientSession(connector=connector) as session:
|
|
||||||
for batch_idx in range(num_batches):
|
|
||||||
start_idx = batch_idx * max_concurrent
|
|
||||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
|
||||||
|
|
||||||
tasks = []
|
|
||||||
for i in range(start_idx, end_idx):
|
|
||||||
_, _, prompt = tasks_def[i]
|
|
||||||
task = asyncio.create_task(
|
|
||||||
self._generate_single_task(
|
|
||||||
session=session,
|
|
||||||
prompt=prompt,
|
|
||||||
model=model,
|
|
||||||
resolution=resolution,
|
|
||||||
aspect_ratio=aspect_ratio,
|
|
||||||
images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
global_task_index=i,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
tasks.append(task)
|
|
||||||
|
|
||||||
batch_results = []
|
|
||||||
for coro in asyncio.as_completed(tasks):
|
|
||||||
result_data = None
|
|
||||||
try:
|
|
||||||
result = await coro
|
|
||||||
if isinstance(result, Exception):
|
|
||||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
|
||||||
else:
|
|
||||||
result_data = result
|
|
||||||
batch_results.append(result_data)
|
|
||||||
except Exception as e:
|
|
||||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""}
|
|
||||||
batch_results.append(result_data)
|
|
||||||
|
|
||||||
completed += 1
|
|
||||||
prompt_snippet = (result_data.get("prompt", "") or "")[:30]
|
|
||||||
|
|
||||||
if result_data and result_data.get("success", False):
|
|
||||||
success_count += 1
|
|
||||||
count = result_data.get("generated_count", 1)
|
|
||||||
print(f"全能生图: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
|
||||||
print(f"全能生图: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
|
||||||
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
|
|
||||||
all_results.extend(batch_results)
|
|
||||||
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
await asyncio.sleep(0.1)
|
|
||||||
|
|
||||||
return all_results
|
|
||||||
|
|
||||||
def generate(
|
|
||||||
self,
|
|
||||||
提示词: str,
|
|
||||||
模型: str,
|
|
||||||
宽高比: str,
|
|
||||||
分辨率: str,
|
|
||||||
生图数量: int,
|
|
||||||
像素缩放: bool,
|
|
||||||
分辨率像素: float,
|
|
||||||
seed: int,
|
|
||||||
跳过错误: bool = False,
|
|
||||||
**kwargs
|
|
||||||
) -> Tuple[torch.Tensor]:
|
|
||||||
"""
|
|
||||||
生成图像
|
|
||||||
|
|
||||||
Args:
|
|
||||||
prompt: 提示词
|
|
||||||
模型: 模型名称
|
|
||||||
宽高比: 宽高比
|
|
||||||
分辨率: 分辨率
|
|
||||||
生图数量: 批次大小
|
|
||||||
像素缩放: 是否启用像素缩放
|
|
||||||
分辨率像素: 目标像素数(百万像素)
|
|
||||||
seed: 随机种子
|
|
||||||
**kwargs: 动态参考图输入 (参考图1-9)
|
|
||||||
|
|
||||||
注意:
|
|
||||||
调试日志功能已移至文件顶部配置,通过修改 DEBUG_LOG_ENABLED 常量控制
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
生成的图像张量 (IMAGE,)
|
|
||||||
"""
|
|
||||||
start_time = time.time()
|
|
||||||
|
|
||||||
# 创建 ComfyUI 原生进度条
|
|
||||||
pbar = None
|
|
||||||
if PROGRESS_BAR_AVAILABLE:
|
|
||||||
pbar = ProgressBar(生图数量)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# 设置随机种子(用于本地随机操作)
|
|
||||||
random.seed(seed)
|
|
||||||
np.random.seed(seed % (2**32))
|
|
||||||
|
|
||||||
# 内存监控初始化
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
|
||||||
import psutil
|
|
||||||
process = psutil.Process()
|
|
||||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
print(f"全能生图: 初始内存使用: {initial_memory:.1f} MB")
|
|
||||||
|
|
||||||
# 初始化 API 客户端
|
|
||||||
if self.client is None:
|
|
||||||
try:
|
|
||||||
self.client = OpenAIAPIClient()
|
|
||||||
except ValueError as e:
|
|
||||||
raise ValueError(f"初始化失败: {str(e)}")
|
|
||||||
|
|
||||||
# 校验分辨率与模型的兼容性
|
|
||||||
supported_resolutions = get_model_supported_resolutions(模型)
|
|
||||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
|
||||||
raise ValueError(
|
|
||||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的分辨率:{', '.join(supported_resolutions)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 校验宽高比与模型的兼容性
|
|
||||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
|
||||||
if supported_ratios and 宽高比 not in supported_ratios:
|
|
||||||
raise ValueError(
|
|
||||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
|
||||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 收集独立输入的参考图
|
|
||||||
input_images = []
|
|
||||||
for i in range(1, 10): # 1-9
|
|
||||||
key = f"参考图{i}"
|
|
||||||
if key in kwargs and kwargs[key] is not None:
|
|
||||||
pil_imgs = tensor_to_pil(kwargs[key])
|
|
||||||
input_images.extend(pil_imgs)
|
|
||||||
|
|
||||||
# 验证输入图像数量
|
|
||||||
if input_images:
|
|
||||||
if len(input_images) > 14:
|
|
||||||
raise ValueError(
|
|
||||||
f"输入图像数量 {len(input_images)} 超过限制 14 张,请减少输入图像数量"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 应用像素缩放(如果启用)
|
|
||||||
if input_images and 像素缩放:
|
|
||||||
scaled_images = []
|
|
||||||
for img in input_images:
|
|
||||||
scaled = self.resize_to_megapixels(img, 分辨率像素)
|
|
||||||
scaled_images.append(scaled)
|
|
||||||
input_images = scaled_images
|
|
||||||
|
|
||||||
# 解析批量提示词
|
|
||||||
batch_prompts = parse_batch_prompts(提示词)
|
|
||||||
|
|
||||||
# 打印首行概览
|
|
||||||
if batch_prompts:
|
|
||||||
# 批量提示词模式
|
|
||||||
num_prompts = len(batch_prompts)
|
|
||||||
total_images = num_prompts * 生图数量
|
|
||||||
mode_str = f"批量提示词模式 ({num_prompts}个提示词)"
|
|
||||||
if input_images:
|
|
||||||
mode_str += f" (输入{len(input_images)}张)"
|
|
||||||
print(f"全能生图: {mode_str} | {分辨率} {宽高比} | 共{total_images}张")
|
|
||||||
|
|
||||||
# 大批量警告
|
|
||||||
if total_images > 100:
|
|
||||||
print(f"⚠️ 全能生图: 警告!批量生成 {total_images} 张图片,内存占用可能较高")
|
|
||||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
|
||||||
else:
|
|
||||||
# 单提示词模式
|
|
||||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
|
||||||
print(f"全能生图: {mode_str} | {分辨率} {宽高比} | {生图数量}张")
|
|
||||||
|
|
||||||
# 大批量警告
|
|
||||||
if 生图数量 > 100:
|
|
||||||
print(f"⚠️ 全能生图: 警告!批量生成 {生图数量} 张图片,内存占用可能较高")
|
|
||||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
|
||||||
|
|
||||||
# 统计变量
|
|
||||||
success_count = 0
|
|
||||||
fail_count = 0
|
|
||||||
|
|
||||||
# 进度回调 - 打印错误信息并更新进度条,添加内存监控
|
|
||||||
def progress_callback(current, total, success, error_msg=None):
|
|
||||||
nonlocal success_count, fail_count
|
|
||||||
if success:
|
|
||||||
success_count += 1
|
|
||||||
print(f"全能生图: 任务 {current}/{total} 成功 ✓")
|
|
||||||
else:
|
|
||||||
fail_count += 1
|
|
||||||
if error_msg:
|
|
||||||
print(f"全能生图: 任务 {current}/{total} 失败 ✗")
|
|
||||||
print(f"原始错误详情:\n{error_msg}")
|
|
||||||
else:
|
|
||||||
print(f"全能生图: 任务 {current}/{total} 失败 ✗")
|
|
||||||
|
|
||||||
# 更新 ComfyUI 原生进度条
|
|
||||||
if pbar is not None:
|
|
||||||
pbar.update(1)
|
|
||||||
|
|
||||||
# 内存监控(每完成10个任务检查一次)
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
|
||||||
import gc
|
|
||||||
gc.collect() # 强制垃圾回收
|
|
||||||
current_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
memory_increase = current_memory - initial_memory
|
|
||||||
print(f"全能生图: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
|
||||||
|
|
||||||
# 内存警告阈值(2GB)
|
|
||||||
if current_memory > 2000:
|
|
||||||
print(f"⚠️ 全能生图: 内存使用过高!建议减少生图数量或分批执行")
|
|
||||||
|
|
||||||
# 根据是否有批量提示词选择生成模式
|
|
||||||
if batch_prompts:
|
|
||||||
num_prompts = len(batch_prompts)
|
|
||||||
total_images = num_prompts * 生图数量
|
|
||||||
|
|
||||||
# ===== 批量提示词模式:异步并发+磁盘保存 =====
|
|
||||||
if pbar is not None:
|
|
||||||
pbar = ProgressBar(total_images)
|
|
||||||
|
|
||||||
# 确定保存路径
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
|
|
||||||
import os
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
def run_async_in_thread():
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
asyncio.set_event_loop(loop)
|
|
||||||
try:
|
|
||||||
return loop.run_until_complete(
|
|
||||||
self._process_batch_async(
|
|
||||||
prompts=batch_prompts,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
images_per_prompt=生图数量,
|
|
||||||
input_images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
pbar=pbar,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
loop.close()
|
|
||||||
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
||||||
future = executor.submit(run_async_in_thread)
|
|
||||||
try:
|
|
||||||
results = future.result(timeout=3600)
|
|
||||||
except TimeoutError:
|
|
||||||
raise RuntimeError("任务执行超时(1小时),请减少提示词数量或检查网络连接")
|
|
||||||
|
|
||||||
# 统计结果
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
|
||||||
all_saved_files = []
|
|
||||||
for r in results:
|
|
||||||
all_saved_files.extend(r.get("saved_files", []))
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}")
|
|
||||||
|
|
||||||
# 失败详情
|
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
|
||||||
if failed_results:
|
|
||||||
for fr in failed_results:
|
|
||||||
idx = fr.get("global_task_index", -1) + 1
|
|
||||||
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
|
||||||
error_msg = fr.get("error", "未知错误")
|
|
||||||
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
|
||||||
|
|
||||||
# 从磁盘加载最后 10 张图片
|
|
||||||
output_images = []
|
|
||||||
max_output_images = 10
|
|
||||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
|
||||||
for file_path in recent_files:
|
|
||||||
try:
|
|
||||||
img = Image.open(file_path)
|
|
||||||
output_images.append(img)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"全能生图: 无法加载 {file_path} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
output_images = [placeholder]
|
|
||||||
|
|
||||||
output_tensor = pil_to_tensor(output_images)
|
|
||||||
print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
else:
|
|
||||||
# 单提示词模式
|
|
||||||
if 生图数量 == 1:
|
|
||||||
# 单张:同步生成 + 保存到磁盘 + 输出 tensor
|
|
||||||
generated_images = self.client.generate_sync(
|
|
||||||
prompt=提示词,
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
batch_size=1,
|
|
||||||
images=input_images,
|
|
||||||
progress_callback=progress_callback,
|
|
||||||
debug=DEBUG_LOG_ENABLED,
|
|
||||||
debug_request=REQUEST_LOG_ENABLED,
|
|
||||||
enable_grounding=False,
|
|
||||||
enable_image_search=False
|
|
||||||
)
|
|
||||||
# 单张:保存到磁盘
|
|
||||||
import os
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
for gen_img in generated_images:
|
|
||||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
|
||||||
save_image(gen_img, output_path)
|
|
||||||
else:
|
|
||||||
# 多张:异步并发 + 磁盘保存(与批量提示词逻辑一致)
|
|
||||||
print(f"全能生图: 单提示词×{生图数量}张 → 异步并发模式")
|
|
||||||
|
|
||||||
if pbar is not None:
|
|
||||||
pbar = ProgressBar(生图数量)
|
|
||||||
|
|
||||||
output_folder = ""
|
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
|
||||||
output_folder = folder_paths.get_output_directory()
|
|
||||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
|
||||||
else:
|
|
||||||
raise ValueError("无法获取 ComfyUI output 目录,请检查 folder_paths 是否可用")
|
|
||||||
|
|
||||||
import os
|
|
||||||
os.makedirs(output_folder, exist_ok=True)
|
|
||||||
|
|
||||||
def run_async_in_thread():
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
asyncio.set_event_loop(loop)
|
|
||||||
try:
|
|
||||||
return loop.run_until_complete(
|
|
||||||
self._process_batch_async(
|
|
||||||
prompts=[提示词],
|
|
||||||
model=模型,
|
|
||||||
resolution=分辨率,
|
|
||||||
aspect_ratio=宽高比,
|
|
||||||
images_per_prompt=生图数量,
|
|
||||||
input_images=input_images,
|
|
||||||
output_folder=output_folder,
|
|
||||||
pbar=pbar,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
loop.close()
|
|
||||||
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
||||||
future = executor.submit(run_async_in_thread)
|
|
||||||
try:
|
|
||||||
results = future.result(timeout=3600)
|
|
||||||
except TimeoutError:
|
|
||||||
raise RuntimeError("任务执行超时(1小时),请减少生图数量或检查网络连接")
|
|
||||||
|
|
||||||
success_count = sum(1 for r in results if r.get("success", False))
|
|
||||||
fail_count = len(results) - success_count
|
|
||||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
|
||||||
all_saved_files = []
|
|
||||||
for r in results:
|
|
||||||
all_saved_files.extend(r.get("saved_files", []))
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
|
||||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}")
|
|
||||||
|
|
||||||
# 失败详情
|
|
||||||
failed_results = [r for r in results if not r.get("success", False)]
|
|
||||||
if failed_results:
|
|
||||||
for fr in failed_results:
|
|
||||||
idx = fr.get("global_task_index", -1) + 1
|
|
||||||
error_msg = fr.get("error", "未知错误")
|
|
||||||
print(f" 失败 #{idx}: {提示词[:30]}{'...' if len(提示词) >= 30 else ''} → {error_msg}")
|
|
||||||
|
|
||||||
# 从磁盘加载最后 10 张图片
|
|
||||||
output_images = []
|
|
||||||
max_output_images = 10
|
|
||||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
|
||||||
for file_path in recent_files:
|
|
||||||
try:
|
|
||||||
img = Image.open(file_path)
|
|
||||||
output_images.append(img)
|
|
||||||
except Exception as e:
|
|
||||||
print(f"全能生图: 无法加载 {file_path} - {e}")
|
|
||||||
|
|
||||||
if not output_images:
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
output_images = [placeholder]
|
|
||||||
|
|
||||||
output_tensor = pil_to_tensor(output_images)
|
|
||||||
print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘,节点输出最后 {len(output_images)} 张")
|
|
||||||
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
|
|
||||||
# 优化:限制输出图片数量,避免内存爆炸
|
|
||||||
max_output_images = 20 # 最多输出20张图片到ComfyUI
|
|
||||||
|
|
||||||
if len(generated_images) > max_output_images:
|
|
||||||
print(f"全能生图: 生成 {len(generated_images)} 张图片,限制输出前 {max_output_images} 张到ComfyUI")
|
|
||||||
output_images = generated_images[:max_output_images]
|
|
||||||
else:
|
|
||||||
output_images = generated_images
|
|
||||||
|
|
||||||
# 转换输出图像
|
|
||||||
output_tensor = pil_to_tensor(output_images)
|
|
||||||
|
|
||||||
# 计算耗时并打印最终统计
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
if elapsed < 1:
|
|
||||||
time_str = f"{elapsed:.3f}s"
|
|
||||||
else:
|
|
||||||
time_str = f"{elapsed:.2f}s"
|
|
||||||
|
|
||||||
# 打印最终汇总
|
|
||||||
if fail_count > 0:
|
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
|
||||||
else:
|
|
||||||
print(f"[4/4] 完成!总耗时 {time_str} | 成功 {len(generated_images)}张")
|
|
||||||
|
|
||||||
# 最终内存清理
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
|
||||||
final_memory = process.memory_info().rss / 1024 / 1024
|
|
||||||
print(f"全能生图: 最终内存使用: {final_memory:.1f} MB")
|
|
||||||
|
|
||||||
return (output_tensor,)
|
|
||||||
|
|
||||||
except ValueError as e:
|
|
||||||
# 检测是否为授权错误
|
|
||||||
if str(e) == "未授权!":
|
|
||||||
print("请联系作者授权后方可使用!")
|
|
||||||
raise ValueError("未授权!") from None
|
|
||||||
else:
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"全能生图: ❌ {error_msg}")
|
|
||||||
if 跳过错误:
|
|
||||||
print("全能生图: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise ValueError(error_msg) from None
|
|
||||||
|
|
||||||
except RuntimeError as e:
|
|
||||||
error_full = str(e)
|
|
||||||
print(f"全能生图: ❌ {error_full}")
|
|
||||||
if 跳过错误:
|
|
||||||
print("全能生图: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise RuntimeError(error_full) from None
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
error_msg = str(e)
|
|
||||||
print(f"全能生图: ❌ {error_msg}")
|
|
||||||
if 跳过错误:
|
|
||||||
print("全能生图: ⚠️ 跳过错误已开启,返回占位图继续执行队列")
|
|
||||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
||||||
return (pil_to_tensor([placeholder]),)
|
|
||||||
raise type(e)(error_msg) from None
|
|
||||||
|
|
||||||
finally:
|
|
||||||
# 查询余额
|
|
||||||
if self.client is not None:
|
|
||||||
try:
|
|
||||||
balance_data = self.client.query_balance_sync()
|
|
||||||
balance_info = self.client.format_balance_info(balance_data)
|
|
||||||
print(f"全能生图: {balance_info}")
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
# 最终内存清理
|
|
||||||
import gc
|
|
||||||
gc.collect()
|
|
||||||
print(f"全能生图: 最终内存清理完成")
|
|
||||||
+264
-2
@@ -73,6 +73,79 @@ def _tensor_to_base64_url(tensor) -> str:
|
|||||||
return f"data:image/png;base64,{b64}"
|
return f"data:image/png;base64,{b64}"
|
||||||
|
|
||||||
|
|
||||||
|
def _video_to_base64_url(video) -> str:
|
||||||
|
"""ComfyUI VIDEO 对象 → data:video/<ext>;base64,xxx"""
|
||||||
|
import base64
|
||||||
|
import io as _io
|
||||||
|
|
||||||
|
source = video.get_stream_source()
|
||||||
|
|
||||||
|
if isinstance(source, _io.BytesIO):
|
||||||
|
source.seek(0)
|
||||||
|
data = source.read()
|
||||||
|
ext = "mp4"
|
||||||
|
else:
|
||||||
|
video_path = source
|
||||||
|
if not video_path or not os.path.isfile(video_path):
|
||||||
|
raise ValueError(f"无法获取参考视频文件路径(当前路径:{video_path})")
|
||||||
|
ext = os.path.splitext(video_path)[1].lower().lstrip(".")
|
||||||
|
if ext not in ("mp4", "mov"):
|
||||||
|
raise ValueError(f"参考视频格式须为 mp4 或 mov,当前为 .{ext}")
|
||||||
|
with open(video_path, "rb") as f:
|
||||||
|
data = f.read()
|
||||||
|
|
||||||
|
b64 = base64.b64encode(data).decode("utf-8")
|
||||||
|
return f"data:video/{ext};base64,{b64}"
|
||||||
|
|
||||||
|
|
||||||
|
def _audio_to_base64_url(audio) -> str:
|
||||||
|
"""ComfyUI AUDIO dict(waveform tensor + sample_rate)→ data:audio/wav;base64,xxx"""
|
||||||
|
import base64
|
||||||
|
import io
|
||||||
|
import struct
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
waveform = audio["waveform"] # shape: [B, C, N] or [C, N]
|
||||||
|
sample_rate = int(audio["sample_rate"])
|
||||||
|
|
||||||
|
# 统一为 [C, N]
|
||||||
|
if waveform.dim() == 3:
|
||||||
|
waveform = waveform[0]
|
||||||
|
|
||||||
|
# 转为 numpy float32,然后转 int16 PCM
|
||||||
|
wav_np = waveform.cpu().numpy()
|
||||||
|
if wav_np.ndim == 2:
|
||||||
|
# 多声道 → 单声道(取均值)
|
||||||
|
wav_np = wav_np.mean(axis=0)
|
||||||
|
wav_np = np.clip(wav_np, -1.0, 1.0)
|
||||||
|
pcm = (wav_np * 32767).astype(np.int16)
|
||||||
|
|
||||||
|
# 写 WAV 文件到内存
|
||||||
|
buf = io.BytesIO()
|
||||||
|
num_samples = len(pcm)
|
||||||
|
num_channels = 1
|
||||||
|
bits_per_sample = 16
|
||||||
|
byte_rate = sample_rate * num_channels * bits_per_sample // 8
|
||||||
|
block_align = num_channels * bits_per_sample // 8
|
||||||
|
data_size = num_samples * block_align
|
||||||
|
|
||||||
|
# RIFF header
|
||||||
|
buf.write(b"RIFF")
|
||||||
|
buf.write(struct.pack("<I", 36 + data_size))
|
||||||
|
buf.write(b"WAVE")
|
||||||
|
# fmt chunk
|
||||||
|
buf.write(b"fmt ")
|
||||||
|
buf.write(struct.pack("<IHHIIHH", 16, 1, num_channels, sample_rate,
|
||||||
|
byte_rate, block_align, bits_per_sample))
|
||||||
|
# data chunk
|
||||||
|
buf.write(b"data")
|
||||||
|
buf.write(struct.pack("<I", data_size))
|
||||||
|
buf.write(pcm.tobytes())
|
||||||
|
|
||||||
|
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||||
|
return f"data:audio/wav;base64,{b64}"
|
||||||
|
|
||||||
|
|
||||||
async def _url_to_tensor(url: str) -> torch.Tensor:
|
async def _url_to_tensor(url: str) -> torch.Tensor:
|
||||||
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None"""
|
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None"""
|
||||||
try:
|
try:
|
||||||
@@ -286,12 +359,201 @@ class Seedance:
|
|||||||
_show_balance()
|
_show_balance()
|
||||||
|
|
||||||
|
|
||||||
|
# ── 多模态参考生视频节点 ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
class SeedanceMultiModal:
|
||||||
|
"""Seedance 2.0 多模态参考生视频(参考图片 + 参考视频 + 参考音频 + 文本)"""
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||||
|
"模型": (_MODELS, {"default": "doubao-seedance-2-0-260128"}),
|
||||||
|
"分辨率": (_RESOLUTIONS, {"default": "720p"}),
|
||||||
|
"宽高比": (["adaptive", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"],
|
||||||
|
{"default": "adaptive"}),
|
||||||
|
"时长秒(-1=自动)": ("INT", {"default": 5, "min": -1, "max": 15, "step": 1}),
|
||||||
|
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||||||
|
"联网搜索": (["关闭", "打开"], {"default": "关闭"}),
|
||||||
|
"返回末帧图片": (["关闭", "打开"], {"default": "关闭"}),
|
||||||
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
|
||||||
|
},
|
||||||
|
"optional": {
|
||||||
|
"参考图片": ("IMAGE",),
|
||||||
|
"参考视频1": ("VIDEO",),
|
||||||
|
"参考视频2": ("VIDEO",),
|
||||||
|
"参考视频3": ("VIDEO",),
|
||||||
|
"参考音频1": ("AUDIO",),
|
||||||
|
"参考音频2": ("AUDIO",),
|
||||||
|
"参考音频3": ("AUDIO",),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("VIDEO", "IMAGE")
|
||||||
|
RETURN_NAMES = ("视频", "末帧图片")
|
||||||
|
FUNCTION = "generate"
|
||||||
|
CATEGORY = "comfyui_o1key/Seedance"
|
||||||
|
INPUT_IS_LIST = True
|
||||||
|
|
||||||
|
async def generate(self, **kwargs):
|
||||||
|
# INPUT_IS_LIST=True 时所有参数都是列表,取第一个元素
|
||||||
|
def _first(v, default=None):
|
||||||
|
if isinstance(v, list):
|
||||||
|
return v[0] if v else default
|
||||||
|
return v if v is not None else default
|
||||||
|
|
||||||
|
prompt = _first(kwargs.get("提示词"), "").strip()
|
||||||
|
model = _first(kwargs.get("模型"))
|
||||||
|
resolution = _first(kwargs.get("分辨率"))
|
||||||
|
ratio = _first(kwargs.get("宽高比"))
|
||||||
|
duration = _first(kwargs.get("时长秒(-1=自动)"), 5)
|
||||||
|
gen_audio = _first(kwargs.get("生成音频"), "关闭") == "打开"
|
||||||
|
web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开"
|
||||||
|
return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开"
|
||||||
|
seed = _first(kwargs.get("seed"), 0)
|
||||||
|
|
||||||
|
# 参考图片:INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表,直接保留
|
||||||
|
raw_images = kwargs.get("参考图片", None)
|
||||||
|
ref_images = [img for img in raw_images if img is not None] if raw_images else None
|
||||||
|
|
||||||
|
ref_videos = [_first(kwargs.get(f"参考视频{i}")) for i in range(1, 4)]
|
||||||
|
ref_audios = [_first(kwargs.get(f"参考音频{i}")) for i in range(1, 4)]
|
||||||
|
|
||||||
|
ref_videos = [v for v in ref_videos if v is not None]
|
||||||
|
ref_audios = [a for a in ref_audios if a is not None]
|
||||||
|
|
||||||
|
# ── 校验 ──────────────────────────────────────────────────────────
|
||||||
|
has_image = bool(ref_images)
|
||||||
|
has_video = len(ref_videos) > 0
|
||||||
|
has_audio = len(ref_audios) > 0
|
||||||
|
|
||||||
|
if not has_image and not has_video and not has_audio and not prompt:
|
||||||
|
raise ValueError("至少需要提供参考图片、参考视频或提示词之一。")
|
||||||
|
if has_audio and not has_image and not has_video:
|
||||||
|
raise ValueError("不可单独输入音频,请至少连接一张参考图片或一个参考视频。")
|
||||||
|
|
||||||
|
# ── 构建 content 列表 ─────────────────────────────────────────────
|
||||||
|
content = []
|
||||||
|
|
||||||
|
# 参考图片(批次,最多9张)
|
||||||
|
if has_image:
|
||||||
|
imgs = ref_images[:9]
|
||||||
|
if len(ref_images) > 9:
|
||||||
|
print(f"[SeedanceMultiModal] 参考图片超过9张,仅取前9张(共{len(ref_images)}张)")
|
||||||
|
for img_tensor in imgs:
|
||||||
|
# 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C],统一确保有 batch 维
|
||||||
|
if img_tensor.dim() == 3:
|
||||||
|
img_tensor = img_tensor.unsqueeze(0)
|
||||||
|
url = _tensor_to_base64_url(img_tensor)
|
||||||
|
content.append({
|
||||||
|
"type": "image_url",
|
||||||
|
"image_url": {"url": url},
|
||||||
|
"role": "reference_image",
|
||||||
|
})
|
||||||
|
|
||||||
|
# 参考视频(最多3个)
|
||||||
|
for v in ref_videos:
|
||||||
|
url = _video_to_base64_url(v)
|
||||||
|
content.append({
|
||||||
|
"type": "video_url",
|
||||||
|
"video_url": {"url": url},
|
||||||
|
"role": "reference_video",
|
||||||
|
})
|
||||||
|
|
||||||
|
# 参考音频(最多3段)
|
||||||
|
for a in ref_audios:
|
||||||
|
url = _audio_to_base64_url(a)
|
||||||
|
content.append({
|
||||||
|
"type": "audio_url",
|
||||||
|
"audio_url": {"url": url},
|
||||||
|
"role": "reference_audio",
|
||||||
|
})
|
||||||
|
|
||||||
|
# 文本提示词(放最后)
|
||||||
|
if prompt:
|
||||||
|
content.append({"type": "text", "text": prompt})
|
||||||
|
|
||||||
|
if not content:
|
||||||
|
raise ValueError("content 为空,请至少提供参考图片、参考视频或提示词。")
|
||||||
|
|
||||||
|
# ── 构建请求体(new-api 兼容格式)──────────────────────────────────
|
||||||
|
metadata: dict = {
|
||||||
|
"resolution": resolution,
|
||||||
|
"watermark": False,
|
||||||
|
"content": content,
|
||||||
|
}
|
||||||
|
|
||||||
|
if ratio != "adaptive":
|
||||||
|
metadata["ratio"] = ratio
|
||||||
|
if duration != -1:
|
||||||
|
metadata["duration"] = duration
|
||||||
|
if gen_audio:
|
||||||
|
metadata["generate_audio"] = True
|
||||||
|
if return_last:
|
||||||
|
metadata["return_last_frame"] = True
|
||||||
|
if seed != 0:
|
||||||
|
metadata["seed"] = seed
|
||||||
|
if web_search:
|
||||||
|
metadata["tools"] = [{"type": "web_search"}]
|
||||||
|
|
||||||
|
# 顶层 image:取第一张参考图的 base64(new-api 单图字段)
|
||||||
|
first_image_url = next(
|
||||||
|
(item["image_url"]["url"] for item in content if item["type"] == "image_url"),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
|
body = {
|
||||||
|
"model": model,
|
||||||
|
"prompt": prompt if prompt else " ",
|
||||||
|
"metadata": metadata,
|
||||||
|
}
|
||||||
|
if first_image_url:
|
||||||
|
body["image"] = first_image_url
|
||||||
|
|
||||||
|
# ── 打印请求体结构(base64 截断显示)────────────────────────────────
|
||||||
|
import json as _json, copy as _copy
|
||||||
|
def _truncate_body(obj):
|
||||||
|
if isinstance(obj, dict):
|
||||||
|
return {k: _truncate_body(v) for k, v in obj.items()}
|
||||||
|
if isinstance(obj, list):
|
||||||
|
return [_truncate_body(i) for i in obj]
|
||||||
|
if isinstance(obj, str) and obj.startswith("data:") and len(obj) > 80:
|
||||||
|
return obj[:60] + f"...[{len(obj)}chars]"
|
||||||
|
return obj
|
||||||
|
print("[SeedanceMultiModal] 请求体预览:")
|
||||||
|
print(_json.dumps(_truncate_body(_copy.deepcopy(body)), ensure_ascii=False, indent=2))
|
||||||
|
|
||||||
|
# ── 保存路径 ──────────────────────────────────────────────────────
|
||||||
|
video_dir = _get_video_output_dir()
|
||||||
|
counter = _get_next_counter(video_dir, "seedance_mm")
|
||||||
|
save_path = os.path.join(video_dir, f"seedance_mm_{counter:05d}.mp4")
|
||||||
|
|
||||||
|
client = SeedanceClient()
|
||||||
|
pbar = _make_pbar()
|
||||||
|
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)
|
||||||
|
|
||||||
|
try:
|
||||||
|
result_path, last_frame_url = await client.generate_async(
|
||||||
|
body=body, save_path=save_path,
|
||||||
|
on_stage=on_stage, on_progress=on_prog,
|
||||||
|
)
|
||||||
|
last_frame_tensor = None
|
||||||
|
if return_last and last_frame_url:
|
||||||
|
last_frame_tensor = await _url_to_tensor(last_frame_url)
|
||||||
|
return (InputImpl.VideoFromFile(result_path), last_frame_tensor)
|
||||||
|
finally:
|
||||||
|
_show_balance()
|
||||||
|
|
||||||
|
|
||||||
# ── 节点注册 ──────────────────────────────────────────────────────────────────
|
# ── 节点注册 ──────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
"Seedance": Seedance,
|
"Seedance": Seedance,
|
||||||
|
"SeedanceMultiModal": SeedanceMultiModal,
|
||||||
}
|
}
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
"Seedance": "Seedance 视频生成",
|
"Seedance": "Seedance 视频生成",
|
||||||
|
"SeedanceMultiModal": "Seedance 多模态参考生视频",
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,28 @@
|
|||||||
|
"""
|
||||||
|
流式文本预览节点
|
||||||
|
接收文本输入,支持 markdown 渲染,通过 ComfyUI 事件系统实时推送内容
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class StreamPreview:
|
||||||
|
"""
|
||||||
|
流式 Markdown 预览节点
|
||||||
|
接收任意文本,在节点面板中实时渲染为 Markdown 格式
|
||||||
|
"""
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"文本": ("STRING", {"forceInput": True}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("STRING",)
|
||||||
|
RETURN_NAMES = ("文本",)
|
||||||
|
FUNCTION = "preview"
|
||||||
|
CATEGORY = "text/preview"
|
||||||
|
OUTPUT_NODE = True
|
||||||
|
|
||||||
|
def preview(self, 文本: str):
|
||||||
|
return {"ui": {"text": [文本]}, "result": (文本,)}
|
||||||
+227
-30
@@ -6,17 +6,19 @@ ComfyUI 自定义节点,通过 OpenAI 兼容协议调用市面上主流的 AI
|
|||||||
API 密钥和地址通过插件统一配置(环境变量或 .config 文件),与 Google Gemini 节点一致
|
API 密钥和地址通过插件统一配置(环境变量或 .config 文件),与 Google Gemini 节点一致
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
import time
|
import time
|
||||||
import base64
|
import base64
|
||||||
import json
|
import json
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from typing import Optional, Tuple
|
from typing import Optional, Tuple, List
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
from ..utils.image_utils import tensor_to_pil
|
from ..utils.image_utils import tensor_to_pil
|
||||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||||
|
from ..utils.file_types import FileList
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
# 模型配置
|
# 模型配置
|
||||||
@@ -76,7 +78,12 @@ class UniversalLLMChat:
|
|||||||
},
|
},
|
||||||
"optional": {
|
"optional": {
|
||||||
"图片": ("IMAGE",),
|
"图片": ("IMAGE",),
|
||||||
}
|
"视频": ("VIDEO",),
|
||||||
|
"文件": ("FILE_LIST",),
|
||||||
|
},
|
||||||
|
"hidden": {
|
||||||
|
"node_id": "UNIQUE_ID",
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
RETURN_TYPES = ("STRING",)
|
RETURN_TYPES = ("STRING",)
|
||||||
@@ -113,12 +120,90 @@ class UniversalLLMChat:
|
|||||||
b64 = base64.b64encode(data).decode('utf-8')
|
b64 = base64.b64encode(data).decode('utf-8')
|
||||||
return f"data:image/jpeg;base64,{b64}"
|
return f"data:image/jpeg;base64,{b64}"
|
||||||
|
|
||||||
def _build_messages(
|
# 文件大小限制
|
||||||
|
MAX_FILE_SIZE = 50 * 1024 * 1024 # 单文件 50MB
|
||||||
|
MAX_TOTAL_FILE_SIZE = 50 * 1024 * 1024 # 所有文件总计 50MB
|
||||||
|
|
||||||
|
# 常见 MIME 类型映射
|
||||||
|
MIME_MAP = {
|
||||||
|
".pdf": "application/pdf",
|
||||||
|
".txt": "text/plain",
|
||||||
|
".md": "text/markdown",
|
||||||
|
".csv": "text/csv",
|
||||||
|
".json": "application/json",
|
||||||
|
".py": "text/x-python",
|
||||||
|
".js": "text/javascript",
|
||||||
|
".html": "text/html",
|
||||||
|
".xml": "application/xml",
|
||||||
|
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||||
|
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||||
|
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||||
|
".zip": "application/zip",
|
||||||
|
}
|
||||||
|
|
||||||
|
# 纯文本类型,直接读取内容
|
||||||
|
TEXT_EXTS = {".txt", ".md", ".csv", ".json", ".py", ".js", ".ts", ".html",
|
||||||
|
".xml", ".yaml", ".yml", ".toml", ".ini", ".cfg", ".log",
|
||||||
|
".sh", ".bat", ".sql", ".css", ".scss", ".jsx", ".tsx"}
|
||||||
|
|
||||||
|
def _load_files(self, file_paths_str: str) -> List[dict]:
|
||||||
|
"""读取文件列表,返回 content part 数组"""
|
||||||
|
if not file_paths_str or not file_paths_str.strip():
|
||||||
|
return []
|
||||||
|
|
||||||
|
paths = [p.strip() for p in file_paths_str.split(",") if p.strip()]
|
||||||
|
parts = []
|
||||||
|
total_size = 0
|
||||||
|
|
||||||
|
for path in paths:
|
||||||
|
if not os.path.isfile(path):
|
||||||
|
raise ValueError(f"文件不存在: {path}")
|
||||||
|
|
||||||
|
file_size = os.path.getsize(path)
|
||||||
|
if file_size > self.MAX_FILE_SIZE:
|
||||||
|
raise ValueError(f"文件 {os.path.basename(path)} 大小 {file_size / 1024 / 1024:.1f}MB 超过单文件 50MB 限制")
|
||||||
|
|
||||||
|
total_size += file_size
|
||||||
|
if total_size > self.MAX_TOTAL_FILE_SIZE:
|
||||||
|
raise ValueError(f"所有文件总大小超过 50MB 限制")
|
||||||
|
|
||||||
|
ext = os.path.splitext(path)[1].lower()
|
||||||
|
mime = self.MIME_MAP.get(ext, "application/octet-stream")
|
||||||
|
filename = os.path.basename(path)
|
||||||
|
|
||||||
|
if ext in self.TEXT_EXTS:
|
||||||
|
# 文本文件直接读取内容
|
||||||
|
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
||||||
|
text_content = f.read()
|
||||||
|
parts.append({
|
||||||
|
"type": "text",
|
||||||
|
"text": f"[文件: {filename}]\n```\n{text_content}\n```",
|
||||||
|
})
|
||||||
|
else:
|
||||||
|
# 二进制文件转 base64,使用 file 格式(OpenAI 兼容协议)
|
||||||
|
with open(path, "rb") as f:
|
||||||
|
file_data = base64.b64encode(f.read()).decode("utf-8")
|
||||||
|
parts.append({
|
||||||
|
"type": "file",
|
||||||
|
"file": {
|
||||||
|
"filename": filename,
|
||||||
|
"file_data": f"data:{mime};base64,{file_data}",
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
print(f"全能LLM: 加载文件 {filename} ({file_size / 1024:.1f}KB, {mime})")
|
||||||
|
|
||||||
|
return parts
|
||||||
|
|
||||||
|
def _build_input(
|
||||||
self,
|
self,
|
||||||
prompt: str,
|
prompt: str,
|
||||||
images: Optional[torch.Tensor] = None,
|
images: Optional[torch.Tensor] = None,
|
||||||
|
file_paths: str = "",
|
||||||
|
file_list: Optional[FileList] = None,
|
||||||
|
video=None,
|
||||||
) -> list:
|
) -> list:
|
||||||
"""构建 OpenAI 格式的 messages 数组"""
|
"""构建 chat/completions 格式的 messages 数组"""
|
||||||
image_data_urls = []
|
image_data_urls = []
|
||||||
pil_images_cache = [] # 保留 PIL Image 用于总体积重新编码
|
pil_images_cache = [] # 保留 PIL Image 用于总体积重新编码
|
||||||
|
|
||||||
@@ -177,49 +262,152 @@ class UniversalLLMChat:
|
|||||||
print(f"全能LLM: 无法将 {len(pil_images_cache)} 张图片压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内,请减少图片数量或降低分辨率")
|
print(f"全能LLM: 无法将 {len(pil_images_cache)} 张图片压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内,请减少图片数量或降低分辨率")
|
||||||
raise ValueError(f"图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过限制,无法压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内")
|
raise ValueError(f"图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过限制,无法压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内")
|
||||||
|
|
||||||
if not image_data_urls:
|
# 处理视频输入(ComfyUI VIDEO 类型)
|
||||||
|
video_url_str = ""
|
||||||
|
if video is not None:
|
||||||
|
# 从 VIDEO 对象中提取文件路径
|
||||||
|
vp = None
|
||||||
|
if isinstance(video, dict):
|
||||||
|
vp = video.get("video") or video.get("path") or video.get("file") or video.get("filename")
|
||||||
|
if not vp:
|
||||||
|
for val in video.values():
|
||||||
|
if isinstance(val, str) and os.path.exists(val):
|
||||||
|
vp = val
|
||||||
|
break
|
||||||
|
elif isinstance(video, str):
|
||||||
|
vp = video
|
||||||
|
else:
|
||||||
|
for attr in ("video", "path", "filename"):
|
||||||
|
if hasattr(video, attr):
|
||||||
|
vp = getattr(video, attr)
|
||||||
|
break
|
||||||
|
if not vp and hasattr(video, "__dict__"):
|
||||||
|
for attr_val in video.__dict__.values():
|
||||||
|
if isinstance(attr_val, str) and os.path.isfile(attr_val):
|
||||||
|
vp = attr_val
|
||||||
|
break
|
||||||
|
|
||||||
|
if not vp or not os.path.isfile(vp):
|
||||||
|
raise ValueError(f"视频文件不存在或路径无效: {vp}")
|
||||||
|
|
||||||
|
mime_map = {
|
||||||
|
".mp4": "video/mp4", ".mpeg": "video/mpeg", ".mpg": "video/mpg",
|
||||||
|
".mov": "video/quicktime", ".avi": "video/x-msvideo",
|
||||||
|
".flv": "video/x-flv", ".webm": "video/webm",
|
||||||
|
".wmv": "video/x-ms-wmv", ".mkv": "video/x-matroska",
|
||||||
|
}
|
||||||
|
ext = os.path.splitext(vp)[1].lower()
|
||||||
|
mime = mime_map.get(ext, "video/mp4")
|
||||||
|
file_size = os.path.getsize(vp)
|
||||||
|
print(f"全能LLM: 加载视频 {os.path.basename(vp)} ({file_size / 1024 / 1024:.1f}MB, {mime})")
|
||||||
|
with open(vp, "rb") as f:
|
||||||
|
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||||
|
video_url_str = f"data:{mime};base64,{b64}"
|
||||||
|
|
||||||
|
# 加载文件:优先使用 FILE_LIST,其次使用字符串路径
|
||||||
|
file_parts = []
|
||||||
|
if file_list:
|
||||||
|
for fd in file_list:
|
||||||
|
print(f"全能LLM: 使用文件 {fd.filename}{fd.extension} ({fd.size / 1024:.1f}KB)")
|
||||||
|
file_parts.append({
|
||||||
|
"type": "file",
|
||||||
|
"file": {
|
||||||
|
"filename": fd.filename + fd.extension,
|
||||||
|
"file_data": f"data:{fd.mime_type};base64,{fd.data}",
|
||||||
|
},
|
||||||
|
})
|
||||||
|
elif file_paths:
|
||||||
|
file_parts = self._load_files(file_paths)
|
||||||
|
|
||||||
|
# 纯文本,无图片无文件无视频
|
||||||
|
if not image_data_urls and not file_parts and not video_url_str:
|
||||||
return [{"role": "user", "content": prompt}]
|
return [{"role": "user", "content": prompt}]
|
||||||
|
|
||||||
content_parts = []
|
content_parts = []
|
||||||
|
|
||||||
|
# 图片
|
||||||
for url in image_data_urls:
|
for url in image_data_urls:
|
||||||
content_parts.append({
|
content_parts.append({
|
||||||
"type": "image_url",
|
"type": "image_url",
|
||||||
"image_url": {"url": url}
|
"image_url": {"url": url},
|
||||||
})
|
})
|
||||||
|
|
||||||
|
# 视频:用 image_url 类型传 data URL(Gemini OpenAI 兼容层支持此格式)
|
||||||
|
# 同时保留 video_url 类型作为备用(其他支持 video_url 的模型)
|
||||||
|
if video_url_str:
|
||||||
|
content_parts.append({
|
||||||
|
"type": "image_url",
|
||||||
|
"image_url": {"url": video_url_str},
|
||||||
|
})
|
||||||
|
|
||||||
|
# 文件
|
||||||
|
for fp in file_parts:
|
||||||
|
content_parts.append(fp)
|
||||||
|
|
||||||
content_parts.append({
|
content_parts.append({
|
||||||
"type": "text",
|
"type": "text",
|
||||||
"text": prompt
|
"text": prompt,
|
||||||
})
|
})
|
||||||
|
|
||||||
return [{"role": "user", "content": content_parts}]
|
return [{"role": "user", "content": content_parts}]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _send_stream_token(node_id, token, done=False):
|
||||||
|
"""通过 PromptServer 向前端推送流式 token"""
|
||||||
|
try:
|
||||||
|
from server import PromptServer
|
||||||
|
PromptServer.instance.send_sync(
|
||||||
|
"o1key.stream_token",
|
||||||
|
{"node_id": str(node_id), "token": token, "done": done},
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
def generate(
|
def generate(
|
||||||
self,
|
self,
|
||||||
模型: str,
|
模型: str,
|
||||||
提示词: str,
|
提示词: str,
|
||||||
图片: Optional[torch.Tensor] = None,
|
图片: Optional[torch.Tensor] = None,
|
||||||
|
视频=None,
|
||||||
|
文件: Optional[FileList] = None,
|
||||||
|
node_id: str = "",
|
||||||
) -> Tuple[str]:
|
) -> Tuple[str]:
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
self._ensure_config()
|
self._ensure_config()
|
||||||
|
|
||||||
# 构建 messages
|
# 构建 input
|
||||||
messages = self._build_messages(提示词, 图片)
|
input_data = self._build_input(提示词, 图片, "", 文件, 视频)
|
||||||
|
|
||||||
img_count = len(tensor_to_pil(图片)) if 图片 is not None else 0
|
img_count = len(tensor_to_pil(图片)) if 图片 is not None else 0
|
||||||
input_desc = "文本" + (f" + {img_count}张图片" if img_count > 0 else "")
|
file_count = len(文件) if 文件 else 0
|
||||||
|
input_desc = "文本"
|
||||||
|
if img_count: input_desc += f" + {img_count}张图片"
|
||||||
|
if 视频 is not None: input_desc += " + 视频"
|
||||||
|
if file_count: input_desc += f" + {file_count}个文件"
|
||||||
|
|
||||||
print(f"全能LLM: 模型 = {模型}")
|
print(f"全能LLM: 模型 = {模型}")
|
||||||
print(f"全能LLM: 输入 = {input_desc}")
|
print(f"全能LLM: 输入 = {input_desc}")
|
||||||
|
|
||||||
# 构建请求体
|
# 构建请求体(chat/completions 格式)
|
||||||
request_body = {
|
request_body = {
|
||||||
"model": 模型,
|
"model": 模型,
|
||||||
"messages": messages,
|
"messages": input_data,
|
||||||
"stream": False,
|
"stream": True,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# 打印请求体,base64 截断显示
|
||||||
|
def _truncate_for_log(obj):
|
||||||
|
if isinstance(obj, dict):
|
||||||
|
return {k: _truncate_for_log(v) for k, v in obj.items()}
|
||||||
|
if isinstance(obj, list):
|
||||||
|
return [_truncate_for_log(i) for i in obj]
|
||||||
|
if isinstance(obj, str) and (obj.startswith("data:image") or obj.startswith("data:application") or obj.startswith("data:text")):
|
||||||
|
return obj[:60] + f"...[{len(obj)}chars]"
|
||||||
|
return obj
|
||||||
|
print(f"全能LLM: 请求原始内容 = {json.dumps(_truncate_for_log(request_body), ensure_ascii=False)}")
|
||||||
|
|
||||||
# 发送请求(在独立线程中运行异步请求,避免与 ComfyUI 事件循环冲突)
|
# 发送请求(在独立线程中运行异步请求,避免与 ComfyUI 事件循环冲突)
|
||||||
import aiohttp
|
import aiohttp
|
||||||
import asyncio
|
import asyncio
|
||||||
@@ -236,9 +424,9 @@ class UniversalLLMChat:
|
|||||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||||
async with session.post(url, headers=headers, json=request_body) as resp:
|
async with session.post(url, headers=headers, json=request_body) as resp:
|
||||||
status = resp.status
|
status = resp.status
|
||||||
body = await resp.text()
|
|
||||||
|
|
||||||
if status != 200:
|
if status != 200:
|
||||||
|
body = await resp.text()
|
||||||
try:
|
try:
|
||||||
err_data = json.loads(body)
|
err_data = json.loads(body)
|
||||||
err_msg = err_data.get("error", {}).get("message", body[:200])
|
err_msg = err_data.get("error", {}).get("message", body[:200])
|
||||||
@@ -256,7 +444,30 @@ class UniversalLLMChat:
|
|||||||
else:
|
else:
|
||||||
raise RuntimeError(f"API 错误 ({status}): {err_msg}")
|
raise RuntimeError(f"API 错误 ({status}): {err_msg}")
|
||||||
|
|
||||||
return json.loads(body)
|
# 流式读取,拼接 delta content
|
||||||
|
reply_parts = []
|
||||||
|
async for raw_line in resp.content:
|
||||||
|
line = raw_line.decode("utf-8").strip()
|
||||||
|
if not line or not line.startswith("data:"):
|
||||||
|
continue
|
||||||
|
data_str = line[len("data:"):].strip()
|
||||||
|
if data_str == "[DONE]":
|
||||||
|
break
|
||||||
|
try:
|
||||||
|
chunk = json.loads(data_str)
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
choices = chunk.get("choices")
|
||||||
|
if not choices:
|
||||||
|
continue
|
||||||
|
delta = choices[0].get("delta", {})
|
||||||
|
content = delta.get("content")
|
||||||
|
if content:
|
||||||
|
reply_parts.append(content)
|
||||||
|
UniversalLLMChat._send_stream_token(node_id, content)
|
||||||
|
|
||||||
|
UniversalLLMChat._send_stream_token(node_id, "", done=True)
|
||||||
|
return "".join(reply_parts)
|
||||||
|
|
||||||
def _run_in_thread():
|
def _run_in_thread():
|
||||||
loop = asyncio.new_event_loop()
|
loop = asyncio.new_event_loop()
|
||||||
@@ -266,24 +477,10 @@ class UniversalLLMChat:
|
|||||||
loop.close()
|
loop.close()
|
||||||
|
|
||||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||||
response_data = pool.submit(_run_in_thread).result()
|
reply = pool.submit(_run_in_thread).result()
|
||||||
|
|
||||||
# 解析响应
|
|
||||||
choices = response_data.get("choices", [])
|
|
||||||
if not choices:
|
|
||||||
raise RuntimeError("API 返回了空响应(无 choices)")
|
|
||||||
|
|
||||||
reply = choices[0].get("message", {}).get("content", "")
|
|
||||||
|
|
||||||
# Token 用量
|
|
||||||
usage = response_data.get("usage", {})
|
|
||||||
prompt_tokens = usage.get("prompt_tokens", 0)
|
|
||||||
completion_tokens = usage.get("completion_tokens", 0)
|
|
||||||
total_tokens = usage.get("total_tokens", 0)
|
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
elapsed = time.time() - start_time
|
||||||
print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)")
|
print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)")
|
||||||
print(f"全能LLM: Token 用量 — 输入: {prompt_tokens}, 输出: {completion_tokens}, 合计: {total_tokens}")
|
|
||||||
if reply:
|
if reply:
|
||||||
preview = reply[:100] + "..." if len(reply) > 100 else reply
|
preview = reply[:100] + "..." if len(reply) > 100 else reply
|
||||||
print(f"全能LLM: 回复预览: {preview}")
|
print(f"全能LLM: 回复预览: {preview}")
|
||||||
|
|||||||
+46
-63
@@ -1,9 +1,10 @@
|
|||||||
"""
|
"""
|
||||||
通用视频预览节点
|
视频预览节点
|
||||||
ComfyUI 自定义节点,接收视频文件路径并在前端展示预览
|
接收 VIDEO 类型,在前端内嵌播放器预览
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import io
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import folder_paths
|
import folder_paths
|
||||||
@@ -11,8 +12,6 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
FOLDER_PATHS_AVAILABLE = False
|
FOLDER_PATHS_AVAILABLE = False
|
||||||
|
|
||||||
SUPPORTED_VIDEO_EXTENSIONS = {".mp4", ".webm", ".mov", ".avi", ".mkv", ".flv", ".wmv", ".3gp"}
|
|
||||||
|
|
||||||
|
|
||||||
def _get_output_dir() -> str:
|
def _get_output_dir() -> str:
|
||||||
if FOLDER_PATHS_AVAILABLE:
|
if FOLDER_PATHS_AVAILABLE:
|
||||||
@@ -22,86 +21,70 @@ def _get_output_dir() -> str:
|
|||||||
|
|
||||||
|
|
||||||
class VideoPreview:
|
class VideoPreview:
|
||||||
"""
|
|
||||||
通用视频预览节点
|
|
||||||
|
|
||||||
功能:
|
|
||||||
- 接收视频文件路径(STRING)
|
|
||||||
- 在 ComfyUI 前端节点上内嵌 <video> 播放器进行预览
|
|
||||||
- 支持 mp4, webm, mov, avi, mkv 等主流格式
|
|
||||||
"""
|
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(cls):
|
def INPUT_TYPES(cls):
|
||||||
return {
|
return {
|
||||||
"required": {
|
"required": {
|
||||||
"预览视频": ("STRING", {"forceInput": True}),
|
"视频": ("VIDEO",),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
RETURN_TYPES = ()
|
RETURN_TYPES = ()
|
||||||
OUTPUT_NODE = True
|
OUTPUT_NODE = True
|
||||||
FUNCTION = "preview"
|
FUNCTION = "preview"
|
||||||
CATEGORY = "video"
|
CATEGORY = "comfyui_o1key/Utils"
|
||||||
|
|
||||||
DESCRIPTION = (
|
def preview(self, 视频) -> dict:
|
||||||
"通用视频预览节点。\n"
|
# 用官方接口取文件路径
|
||||||
"接收视频文件路径,在节点上显示视频播放器。\n"
|
source = 视频.get_stream_source()
|
||||||
"支持 mp4, webm, mov, avi, mkv 等主流视频格式。"
|
|
||||||
)
|
|
||||||
|
|
||||||
def preview(self, **kwargs) -> dict:
|
|
||||||
video_path = kwargs.get("预览视频", "")
|
|
||||||
if not video_path or not video_path.strip():
|
|
||||||
raise ValueError("视频路径为空")
|
|
||||||
|
|
||||||
video_path = video_path.strip()
|
|
||||||
|
|
||||||
if not os.path.isfile(video_path):
|
|
||||||
raise ValueError(f"视频文件不存在: {video_path}")
|
|
||||||
|
|
||||||
ext = os.path.splitext(video_path)[1].lower()
|
|
||||||
if ext not in SUPPORTED_VIDEO_EXTENSIONS:
|
|
||||||
raise ValueError(
|
|
||||||
f"不支持的视频格式 '{ext}',"
|
|
||||||
f"支持: {', '.join(sorted(SUPPORTED_VIDEO_EXTENSIONS))}"
|
|
||||||
)
|
|
||||||
|
|
||||||
output_dir = _get_output_dir()
|
|
||||||
abs_video = os.path.abspath(video_path)
|
|
||||||
abs_output = os.path.abspath(output_dir)
|
|
||||||
|
|
||||||
if abs_video.startswith(abs_output):
|
|
||||||
rel_path = os.path.relpath(abs_video, abs_output)
|
|
||||||
subfolder = os.path.dirname(rel_path).replace("\\", "/")
|
|
||||||
filename = os.path.basename(rel_path)
|
|
||||||
file_type = "output"
|
|
||||||
else:
|
|
||||||
filename = os.path.basename(abs_video)
|
|
||||||
subfolder = ""
|
|
||||||
file_type = "output"
|
|
||||||
|
|
||||||
# 如果文件不在 output 目录下,复制一份到 output/video/
|
|
||||||
target_dir = os.path.join(output_dir, "video")
|
|
||||||
os.makedirs(target_dir, exist_ok=True)
|
|
||||||
target_path = os.path.join(target_dir, filename)
|
|
||||||
|
|
||||||
if not os.path.exists(target_path) or abs_video != os.path.abspath(target_path):
|
|
||||||
import shutil
|
|
||||||
shutil.copy2(abs_video, target_path)
|
|
||||||
|
|
||||||
|
if isinstance(source, io.BytesIO):
|
||||||
|
# BytesIO 情况:写到 output/video/ 临时文件
|
||||||
|
output_dir = os.path.join(_get_output_dir(), "video")
|
||||||
|
os.makedirs(output_dir, exist_ok=True)
|
||||||
|
filename = "preview_tmp.mp4"
|
||||||
|
tmp_path = os.path.join(output_dir, filename)
|
||||||
|
source.seek(0)
|
||||||
|
with open(tmp_path, "wb") as f:
|
||||||
|
f.write(source.read())
|
||||||
subfolder = "video"
|
subfolder = "video"
|
||||||
|
else:
|
||||||
|
video_path = source
|
||||||
|
output_dir = _get_output_dir()
|
||||||
|
abs_video = os.path.abspath(video_path)
|
||||||
|
abs_output = os.path.abspath(output_dir)
|
||||||
|
|
||||||
file_size = os.path.getsize(abs_video)
|
if abs_video.startswith(abs_output):
|
||||||
size_mb = file_size / (1024 * 1024)
|
rel_path = os.path.relpath(abs_video, abs_output)
|
||||||
print(f"视频预览: {filename} ({size_mb:.1f}MB)")
|
subfolder = os.path.dirname(rel_path).replace("\\", "/")
|
||||||
|
filename = os.path.basename(rel_path)
|
||||||
|
else:
|
||||||
|
# 文件在 output 目录外,复制一份
|
||||||
|
target_dir = os.path.join(output_dir, "video")
|
||||||
|
os.makedirs(target_dir, exist_ok=True)
|
||||||
|
filename = os.path.basename(abs_video)
|
||||||
|
target_path = os.path.join(target_dir, filename)
|
||||||
|
if not os.path.exists(target_path):
|
||||||
|
import shutil
|
||||||
|
shutil.copy2(abs_video, target_path)
|
||||||
|
subfolder = "video"
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"ui": {
|
"ui": {
|
||||||
"videos": [{
|
"videos": [{
|
||||||
"filename": filename,
|
"filename": filename,
|
||||||
"subfolder": subfolder,
|
"subfolder": subfolder,
|
||||||
"type": file_type,
|
"type": "output",
|
||||||
}],
|
}],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"VideoPreview": VideoPreview,
|
||||||
|
}
|
||||||
|
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
|
"VideoPreview": "预览视频",
|
||||||
|
}
|
||||||
|
|||||||
+50
-8
@@ -3,13 +3,13 @@
|
|||||||
用于在 ComfyUI 节点间传递文件数据
|
用于在 ComfyUI 节点间传递文件数据
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import NamedTuple
|
from typing import NamedTuple, List
|
||||||
|
|
||||||
|
|
||||||
class FileData(NamedTuple):
|
class FileData(NamedTuple):
|
||||||
"""
|
"""
|
||||||
文件数据类型,用于在节点间传递
|
单个文件数据,用于节点间传递
|
||||||
|
|
||||||
Attributes:
|
Attributes:
|
||||||
path: 文件完整路径
|
path: 文件完整路径
|
||||||
filename: 文件名(不含扩展名)
|
filename: 文件名(不含扩展名)
|
||||||
@@ -26,15 +26,57 @@ class FileData(NamedTuple):
|
|||||||
size: int
|
size: int
|
||||||
|
|
||||||
|
|
||||||
# 支持的文档 MIME 类型映射
|
# FILE_LIST 类型:FileData 的列表,用于多文件传递
|
||||||
|
# ComfyUI 自定义类型名,节点 RETURN_TYPES / INPUT_TYPES 中使用 "FILE_LIST"
|
||||||
|
FileList = List[FileData]
|
||||||
|
|
||||||
|
|
||||||
|
# 支持的文件 MIME 类型映射(与 universal_llm.py 的 MIME_MAP 保持一致)
|
||||||
DOCUMENT_MIME_TYPES = {
|
DOCUMENT_MIME_TYPES = {
|
||||||
".pdf": "application/pdf",
|
".pdf": "application/pdf",
|
||||||
".txt": "text/plain"
|
".txt": "text/plain",
|
||||||
|
".md": "text/markdown",
|
||||||
|
".csv": "text/csv",
|
||||||
|
".json": "application/json",
|
||||||
|
".py": "text/x-python",
|
||||||
|
".js": "text/javascript",
|
||||||
|
".ts": "text/javascript",
|
||||||
|
".html": "text/html",
|
||||||
|
".xml": "application/xml",
|
||||||
|
".yaml": "text/plain",
|
||||||
|
".yml": "text/plain",
|
||||||
|
".toml": "text/plain",
|
||||||
|
".ini": "text/plain",
|
||||||
|
".cfg": "text/plain",
|
||||||
|
".log": "text/plain",
|
||||||
|
".sh": "text/plain",
|
||||||
|
".bat": "text/plain",
|
||||||
|
".sql": "text/plain",
|
||||||
|
".css": "text/plain",
|
||||||
|
".scss": "text/plain",
|
||||||
|
".jsx": "text/javascript",
|
||||||
|
".tsx": "text/javascript",
|
||||||
|
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||||
|
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||||
|
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||||
|
".zip": "application/zip",
|
||||||
|
".png": "image/png",
|
||||||
|
".jpg": "image/jpeg",
|
||||||
|
".jpeg": "image/jpeg",
|
||||||
|
".webp": "image/webp",
|
||||||
|
".wav": "audio/wav",
|
||||||
|
".mp3": "audio/mpeg",
|
||||||
|
".mp4": "video/mp4",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# 单文件大小上限:50MB
|
||||||
|
FILE_SIZE_LIMIT = 50 * 1024 * 1024
|
||||||
|
|
||||||
# 文件大小限制(字节)
|
# 所有文件总大小上限:50MB
|
||||||
|
TOTAL_FILE_SIZE_LIMIT = 50 * 1024 * 1024
|
||||||
|
|
||||||
|
# 兼容旧代码
|
||||||
FILE_SIZE_LIMITS = {
|
FILE_SIZE_LIMITS = {
|
||||||
".pdf": 50 * 1024 * 1024, # 50MB (Gemini API 官方限制)
|
".pdf": FILE_SIZE_LIMIT,
|
||||||
".txt": 20 * 1024 * 1024 # 20MB (保守限制)
|
".txt": FILE_SIZE_LIMIT,
|
||||||
}
|
}
|
||||||
|
|||||||
+22
-10
@@ -72,15 +72,27 @@ def check_for_updates() -> bool:
|
|||||||
|
|
||||||
|
|
||||||
def notify_update_available():
|
def notify_update_available():
|
||||||
"""通知用户有更新可用"""
|
"""通知用户有更新可用(前端弹窗 + 控制台)"""
|
||||||
current_version = get_current_version()
|
current_version = get_current_version()
|
||||||
version_str = f" (当前版本: {current_version})" if current_version else ""
|
version_str = f" (当前版本: {current_version})" if current_version else ""
|
||||||
|
|
||||||
print("\n" + "="*60)
|
print(f"[comfyui_o1key] 有新版本可用{version_str}")
|
||||||
print(f"🎉 Comfyui_o1key 有新版本可用{version_str}")
|
|
||||||
print("="*60)
|
try:
|
||||||
print("更新方法:")
|
import threading
|
||||||
print(" Windows: 双击运行 update.bat")
|
from server import PromptServer
|
||||||
print(" Linux/Mac: 运行 ./update.sh")
|
|
||||||
print("或手动执行: git pull origin main")
|
def _send():
|
||||||
print("="*60 + "\n")
|
try:
|
||||||
|
PromptServer.instance.send_sync(
|
||||||
|
"o1key.update_available",
|
||||||
|
{"message": "欢迎使用o1key工作流,祝您马年,马上有福,马上有钱,马到成功!!!"}
|
||||||
|
)
|
||||||
|
print("[o1key] 更新通知已发送到前端")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[o1key] 发送通知失败: {e}")
|
||||||
|
|
||||||
|
# 延迟 5 秒发送,确保前端 WebSocket 已连接
|
||||||
|
threading.Timer(5.0, _send).start()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|||||||
@@ -0,0 +1,122 @@
|
|||||||
|
import { app } from "../../../scripts/app.js";
|
||||||
|
import { api } from "../../../scripts/api.js";
|
||||||
|
|
||||||
|
// 上传单个文件到 ComfyUI input 目录,返回服务端绝对路径
|
||||||
|
async function uploadFile(file) {
|
||||||
|
const formData = new FormData();
|
||||||
|
formData.append("image", file, file.name);
|
||||||
|
const resp = await api.fetchApi("/upload/image", { method: "POST", body: formData });
|
||||||
|
if (!resp.ok) throw new Error(`上传失败: ${file.name}`);
|
||||||
|
const data = await resp.json();
|
||||||
|
const inputDir = await getInputDir();
|
||||||
|
// 拼成绝对路径(Windows 用反斜杠也可以,用正斜杠 Python 也认)
|
||||||
|
return inputDir ? inputDir.replace(/\\/g, "/") + "/" + data.name : data.name;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取 ComfyUI input 目录绝对路径(缓存)
|
||||||
|
let _inputDir = null;
|
||||||
|
async function getInputDir() {
|
||||||
|
if (_inputDir !== null) return _inputDir;
|
||||||
|
try {
|
||||||
|
const resp = await api.fetchApi("/o1key/input_dir");
|
||||||
|
if (resp.ok) _inputDir = (await resp.json()).path;
|
||||||
|
else _inputDir = "";
|
||||||
|
} catch { _inputDir = ""; }
|
||||||
|
return _inputDir;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 创建一个"选择文件"按钮,点击后弹出文件选择框
|
||||||
|
// onPaths(paths: string[]) 回调拿到上传后的路径列表
|
||||||
|
function makeUploadButton(label, accept, multiple, onPaths) {
|
||||||
|
const btn = document.createElement("button");
|
||||||
|
btn.textContent = label;
|
||||||
|
btn.style.cssText =
|
||||||
|
"width:100%;padding:4px 8px;cursor:pointer;margin-top:2px;" +
|
||||||
|
"background:#3a5a3a;color:#ddd;border:1px solid #666;" +
|
||||||
|
"border-radius:4px;font-size:12px;";
|
||||||
|
|
||||||
|
const fileInput = document.createElement("input");
|
||||||
|
fileInput.type = "file";
|
||||||
|
fileInput.multiple = multiple;
|
||||||
|
fileInput.accept = accept;
|
||||||
|
fileInput.style.display = "none";
|
||||||
|
document.body.appendChild(fileInput);
|
||||||
|
|
||||||
|
btn.addEventListener("click", () => fileInput.click());
|
||||||
|
|
||||||
|
fileInput.addEventListener("change", async () => {
|
||||||
|
const files = Array.from(fileInput.files);
|
||||||
|
if (!files.length) return;
|
||||||
|
btn.textContent = "⏳ 上传中...";
|
||||||
|
btn.disabled = true;
|
||||||
|
try {
|
||||||
|
const paths = [];
|
||||||
|
for (const f of files) paths.push(await uploadFile(f));
|
||||||
|
onPaths(paths);
|
||||||
|
btn.textContent = `✅ 已上传 ${files.length} 个`;
|
||||||
|
setTimeout(() => { btn.textContent = label; }, 2000);
|
||||||
|
} catch (e) {
|
||||||
|
console.error("[o1key fileUpload]", e);
|
||||||
|
btn.textContent = "❌ 上传失败";
|
||||||
|
setTimeout(() => { btn.textContent = label; }, 2000);
|
||||||
|
} finally {
|
||||||
|
btn.disabled = false;
|
||||||
|
fileInput.value = "";
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
return btn;
|
||||||
|
}
|
||||||
|
|
||||||
|
const ACCEPT = ".pdf,.txt,.md,.csv,.json,.py,.js,.ts,.html,.xml,.docx,.xlsx,.pptx,.zip,.wav,.mp3,.png,.jpg,.jpeg,.webp";
|
||||||
|
|
||||||
|
app.registerExtension({
|
||||||
|
name: "o1key.fileUpload",
|
||||||
|
|
||||||
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||||
|
if (nodeData.name !== "LoadFile") return;
|
||||||
|
|
||||||
|
const origCreated = nodeType.prototype.onNodeCreated;
|
||||||
|
nodeType.prototype.onNodeCreated = function () {
|
||||||
|
origCreated?.call(this);
|
||||||
|
|
||||||
|
const singleWidget = this.widgets?.find(w => w.name === "单文件路径");
|
||||||
|
const folderWidget = this.widgets?.find(w => w.name === "文件夹路径");
|
||||||
|
|
||||||
|
// "单文件路径"下方加按钮(支持多选,追加路径)
|
||||||
|
if (singleWidget) {
|
||||||
|
const btn = makeUploadButton("📂 选择文件(可多选)", ACCEPT, true, (paths) => {
|
||||||
|
const existing = singleWidget.value?.trim();
|
||||||
|
singleWidget.value = existing
|
||||||
|
? existing + ", " + paths.join(", ")
|
||||||
|
: paths.join(", ");
|
||||||
|
singleWidget.callback?.(singleWidget.value);
|
||||||
|
app.graph.setDirtyCanvas(true);
|
||||||
|
});
|
||||||
|
this.addDOMWidget("upload_single_btn", "btn", btn, {
|
||||||
|
getValue() { return null; },
|
||||||
|
setValue() {},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 清空按钮:同时清空单文件路径和文件夹路径
|
||||||
|
if (singleWidget || folderWidget) {
|
||||||
|
const clearBtn = document.createElement("button");
|
||||||
|
clearBtn.textContent = "🗑 清空文件路径";
|
||||||
|
clearBtn.style.cssText =
|
||||||
|
"width:100%;padding:4px 8px;cursor:pointer;margin-top:2px;" +
|
||||||
|
"background:#5a3a3a;color:#ddd;border:1px solid #666;" +
|
||||||
|
"border-radius:4px;font-size:12px;";
|
||||||
|
clearBtn.addEventListener("click", () => {
|
||||||
|
if (singleWidget) { singleWidget.value = ""; singleWidget.callback?.(""); }
|
||||||
|
if (folderWidget) { folderWidget.value = ""; folderWidget.callback?.(""); }
|
||||||
|
app.graph.setDirtyCanvas(true);
|
||||||
|
});
|
||||||
|
this.addDOMWidget("clear_paths_btn", "btn", clearBtn, {
|
||||||
|
getValue() { return null; },
|
||||||
|
setValue() {},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
},
|
||||||
|
});
|
||||||
@@ -0,0 +1,210 @@
|
|||||||
|
import { app } from "../../../scripts/app.js";
|
||||||
|
import { api } from "../../../scripts/api.js";
|
||||||
|
|
||||||
|
// ── marked.js 懒加载 ──────────────────────────────────────────────────────────
|
||||||
|
let markedReady = null;
|
||||||
|
function loadMarked() {
|
||||||
|
if (markedReady) return markedReady;
|
||||||
|
markedReady = new Promise((resolve) => {
|
||||||
|
if (window.marked) { resolve(window.marked); return; }
|
||||||
|
const s = document.createElement("script");
|
||||||
|
s.src = "https://cdn.jsdelivr.net/npm/marked/marked.min.js";
|
||||||
|
s.onload = () => resolve(window.marked);
|
||||||
|
s.onerror = () => resolve(null);
|
||||||
|
document.head.appendChild(s);
|
||||||
|
});
|
||||||
|
return markedReady;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 节点 UI 构建 ──────────────────────────────────────────────────────────────
|
||||||
|
function buildUI(node) {
|
||||||
|
if (node._spContainer) return;
|
||||||
|
|
||||||
|
const container = document.createElement("div");
|
||||||
|
container.style.cssText =
|
||||||
|
"width:100%;height:100%;box-sizing:border-box;padding:6px;" +
|
||||||
|
"display:flex;flex-direction:column;gap:4px;";
|
||||||
|
|
||||||
|
const toolbar = document.createElement("div");
|
||||||
|
toolbar.style.cssText =
|
||||||
|
"display:flex;justify-content:flex-end;gap:6px;align-items:center;";
|
||||||
|
|
||||||
|
const mdToggle = document.createElement("button");
|
||||||
|
mdToggle.textContent = "MD";
|
||||||
|
mdToggle.title = "切换 Markdown / 纯文本";
|
||||||
|
mdToggle.style.cssText =
|
||||||
|
"font-size:10px;padding:2px 6px;border-radius:3px;cursor:pointer;" +
|
||||||
|
"background:#2a5a2a;color:#ccc;border:1px solid #666;";
|
||||||
|
|
||||||
|
const copyBtn = document.createElement("button");
|
||||||
|
copyBtn.textContent = "复制";
|
||||||
|
copyBtn.style.cssText =
|
||||||
|
"font-size:10px;padding:2px 6px;border-radius:3px;cursor:pointer;" +
|
||||||
|
"background:#444;color:#ccc;border:1px solid #666;";
|
||||||
|
|
||||||
|
toolbar.appendChild(mdToggle);
|
||||||
|
toolbar.appendChild(copyBtn);
|
||||||
|
|
||||||
|
const isDedicatedPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
|
||||||
|
|
||||||
|
const content = document.createElement("div");
|
||||||
|
if (isDedicatedPreview) {
|
||||||
|
content.style.cssText =
|
||||||
|
"flex:1;min-height:0;overflow:hidden;" +
|
||||||
|
"background:#1a1a1a;border:1px solid #444;border-radius:4px;" +
|
||||||
|
"padding:8px;box-sizing:border-box;font-size:13px;line-height:1.6;" +
|
||||||
|
"color:#ddd;white-space:pre-wrap;word-break:break-word;";
|
||||||
|
} else {
|
||||||
|
content.style.cssText =
|
||||||
|
"width:100%;min-height:60px;max-height:480px;overflow-y:auto;" +
|
||||||
|
"background:#1a1a1a;border:1px solid #444;border-radius:4px;" +
|
||||||
|
"padding:8px;box-sizing:border-box;font-size:13px;line-height:1.6;" +
|
||||||
|
"color:#ddd;white-space:pre-wrap;word-break:break-word;";
|
||||||
|
}
|
||||||
|
|
||||||
|
const status = document.createElement("div");
|
||||||
|
status.style.cssText =
|
||||||
|
"font-size:10px;color:#888;text-align:right;min-height:14px;";
|
||||||
|
|
||||||
|
container.appendChild(toolbar);
|
||||||
|
container.appendChild(content);
|
||||||
|
container.appendChild(status);
|
||||||
|
|
||||||
|
node._spContainer = container;
|
||||||
|
node._spContent = content;
|
||||||
|
node._spStatus = status;
|
||||||
|
node._spMdToggle = mdToggle;
|
||||||
|
node._spRawText = "";
|
||||||
|
node._spMarkdown = true;
|
||||||
|
|
||||||
|
mdToggle.addEventListener("click", () => {
|
||||||
|
node._spMarkdown = !node._spMarkdown;
|
||||||
|
mdToggle.style.background = node._spMarkdown ? "#2a5a2a" : "#444";
|
||||||
|
renderContent(node);
|
||||||
|
});
|
||||||
|
|
||||||
|
copyBtn.addEventListener("click", () => {
|
||||||
|
navigator.clipboard.writeText(node._spRawText).then(() => {
|
||||||
|
copyBtn.textContent = "已复制";
|
||||||
|
setTimeout(() => { copyBtn.textContent = "复制"; }, 1500);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
const widget = node.addDOMWidget("stream_preview_widget", "preview", container, {
|
||||||
|
getValue() { return node._spRawText; },
|
||||||
|
setValue(v) { },
|
||||||
|
});
|
||||||
|
widget.computeSize = (width) => {
|
||||||
|
const isDedicatedPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
|
||||||
|
if (isDedicatedPreview) {
|
||||||
|
const nodeHeight = node.size?.[1] ?? 320;
|
||||||
|
const overhead = 60;
|
||||||
|
return [width, Math.max(120, nodeHeight - overhead)];
|
||||||
|
}
|
||||||
|
return [width, 320];
|
||||||
|
};
|
||||||
|
|
||||||
|
loadMarked();
|
||||||
|
}
|
||||||
|
|
||||||
|
async function renderContent(node) {
|
||||||
|
const text = node._spRawText;
|
||||||
|
const el = node._spContent;
|
||||||
|
if (!text) { el.innerHTML = ""; return; }
|
||||||
|
|
||||||
|
if (node._spMarkdown) {
|
||||||
|
const marked = await loadMarked();
|
||||||
|
if (marked) {
|
||||||
|
el.style.whiteSpace = "normal";
|
||||||
|
el.innerHTML = marked.parse(text);
|
||||||
|
} else {
|
||||||
|
el.style.whiteSpace = "pre-wrap";
|
||||||
|
el.textContent = text;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
el.style.whiteSpace = "pre-wrap";
|
||||||
|
el.textContent = text;
|
||||||
|
}
|
||||||
|
const isPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
|
||||||
|
if (!isPreview) el.scrollTop = el.scrollHeight;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 流式事件监听 ──────────────────────────────────────────────────────────────
|
||||||
|
api.addEventListener("o1key.stream_token", (event) => {
|
||||||
|
const { node_id, token, done } = event.detail;
|
||||||
|
const node = app.graph.getNodeById(parseInt(node_id));
|
||||||
|
if (!node) return;
|
||||||
|
|
||||||
|
buildUI(node);
|
||||||
|
|
||||||
|
if (done) {
|
||||||
|
node._spStreaming = false;
|
||||||
|
node._spStatus.textContent = "生成完成";
|
||||||
|
node._spStatus.style.color = "#4a4";
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 第一个 token 到来时清空上一次内容
|
||||||
|
if (!node._spStreaming) {
|
||||||
|
node._spStreaming = true;
|
||||||
|
node._spRawText = "";
|
||||||
|
}
|
||||||
|
|
||||||
|
node._spRawText += token;
|
||||||
|
node._spStatus.textContent = "生成中…";
|
||||||
|
node._spStatus.style.color = "#a84";
|
||||||
|
renderContent(node);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ── 节点注册 ──────────────────────────────────────────────────────────────────
|
||||||
|
app.registerExtension({
|
||||||
|
name: "comfyui_o1key.streamPreview",
|
||||||
|
|
||||||
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||||
|
if (nodeData.name !== "StreamPreview") return;
|
||||||
|
|
||||||
|
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||||||
|
nodeType.prototype.onNodeCreated = function () {
|
||||||
|
if (origOnNodeCreated) origOnNodeCreated.apply(this, arguments);
|
||||||
|
buildUI(this);
|
||||||
|
};
|
||||||
|
|
||||||
|
nodeType.prototype.onResize = function () {
|
||||||
|
this.setDirtyCanvas(true, false);
|
||||||
|
};
|
||||||
|
|
||||||
|
const origOnExecuted = nodeType.prototype.onExecuted;
|
||||||
|
nodeType.prototype.onExecuted = function (message) {
|
||||||
|
if (origOnExecuted) origOnExecuted.apply(this, arguments);
|
||||||
|
buildUI(this);
|
||||||
|
|
||||||
|
const texts = message?.text;
|
||||||
|
if (!texts || texts.length === 0) return;
|
||||||
|
|
||||||
|
this._spRawText = texts[0];
|
||||||
|
this._spStatus.textContent = "完成";
|
||||||
|
this._spStatus.style.color = "#4a4";
|
||||||
|
renderContent(this);
|
||||||
|
this.setDirtyCanvas(true, true);
|
||||||
|
};
|
||||||
|
|
||||||
|
const origOnSerialize = nodeType.prototype.onSerialize;
|
||||||
|
nodeType.prototype.onSerialize = function (o) {
|
||||||
|
if (origOnSerialize) origOnSerialize.apply(this, arguments);
|
||||||
|
o.sp_text = this._spRawText || "";
|
||||||
|
o.sp_markdown = this._spMarkdown !== false;
|
||||||
|
};
|
||||||
|
|
||||||
|
const origOnConfigure = nodeType.prototype.onConfigure;
|
||||||
|
nodeType.prototype.onConfigure = function (o) {
|
||||||
|
if (origOnConfigure) origOnConfigure.apply(this, arguments);
|
||||||
|
buildUI(this);
|
||||||
|
if (o.sp_text) {
|
||||||
|
this._spRawText = o.sp_text;
|
||||||
|
this._spMarkdown = o.sp_markdown !== false;
|
||||||
|
this._spMdToggle.style.background = this._spMarkdown ? "#2a5a2a" : "#444";
|
||||||
|
renderContent(this);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
},
|
||||||
|
});
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
import { api } from "../../../scripts/api.js";
|
||||||
|
|
||||||
|
api.addEventListener("o1key.update_available", (event) => {
|
||||||
|
const message = event.detail?.message || "欢迎使用o1key工作流";
|
||||||
|
|
||||||
|
const style = document.createElement("style");
|
||||||
|
style.textContent = `
|
||||||
|
@keyframes o1key-fadein {
|
||||||
|
from { opacity: 0; transform: translateY(16px) scale(0.97); }
|
||||||
|
to { opacity: 1; transform: translateY(0) scale(1); }
|
||||||
|
}
|
||||||
|
.o1key-toast-close:hover { color: #fff !important; }
|
||||||
|
`;
|
||||||
|
document.head.appendChild(style);
|
||||||
|
|
||||||
|
const toast = document.createElement("div");
|
||||||
|
toast.style.cssText = `
|
||||||
|
position: fixed;
|
||||||
|
bottom: 28px;
|
||||||
|
left: 28px;
|
||||||
|
background: linear-gradient(135deg, #0d3320 0%, #145a32 60%, #1e8449 100%);
|
||||||
|
color: #d5f5e3;
|
||||||
|
border: 1px solid #27ae60;
|
||||||
|
border-radius: 12px;
|
||||||
|
padding: 18px 20px 16px 20px;
|
||||||
|
font-size: 15px;
|
||||||
|
z-index: 99999;
|
||||||
|
box-shadow: 0 6px 24px rgba(39,174,96,0.35), 0 2px 8px rgba(0,0,0,0.5);
|
||||||
|
max-width: 340px;
|
||||||
|
animation: o1key-fadein 0.4s cubic-bezier(.22,.68,0,1.2);
|
||||||
|
`;
|
||||||
|
|
||||||
|
const header = document.createElement("div");
|
||||||
|
header.style.cssText = `
|
||||||
|
display: flex;
|
||||||
|
align-items: center;
|
||||||
|
justify-content: space-between;
|
||||||
|
margin-bottom: 10px;
|
||||||
|
`;
|
||||||
|
|
||||||
|
const title = document.createElement("span");
|
||||||
|
title.textContent = "🐴 o1key 工作流";
|
||||||
|
title.style.cssText = `font-weight: bold; font-size: 13px; color: #82e0aa; letter-spacing: 0.5px;`;
|
||||||
|
|
||||||
|
const closeBtn = document.createElement("button");
|
||||||
|
closeBtn.textContent = "×";
|
||||||
|
closeBtn.className = "o1key-toast-close";
|
||||||
|
closeBtn.style.cssText = `
|
||||||
|
background: none;
|
||||||
|
border: none;
|
||||||
|
color: #82e0aa;
|
||||||
|
font-size: 20px;
|
||||||
|
cursor: pointer;
|
||||||
|
padding: 0;
|
||||||
|
line-height: 1;
|
||||||
|
`;
|
||||||
|
closeBtn.onclick = () => toast.remove();
|
||||||
|
|
||||||
|
header.appendChild(title);
|
||||||
|
header.appendChild(closeBtn);
|
||||||
|
|
||||||
|
const divider = document.createElement("div");
|
||||||
|
divider.style.cssText = `height: 1px; background: rgba(39,174,96,0.3); margin-bottom: 10px;`;
|
||||||
|
|
||||||
|
const text = document.createElement("div");
|
||||||
|
text.textContent = message;
|
||||||
|
text.style.cssText = `line-height: 1.7; color: #d5f5e3;`;
|
||||||
|
|
||||||
|
toast.appendChild(header);
|
||||||
|
toast.appendChild(divider);
|
||||||
|
toast.appendChild(text);
|
||||||
|
document.body.appendChild(toast);
|
||||||
|
});
|
||||||
@@ -124,7 +124,7 @@ app.registerExtension({
|
|||||||
// 节点头部 + 其他 widget 的高度
|
// 节点头部 + 其他 widget 的高度
|
||||||
// LiteGraph 节点头部约 30px,每个普通 widget 约 24px
|
// LiteGraph 节点头部约 30px,每个普通 widget 约 24px
|
||||||
const NON_VIDEO_WIDGETS = (this.widgets?.filter(
|
const NON_VIDEO_WIDGETS = (this.widgets?.filter(
|
||||||
(w) => w.name !== "video_preview"
|
(w) => w.name !== "video_preview_widget"
|
||||||
).length ?? 0);
|
).length ?? 0);
|
||||||
const headerH = 58 + NON_VIDEO_WIDGETS * 24;
|
const headerH = 58 + NON_VIDEO_WIDGETS * 24;
|
||||||
|
|
||||||
@@ -138,11 +138,8 @@ app.registerExtension({
|
|||||||
const origOnResize = nodeType.prototype.onResize;
|
const origOnResize = nodeType.prototype.onResize;
|
||||||
nodeType.prototype.onResize = function (size) {
|
nodeType.prototype.onResize = function (size) {
|
||||||
if (origOnResize) origOnResize.apply(this, arguments);
|
if (origOnResize) origOnResize.apply(this, arguments);
|
||||||
if (this._videoAspectRatio && this._videoEl) {
|
if (this._videoAspectRatio) {
|
||||||
// 用新宽度重新计算正确高度,避免拉伸/压缩
|
this._resizeToVideo();
|
||||||
const innerW = Math.max(size[0] - 16, 10);
|
|
||||||
const videoH = Math.round(innerW * this._videoAspectRatio);
|
|
||||||
this._videoEl.style.height = videoH + "px";
|
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
},
|
},
|
||||||
|
|||||||
Reference in New Issue
Block a user