diff --git a/.config b/.config new file mode 100644 index 0000000..cece251 --- /dev/null +++ b/.config @@ -0,0 +1 @@ +O1KEY_API_KEY= \ No newline at end of file diff --git a/.gitignore b/.gitignore index 4b2985f..a614a27 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,3 @@ -# .config 文件包含敏感信息，不提交到版本控制 -# 用户可通过 setup_api_key.bat 自动创建本地配置 -.config - # Python 缓存 __pycache__/ *.py[cod] diff --git a/__init__.py b/__init__.py index d441f6c..739b414 100644 --- a/__init__.py +++ b/__init__.py @@ -12,16 +12,15 @@ Comfyui_o1key - ComfyUI 自定义节点集合 # 检查更新（仅在启动时检查一次） try: from .utils.update_checker import check_for_updates, notify_update_available - - if check_for_updates(): - notify_update_available() + + notify_update_available() # TODO: 测试用，改回 if check_for_updates(): notify_update_available() except Exception: # 静默失败，不影响插件加载 pass 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 请求超时，请稍后重试或检查网络。" @@ -55,8 +54,6 @@ def _wrap_generate_for_error_display(cls, attr="generate"): _wrap_generate_for_error_display(NanoBananaPro) _wrap_generate_for_error_display(BatchNanoBananaPro) -_wrap_generate_for_error_display(QuanNengShengTu) -_wrap_generate_for_error_display(BatchQuanNengShengTu, "process_batch") # ComfyUI 节点注册 NODE_CLASS_MAPPINGS = { @@ -74,12 +71,12 @@ NODE_CLASS_MAPPINGS = { "KlingVideo": KlingVideo, "KlingFirstLastFrame": KlingFirstLastFrame, "KlingMotionControlTest": KlingMotionControlTest, - "QuanNengShengTu": QuanNengShengTu, - "BatchQuanNengShengTu": BatchQuanNengShengTu, "AspectRatioPreset": AspectRatioPreset, "MultiResPreview": MultiResPreview, "BatchImagesO1key": BatchImagesO1key, "Seedance": Seedance, + "SeedanceMultiModal": SeedanceMultiModal, + "StreamPreview": StreamPreview, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -90,21 +87,35 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageStitchPro": "图像拼接 Pro", "SaveCleanImage": "保存图像（防AI识别）", "BatchCleanMetadata": "批量任务（防AI识别）", - "VideoPreview": "视频预览", + "VideoPreview": "预览视频", "GoogleVeo": "Google Veo - ab", "FluxImageEdit": "Flux2 图像编辑", "UniversalLLMChat": "全能LLM对话助手", "KlingVideo": "文/图生视频 自研模型", "KlingFirstLastFrame": "首尾帧生视频 自研模型", "KlingMotionControlTest": "动作控制 自研模型", - "QuanNengShengTu": "全能生图", - "BatchQuanNengShengTu": "全能生图（批量）", "AspectRatioPreset": "图片宽高比预设", "MultiResPreview": "预览图像（v2）", "BatchImagesO1key": "加载图像（批量）", "Seedance": "Seedance 视频生成", + "SeedanceMultiModal": "Seedance 多模态参考生视频", + "StreamPreview": "流式文本预览", } WEB_DIRECTORY = "./web" __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 diff --git a/clients/base_client.py b/clients/base_client.py index 701ea33..cadef18 100644 --- a/clients/base_client.py +++ b/clients/base_client.py @@ -79,6 +79,15 @@ class BaseAPIClient(ABC): """ 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]: """ 获取请求头 @@ -142,67 +151,132 @@ class BaseAPIClient(ABC): ) -> Dict[str, Any]: """ 发送异步 HTTP 请求（带详细计时） - + Args: endpoint: API 端点 request_body: 请求体 session: aiohttp 会话（可选） use_bearer_token: 是否使用 Bearer Token 认证 - timeout: 超时时间（秒）- 已废弃，由服务器端控制 - + timeout: 超时时间（秒），默认 900 秒 + Returns: 响应 JSON - + Raises: RuntimeError: 请求失败时 + InterruptProcessingException: 用户点击终止按钮时 """ 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}" headers = self.get_headers(use_bearer_token) - + # 检查请求大小 self.check_request_size(request_body) - + close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() 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() - - async with session.post(url, json=request_body, headers=headers) as response: + async with session.post(url, json=request_body, headers=headers, timeout=_aiohttp_timeout) as response: connect_time = time.time() - connect_start - + if response.status != 200: error_text = await response.text() raise RuntimeError(error_text) - - # 接收响应体 + wait_start = time.time() response_data = await response.json() download_time = time.time() - wait_start - - # 附加计时信息到响应数据（供上层使用） + response_size = len(str(response_data)) if not isinstance(response_data, dict): response_data = {"data": response_data} - - # 将计时信息存储在响应的元数据中 + response_data["_timing"] = { "connect_time": connect_time, "download_time": download_time, "response_size": response_size } - 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: if close_session: await session.close() - + async def request_get_async( self, endpoint: str, @@ -230,9 +304,9 @@ class BaseAPIClient(ABC): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True - + try: async with session.get(url, headers=headers) as response: if response.status != 200: @@ -327,9 +401,7 @@ class BaseAPIClient(ABC): total = len(requests) # 创建无限制的连接器 - connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) - - async with aiohttp.ClientSession(connector=connector) as session: + async with self._make_session() as session: tasks = [] for req in requests: diff --git a/clients/gemini_client.py b/clients/gemini_client.py index d447b49..58aac61 100644 --- a/clients/gemini_client.py +++ b/clients/gemini_client.py @@ -141,7 +141,6 @@ class GeminiAPIClient(BaseAPIClient): resolution: str = "2K", enable_grounding: bool = False, enable_image_search: bool = False, - candidate_count: int = 1, **kwargs ) -> Dict[str, Any]: """ @@ -154,7 +153,6 @@ class GeminiAPIClient(BaseAPIClient): resolution: 分辨率 enable_grounding: 是否启用 Google Search Grounding enable_image_search: 是否同时启用 Google Image Search（仅 Gemini 3.1 Flash 支持） - candidate_count: 单次请求返回的候选图数量，默认 1 Returns: 请求体字典 @@ -184,7 +182,6 @@ class GeminiAPIClient(BaseAPIClient): } ], "generationConfig": { - "candidateCount": candidate_count, "responseModalities": ["TEXT", "IMAGE"], "imageConfig": { "aspectRatio": aspect_ratio, @@ -303,7 +300,7 @@ class GeminiAPIClient(BaseAPIClient): # 需要关闭 session 的标记 close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -438,7 +435,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request: bool = False, enable_grounding: bool = False, enable_image_search: bool = False, - candidate_count: int = 1 ) -> tuple[List[Image.Image], Dict[str, Any]]: """ 单次异步生成请求（极简单行日志） @@ -465,7 +461,7 @@ class GeminiAPIClient(BaseAPIClient): 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. 构建请求 ========== build_start = time.time() @@ -477,7 +473,6 @@ class GeminiAPIClient(BaseAPIClient): resolution=resolution, enable_grounding=enable_grounding, enable_image_search=enable_image_search, - candidate_count=candidate_count ) build_time = time.time() - build_start @@ -499,7 +494,7 @@ class GeminiAPIClient(BaseAPIClient): safe_request = _truncate_base64_req(request_body) print( f"\n{'='*60}\n" - f"[请求体日志] 任务 {task_prefix or '?'} 发送请求体：\n" + f"[请求体日志] {task_prefix}发送请求体：\n" f"端点: {endpoint}\n" f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n" f"{'='*60}\n" @@ -541,7 +536,7 @@ class GeminiAPIClient(BaseAPIClient): safe_response = _truncate_base64(response) print( 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"{'='*60}\n" ) @@ -554,7 +549,7 @@ class GeminiAPIClient(BaseAPIClient): except Exception as e: parse_time = time.time() - parse_start 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 parse_time = time.time() - parse_start @@ -578,7 +573,7 @@ class GeminiAPIClient(BaseAPIClient): 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 @@ -605,7 +600,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request: bool = False, enable_grounding: bool = False, enable_image_search: bool = False, - candidate_count: int = 1 ) -> List[Image.Image]: """ 批量全并发生成 - 改进版：支持分批处理和内存管理 @@ -622,7 +616,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request: 是否打印发送的请求体 enable_grounding: 是否启用 Google Search Grounding enable_image_search: 是否同时启用 Google Image Search - candidate_count: 单次请求返回的候选图数量 Returns: 生成的图像列表 @@ -635,28 +628,23 @@ class GeminiAPIClient(BaseAPIClient): success_count = 0 fail_count = 0 first_error = None # 保存第一个错误 - + # 分批处理配置 max_concurrent = 10 # 最大并发数 save_batch_size = 10 # 分批保存大小 - + # 计算需要多少批次 num_batches = (batch_size + max_concurrent - 1) // max_concurrent - - print(f"GeminiClient: 批量生成 {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: # 分批执行 for batch_idx in range(num_batches): batch_start = batch_idx * max_concurrent batch_end = min(batch_start + max_concurrent, batch_size) 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 = [] for i in range(batch_size_current): @@ -675,7 +663,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request=debug_request, enable_grounding=enable_grounding, enable_image_search=enable_image_search, - candidate_count=candidate_count ), name=f"task_{task_index}" ) @@ -698,14 +685,11 @@ class GeminiAPIClient(BaseAPIClient): all_images.append(img) success_count += 1 - + # 通知进度 if progress_callback: progress_callback(completed, batch_size, True, None) - - # 每成功生成一张图片就打印日志 - print(f"GeminiClient: 任务 {completed}/{batch_size} 成功生成图片 ✓") - + except Exception as e: fail_count += 1 # 保存第一个错误（用于后续抛出） @@ -719,25 +703,20 @@ class GeminiAPIClient(BaseAPIClient): # 当前批次完成后，立即清理内存 if batch_images: - print(f"GeminiClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张图片") - - # 强制垃圾回收，释放内存 import gc gc.collect() - - # 短暂暂停，让系统处理内存 await asyncio.sleep(0.1) - + # 清空当前批次图片引用，帮助垃圾回收 batch_images = [] - + # 最终结果检查 if not all_images: # 如果有保存的原始错误，直接抛出原始错误 if first_error: raise first_error raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败") - + print(f"GeminiClient: 批量生成完成，成功 {success_count}/{batch_size}，失败 {fail_count}") return all_images @@ -754,7 +733,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request: bool = False, enable_grounding: bool = False, enable_image_search: bool = False, - candidate_count: int = 1 ) -> List[Image.Image]: """ 同步生成接口（用于 ComfyUI） @@ -771,7 +749,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request: 是否打印发送的请求体 enable_grounding: 是否启用 Google Search Grounding enable_image_search: 是否同时启用 Google Image Search - candidate_count: 单次请求返回的候选图数量 Returns: 生成的图像列表 @@ -788,7 +765,6 @@ class GeminiAPIClient(BaseAPIClient): debug_request=debug_request, enable_grounding=enable_grounding, enable_image_search=enable_image_search, - candidate_count=candidate_count ) return self.run_async_in_thread(coro) @@ -837,13 +813,11 @@ class GeminiAPIClient(BaseAPIClient): fail_count = 0 first_error = None # 保存第一个错误 total_tasks = len(prompts) * images_per_prompt - + # 分批处理配置 max_concurrent = 10 # 最大并发数 - - print(f"GeminiClient: 多提示词批量生成，共 {total_tasks} 个任务，{len(prompts)} 个提示词，每个 {images_per_prompt} 张") - - 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: # 创建所有任务 @@ -874,15 +848,12 @@ class GeminiAPIClient(BaseAPIClient): # 分批处理：每10个任务为一组 batch_size = max_concurrent num_batches = (total_tasks + batch_size - 1) // batch_size - + for batch_idx in range(num_batches): batch_start = batch_idx * batch_size batch_end = min(batch_start + batch_size, total_tasks) 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 = [] @@ -898,36 +869,26 @@ class GeminiAPIClient(BaseAPIClient): all_images.append(img) success_count += 1 - + # 通知进度 if progress_callback: progress_callback(completed, total_tasks, True, None) - - # 每成功生成一张图片就打印日志 - print(f"GeminiClient: 任务 {completed}/{total_tasks} 成功生成图片 ✓") - + except Exception as e: fail_count += 1 # 保存第一个错误（用于后续抛出） if first_error is None: first_error = e error_msg = str(e) - + # 传递完整的错误信息（用于排查问题） if progress_callback: progress_callback(completed, total_tasks, False, error_msg) - - print(f"GeminiClient: 任务 {completed}/{total_tasks} 失败 ✗") - + # 当前批次完成后，立即清理内存 if batch_images: - print(f"GeminiClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张图片") - - # 强制垃圾回收，释放内存 import gc gc.collect() - - # 短暂暂停，让系统处理内存 await asyncio.sleep(0.1) # 清空当前批次图片引用，帮助垃圾回收 @@ -938,8 +899,7 @@ class GeminiAPIClient(BaseAPIClient): if first_error: raise first_error raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败") - - print(f"GeminiClient: 多提示词批量生成完成，成功 {success_count}/{total_tasks}，失败 {fail_count}") + return all_images def generate_multi_prompts_sync( diff --git a/clients/kling_client.py b/clients/kling_client.py index 8ba35b6..a965a01 100644 --- a/clients/kling_client.py +++ b/clients/kling_client.py @@ -137,7 +137,7 @@ class KlingClient: on_progress: Optional[Callable[[int], None]] = None, ) -> str: """提交 → 轮询 → 下载，返回本地文件路径""" - connector = aiohttp.TCPConnector(force_close=True) + connector = aiohttp.TCPConnector(ssl=False, force_close=True) async with aiohttp.ClientSession(connector=connector) as session: if on_stage: on_stage("submitting") @@ -196,7 +196,7 @@ class KlingClient: "Content-Type": "application/json"} 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: # 1. 提交 diff --git a/clients/openai_client.py b/clients/openai_client.py index ba99088..72ad2ce 100644 --- a/clients/openai_client.py +++ b/clients/openai_client.py @@ -309,7 +309,7 @@ class OpenAIAPIClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -677,7 +677,7 @@ class OpenAIAPIClient(BaseAPIClient): 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: for batch_idx in range(num_batches): diff --git a/clients/seedance_client.py b/clients/seedance_client.py index 2edcba9..5dd20ef 100644 --- a/clients/seedance_client.py +++ b/clients/seedance_client.py @@ -169,7 +169,7 @@ class SeedanceClient: on_progress: Optional[Callable[[int], None]] = None, ) -> tuple: """提交 → 轮询 → 下载，返回 (本地视频路径, 末帧图片URL或None)""" - connector = aiohttp.TCPConnector(force_close=True) + connector = aiohttp.TCPConnector(ssl=False, force_close=True) async with aiohttp.ClientSession(connector=connector) as session: # 提交 diff --git a/clients/sora_client.py b/clients/sora_client.py index 4d76bba..c477920 100644 --- a/clients/sora_client.py +++ b/clients/sora_client.py @@ -132,7 +132,7 @@ class SoraClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -174,7 +174,7 @@ class SoraClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True interval = self.POLL_INITIAL_INTERVAL @@ -242,7 +242,7 @@ class SoraClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -299,7 +299,7 @@ class SoraClient(BaseAPIClient): """ async def _run(): - connector = aiohttp.TCPConnector(limit=0) + connector = aiohttp.TCPConnector(ssl=False, limit=0) async with aiohttp.ClientSession(connector=connector) as session: # 1. 提交任务 if on_stage: @@ -408,7 +408,7 @@ class SoraClient(BaseAPIClient): 成功生成的视频路径列表 """ 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: tasks = [ diff --git a/clients/veo_client.py b/clients/veo_client.py index 01d3e9d..b009f66 100644 --- a/clients/veo_client.py +++ b/clients/veo_client.py @@ -160,7 +160,7 @@ class VeoClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -202,7 +202,7 @@ class VeoClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True interval = self.POLL_INITIAL_INTERVAL @@ -265,7 +265,7 @@ class VeoClient(BaseAPIClient): close_session = False if session is None: - session = aiohttp.ClientSession() + session = self._make_session() close_session = True try: @@ -318,7 +318,7 @@ class VeoClient(BaseAPIClient): """ async def _run(): - connector = aiohttp.TCPConnector(limit=0) + connector = aiohttp.TCPConnector(ssl=False, limit=0) async with aiohttp.ClientSession(connector=connector) as session: # 1. 提交任务 if on_stage: @@ -383,7 +383,7 @@ class VeoClient(BaseAPIClient): """ async def _run(): batch_size = len(save_paths) - connector = aiohttp.TCPConnector(limit=0) + connector = aiohttp.TCPConnector(ssl=False, limit=0) async def generate_one(save_path: str): return await self._generate_one_video_async( diff --git a/nodes/__init__.py b/nodes/__init__.py index da54a06..5799075 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -3,6 +3,7 @@ 包含所有 ComfyUI 自定义节点的实现 """ +from .stream_preview import StreamPreview from .nano_banana_pro import NanoBananaPro from .batch_nano_banana_pro import BatchNanoBananaPro from .google_gemini import GoogleGemini @@ -14,12 +15,8 @@ from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest from .veo_video import GoogleVeo from .flux_edit import FluxImageEdit 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 .batch_images_o1key import BatchImagesO1key -from .nano_banana_v2 import NanaBananaV2 -from .batch_nano_banana_v2 import BatchNanaBananaV2 -from .seedance_video import Seedance +from .seedance_video import Seedance, SeedanceMultiModal -__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'] diff --git a/nodes/batch_nano_banana_pro.py b/nodes/batch_nano_banana_pro.py index 797f205..eff51b0 100644 --- a/nodes/batch_nano_banana_pro.py +++ b/nodes/batch_nano_banana_pro.py @@ -236,18 +236,6 @@ class BatchNanoBananaPro: "分辨率": (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": "关闭" }), @@ -298,11 +286,6 @@ class BatchNanoBananaPro: "保存路径": ("STRING", { "default": "", "multiline": False - }), - "跳过错误": ("BOOLEAN", { - "default": False, - "label_on": "打开", - "label_off": "关闭" }) }, "optional": optional_inputs @@ -324,8 +307,6 @@ class BatchNanoBananaPro: 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, @@ -334,48 +315,25 @@ class BatchNanoBananaPro: ) -> List[List[ImageInfo]]: """ 加载所有文件夹中的图片 - + Args: folder1-9: 文件夹路径 - enable_scaling: 是否启用像素缩放 - target_megapixels: 目标像素数（百万像素） - + Returns: 图片列表的列表 """ 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 - ) - # 创建新的 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) - else: - # 空文件夹，静默跳过 - pass except ValueError as e: print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}") - + return all_images def _create_pairs( @@ -596,8 +554,7 @@ class BatchNanoBananaPro: total_tasks = len(pairs) - # 保持并发数为10不变（按用户要求） - max_concurrent = 10 + max_concurrent = 50 # 分批保存的批次大小（与并发数一致） save_batch_size = 10 @@ -627,7 +584,7 @@ class BatchNanoBananaPro: if num_batches > 1: 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: # 分批处理：每批最多10个任务 @@ -761,14 +718,11 @@ class BatchNanoBananaPro: 文件夹7: str, 文件夹8: str, 文件夹9: str, - 像素缩放: bool, - 分辨率像素: float, seed: int, 图片配对模式: str, 模型: str, 宽高比: str, 分辨率: str, - 跳过错误: bool = False, 保存路径: str = "", **kwargs ) -> Tuple[torch.Tensor]: @@ -778,8 +732,6 @@ class BatchNanoBananaPro: Args: prompt: 提示词 文件夹1-9: 图片文件夹路径 - 像素缩放: 是否启用像素缩放 - 分辨率像素: 目标像素数（百万像素） seed: 随机种子 保存路径: 输出保存路径 图片配对模式: 1:1 或 1*N @@ -840,7 +792,6 @@ class BatchNanoBananaPro: print("BatchNanoBananaPro: 开始加载图片...") image_lists = self._load_folders( 文件夹1, 文件夹2, 文件夹3, 文件夹4, - 像素缩放, 分辨率像素, 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9 ) @@ -856,10 +807,6 @@ class BatchNanoBananaPro: if key in kwargs and kwargs[key] is not None: pil_images = tensor_to_pil(kwargs[key]) for j, img in enumerate(pil_images): - # 如果启用像素缩放，也对参考图进行缩放 - if 像素缩放: - img = self.resize_to_megapixels(img, 分辨率像素) - manual_images.append( ImageInfo( image=img, @@ -976,9 +923,9 @@ class BatchNanoBananaPro: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_async_in_thread) try: - results = future.result(timeout=3600) # 1小时超时 + results = future.result(timeout=900) # 900秒超时 except TimeoutError: - print("BatchNanoBananaPro: 任务执行超时（1小时）") + print("BatchNanoBananaPro: 任务执行超时（900秒）") raise RuntimeError("任务执行超时，请减少任务数量或检查网络连接") except Exception as e: # 即使失败，也尝试返回部分结果 @@ -1078,24 +1025,12 @@ class BatchNanoBananaPro: if str(e) == "未授权！": print("请联系作者授权后方可使用！") 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 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 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 finally: diff --git a/nodes/batch_nano_banana_v2.py b/nodes/batch_nano_banana_v2.py deleted file mode 100644 index 3168829..0000000 --- a/nodes/batch_nano_banana_v2.py +++ /dev/null @@ -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() diff --git a/nodes/batch_quan_neng_sheng_tu.py b/nodes/batch_quan_neng_sheng_tu.py deleted file mode 100644 index b61fdd9..0000000 --- a/nodes/batch_quan_neng_sheng_tu.py +++ /dev/null @@ -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 - diff --git a/nodes/load_file.py b/nodes/load_file.py index 7d432b3..caf07e2 100644 --- a/nodes/load_file.py +++ b/nodes/load_file.py @@ -1,146 +1,125 @@ """ -LoadFile 节点 -ComfyUI 自定义节点,用于加载文件并转换为 FILE 类型数据 +LoadFile 节点（增强版） +支持单文件路径和文件夹路径，输出 FILE_LIST 类型供全能LLM等节点使用 """ import base64 import os 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: """ - LoadFile 节点 - - 功能： - - 从文件系统加载文件 - - 支持 PDF 和 TXT 文件 - - 转换为 FILE 类型数据（包含 base64 编码内容） - - 验证文件大小和格式 + 加载文件节点 + + - 单文件路径：加载指定文件 + - 文件夹路径：加载文件夹内所有支持的文件（非递归） + - 两者可同时使用，结果合并输出 + - 输出 FILE_LIST 类型，可直接连接到全能LLM对话助手 """ - + @classmethod def INPUT_TYPES(cls): - """ - 定义输入参数 - """ return { - "required": { - "文件路径": ("STRING", { + "required": {}, + "optional": { + "单文件路径": ("STRING", { "default": "", - "multiline": False - }) - } + "multiline": False, + "placeholder": "文件完整路径，多个文件用英文逗号分隔", + }), + "文件夹路径": ("STRING", { + "default": "", + "multiline": False, + "placeholder": "文件夹路径，自动读取其中所有支持的文件", + }), + }, } - - # 返回值类型 - RETURN_TYPES = ("FILE", "STRING") - RETURN_NAMES = ("文件", "文件信息") - - # 执行函数名 + + RETURN_TYPES = ("FILE_LIST", "STRING") + RETURN_NAMES = ("文件列表", "文件信息") FUNCTION = "load_file" - - # 节点分类 CATEGORY = "file/input" - - def load_file(self, 文件路径: str) -> Tuple[FileData, str]: - """ - 加载文件并转换为 FILE 类型 - - Args: - 文件路径: 文件的完整路径（支持绝对路径和相对路径） - - Returns: - (FileData, 文件信息预览) - - Raises: - ValueError: 文件不存在、不支持的文件类型或文件过大 - """ - try: - # 清理路径（去除空格和引号） - file_path = 文件路径.strip().strip('"').strip("'") - - if not file_path: - raise ValueError("文件路径不能为空") - - # 转换为 Path 对象 - path = Path(file_path) - - # 如果是相对路径，转换为绝对路径 - if not path.is_absolute(): - # 相对于当前工作目录 - path = Path.cwd() / path - - # 验证文件是否存在 - if not path.exists(): - raise ValueError(f"文件不存在: {file_path}") - - if not path.is_file(): - raise ValueError(f"路径不是文件: {file_path}") - - # 获取文件信息 - extension = path.suffix.lower() - filename = path.stem - file_size = path.stat().st_size - - # 验证文件类型 - if extension not in DOCUMENT_MIME_TYPES: - supported_types = ", ".join(DOCUMENT_MIME_TYPES.keys()) + + def load_file(self, 单文件路径: str = "", 文件夹路径: str = "") -> Tuple[FileList, str]: + collected: List[Path] = [] + + # 1. 单文件路径（逗号分隔，支持多个） + if 单文件路径.strip(): + for raw in 单文件路径.split(","): + p = Path(raw.strip().strip('"').strip("'")) + if not p.is_absolute(): + p = Path.cwd() / p + if not p.exists(): + raise ValueError(f"文件不存在: {p}") + if not p.is_file(): + raise ValueError(f"路径不是文件: {p}") + collected.append(p) + + # 2. 文件夹路径 + if 文件夹路径.strip(): + folder = Path(文件夹路径.strip().strip('"').strip("'")) + if not folder.is_absolute(): + folder = Path.cwd() / folder + if not folder.exists(): + raise ValueError(f"文件夹不存在: {folder}") + if not folder.is_dir(): + raise ValueError(f"路径不是文件夹: {folder}") + for p in sorted(folder.iterdir()): + if p.is_file() and p.suffix.lower() in DOCUMENT_MIME_TYPES: + collected.append(p) + if not collected: + raise ValueError(f"文件夹中没有支持的文件: {folder}") + + if not collected: + raise ValueError("请至少提供一个文件路径或文件夹路径") + + # 去重（保持顺序） + seen = set() + unique: List[Path] = [] + for p in collected: + key = str(p.resolve()) + if key not in seen: + seen.add(key) + unique.append(p) + + # 大小检查 & 读取 + 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( - f"不支持的文件类型: {extension}\n" - f"支持的类型: {supported_types}" + f"文件 {p.name} 大小 {file_size / 1024 / 1024:.1f}MB 超过单文件 50MB 限制" ) - - # 获取 MIME 类型 - mime_type = DOCUMENT_MIME_TYPES[extension] - - # 验证文件大小 - size_limit = FILE_SIZE_LIMITS.get(extension, 20 * 1024 * 1024) - if file_size > size_limit: - raise ValueError( - f"文件过大 ({file_size / 1024 / 1024:.2f}MB)，" - f"最大支持 {size_limit / 1024 / 1024:.0f}MB" - ) - - # 读取文件并转换为 base64 - print(f"LoadFile: 正在加载文件 {filename}{extension}") - print(f"LoadFile: 文件大小 = {file_size / 1024:.2f}KB") - - with open(path, "rb") as f: - file_bytes = f.read() - - # Base64 编码 - b64_str = base64.b64encode(file_bytes).decode("utf-8") - - # 创建 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 + total_size += file_size + if total_size > TOTAL_FILE_SIZE_LIMIT: + raise ValueError(f"所有文件总大小超过 50MB 限制") + + mime = DOCUMENT_MIME_TYPES[ext] + with open(p, "rb") as f: + b64 = base64.b64encode(f.read()).decode("utf-8") + + file_list.append(FileData( + path=str(p), + filename=p.stem, + extension=ext, + mime_type=mime, + data=b64, + size=file_size, + )) + info_lines.append(f" {p.name} ({file_size / 1024:.1f}KB, {mime})") + print(f"LoadFile: 加载 {p.name} ({file_size / 1024:.1f}KB)") + + info = f"共 {len(file_list)} 个文件，总大小 {total_size / 1024:.1f}KB\n" + "\n".join(info_lines) + return (file_list, info) diff --git a/nodes/nano_banana_pro.py b/nodes/nano_banana_pro.py index 9d4f01f..7ace0eb 100644 --- a/nodes/nano_banana_pro.py +++ b/nodes/nano_banana_pro.py @@ -52,7 +52,7 @@ except ImportError: # ============================================================================ # 是否启用调试日志（打印完整的 API 响应内容） # 设置为 True 以启用调试日志，False 以禁用 -DEBUG_LOG_ENABLED = False +DEBUG_LOG_ENABLED = True # 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断） # 设置为 True 以启用请求体日志，False 以禁用 REQUEST_LOG_ENABLED = False @@ -177,18 +177,6 @@ class NanoBananaPro: "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": "关闭" }), @@ -199,11 +187,6 @@ class NanoBananaPro: "default": 0, "min": 0, "max": 0xffffffffffffffff - }), - "跳过错误": ("BOOLEAN", { - "default": False, - "label_on": "打开", - "label_off": "关闭" }) }, "optional": optional_inputs @@ -309,14 +292,16 @@ class NanoBananaPro: global_task_index: int, enable_grounding: bool = False, enable_image_search: bool = False, + save_to_disk: bool = True, ) -> dict: - """执行单个生成任务，生成后立即保存到磁盘""" + """执行单个生成任务""" result = { "global_task_index": global_task_index, "prompt": prompt, "success": False, "generated_count": 0, "saved_files": [], + "output_images": [], "error": None } @@ -335,20 +320,23 @@ class NanoBananaPro: ) 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 # 释放内存 - + if save_to_disk: + 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 + else: + result["output_images"] = images_list + result["success"] = True result["generated_count"] = len(images_list) except Exception as e: result["error"] = str(e) - + return result async def _process_batch_async( @@ -363,8 +351,9 @@ class NanoBananaPro: pbar=None, enable_grounding: bool = False, enable_image_search: bool = False, + save_to_disk: bool = True, ) -> List[dict]: - """异步批量处理：每个提示词独立调用 API，生成后立即写磁盘""" + """异步批量处理：每个提示词独立调用 API""" # 构建任务列表：(prompt, sub_index) 用于 images_per_prompt > 1 的情况 tasks_def = [] for p_idx, prompt in enumerate(prompts): @@ -373,9 +362,8 @@ class NanoBananaPro: total_tasks = len(tasks_def) num_prompts = len(prompts) - print(f"Nano Banana Pro: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务") - - max_concurrent = 10 + + max_concurrent = 50 num_batches = math.ceil(total_tasks / max_concurrent) all_results = [] @@ -383,7 +371,7 @@ class NanoBananaPro: success_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: for batch_idx in range(num_batches): @@ -405,6 +393,7 @@ class NanoBananaPro: global_task_index=i, enable_grounding=enable_grounding, enable_image_search=enable_image_search, + save_to_disk=save_to_disk, ) ) tasks.append(task) @@ -454,23 +443,18 @@ class NanoBananaPro: 宽高比: str, 分辨率: str, 生图数量: int, - 像素缩放: bool, - 分辨率像素: float, seed: int, - 跳过错误: bool = False, **kwargs ) -> Tuple[torch.Tensor]: """ 生成图像 - + Args: prompt: 提示词 模型: 模型名称 宽高比: 宽高比 分辨率: 分辨率 生图数量: 批次大小 - 像素缩放: 是否启用像素缩放 - 分辨率像素: 目标像素数（百万像素） seed: 随机种子 **kwargs: 搜索开关（谷歌搜索（联网）/ 图片搜索（联网））及动态参考图输入 (参考图1-9) 注：两个搜索参数名含全角括号，不能作为 Python 形参，从 kwargs 中提取 @@ -550,15 +534,7 @@ class NanoBananaPro: 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(prompt) @@ -602,14 +578,13 @@ class NanoBananaPro: nonlocal success_count, fail_count if success: success_count += 1 - print(f"Nano Banana Pro: 任务 {current}/{total} 成功 ✓") else: fail_count += 1 - + # 更新 ComfyUI 原生进度条 if pbar is not None: pbar.update(1) - + # 内存监控（每完成10个任务检查一次） if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0: import gc @@ -627,21 +602,10 @@ class NanoBananaPro: 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"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(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) @@ -654,10 +618,11 @@ class NanoBananaPro: aspect_ratio=宽高比, images_per_prompt=生图数量, input_images=input_images, - output_folder=output_folder, + output_folder="", pbar=pbar, enable_grounding=enable_grounding, enable_image_search=enable_image_search, + save_to_disk=False, ) ) finally: @@ -666,23 +631,20 @@ class NanoBananaPro: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_async_in_thread) try: - results = future.result(timeout=3600) + results = future.result(timeout=900) except TimeoutError: - raise RuntimeError("任务执行超时（1小时），请减少提示词数量或检查网络连接") - + raise RuntimeError("任务执行超时（900秒），请减少提示词数量或检查网络连接") + # 统计结果 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: @@ -691,24 +653,17 @@ class NanoBananaPro: 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"Nano Banana Pro: 无法加载 {file_path} - {e}") - + for r in results: + output_images.extend(r.get("output_images", [])) + if not output_images: placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) output_images = [placeholder] output_tensor = _images_to_tensor_safe(output_images, _NODE) - print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") import gc gc.collect() @@ -716,7 +671,7 @@ class NanoBananaPro: else: # 单提示词模式 if 生图数量 == 1: - # 单张：同步生成 + 保存到磁盘 + 输出 tensor + # 单张：同步生成，输出 tensor generated_images = self.client.generate_sync( prompt=prompt, model=模型, @@ -730,35 +685,12 @@ class NanoBananaPro: enable_grounding=enable_grounding, 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: - # 多张：异步并发 + 磁盘保存（与批量提示词逻辑一致） - print(f"Nano Banana Pro: 单提示词×{生图数量}张 → 异步并发模式") + # 多张：异步并发，内存输出 if pbar is not None: 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(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) @@ -771,10 +703,11 @@ class NanoBananaPro: aspect_ratio=宽高比, images_per_prompt=生图数量, input_images=input_images, - output_folder=output_folder, + output_folder="", pbar=pbar, enable_grounding=enable_grounding, enable_image_search=enable_image_search, + save_to_disk=False, ) ) finally: @@ -783,16 +716,13 @@ class NanoBananaPro: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_async_in_thread) try: - results = future.result(timeout=3600) + results = future.result(timeout=900) except TimeoutError: - raise RuntimeError("任务执行超时（1小时），请减少生图数量或检查网络连接") + raise RuntimeError("任务执行超时（900秒），请减少生图数量或检查网络连接") 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" @@ -806,24 +736,16 @@ class NanoBananaPro: error_msg = fr.get("error", "未知错误") print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 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"Nano Banana Pro: 无法加载 {file_path} - {e}") + for r in results: + output_images.extend(r.get("output_images", [])) if not output_images: placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) output_images = [placeholder] 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 gc.collect() @@ -851,9 +773,9 @@ class NanoBananaPro: # 打印最终汇总 if fail_count > 0: - print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张") + print(f"完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张") else: - print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张") + print(f"完成！总耗时 {time_str} | 成功 {len(generated_images)}张") # 最终内存清理 import gc @@ -861,7 +783,7 @@ class NanoBananaPro: if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: final_memory = process.memory_info().rss / 1024 / 1024 print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB") - + return (output_tensor,) except ValueError as e: @@ -869,24 +791,12 @@ class NanoBananaPro: if str(e) == "未授权！": print("请联系作者授权后方可使用！") 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 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 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 finally: @@ -898,8 +808,7 @@ class NanoBananaPro: print(f"Nano Banana Pro: {balance_info}") except Exception: pass - + # 最终内存清理 import gc - gc.collect() - print(f"Nano Banana Pro: 最终内存清理完成") \ No newline at end of file + gc.collect() \ No newline at end of file diff --git a/nodes/nano_banana_v2.py b/nodes/nano_banana_v2.py deleted file mode 100644 index 0b653db..0000000 --- a/nodes/nano_banana_v2.py +++ /dev/null @@ -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() diff --git a/nodes/quan_neng_sheng_tu.py b/nodes/quan_neng_sheng_tu.py deleted file mode 100644 index 8d91cb6..0000000 --- a/nodes/quan_neng_sheng_tu.py +++ /dev/null @@ -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"全能生图: 最终内存清理完成") diff --git a/nodes/seedance_video.py b/nodes/seedance_video.py index 576cfe7..b2eb7b4 100644 --- a/nodes/seedance_video.py +++ b/nodes/seedance_video.py @@ -73,6 +73,79 @@ def _tensor_to_base64_url(tensor) -> str: return f"data:image/png;base64,{b64}" +def _video_to_base64_url(video) -> str: + """ComfyUI VIDEO 对象 → data:video/;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(" torch.Tensor: """从 URL 下载图片并转为 ComfyUI IMAGE tensor，失败时返回 None""" try: @@ -286,12 +359,201 @@ class Seedance: _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 = { - "Seedance": Seedance, + "Seedance": Seedance, + "SeedanceMultiModal": SeedanceMultiModal, } NODE_DISPLAY_NAME_MAPPINGS = { - "Seedance": "Seedance 视频生成", + "Seedance": "Seedance 视频生成", + "SeedanceMultiModal": "Seedance 多模态参考生视频", } diff --git a/nodes/stream_preview.py b/nodes/stream_preview.py new file mode 100644 index 0000000..f899b2f --- /dev/null +++ b/nodes/stream_preview.py @@ -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": (文本,)} diff --git a/nodes/universal_llm.py b/nodes/universal_llm.py index 0f445cb..2f03be3 100644 --- a/nodes/universal_llm.py +++ b/nodes/universal_llm.py @@ -6,17 +6,19 @@ ComfyUI 自定义节点，通过 OpenAI 兼容协议调用市面上主流的 AI API 密钥和地址通过插件统一配置（环境变量或 .config 文件），与 Google Gemini 节点一致 """ +import os import time import base64 import json from io import BytesIO -from typing import Optional, Tuple +from typing import Optional, Tuple, List import torch from PIL import Image from ..utils.image_utils import tensor_to_pil 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": { "图片": ("IMAGE",), - } + "视频": ("VIDEO",), + "文件": ("FILE_LIST",), + }, + "hidden": { + "node_id": "UNIQUE_ID", + }, } RETURN_TYPES = ("STRING",) @@ -113,12 +120,90 @@ class UniversalLLMChat: b64 = base64.b64encode(data).decode('utf-8') 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, prompt: str, images: Optional[torch.Tensor] = None, + file_paths: str = "", + file_list: Optional[FileList] = None, + video=None, ) -> list: - """构建 OpenAI 格式的 messages 数组""" + """构建 chat/completions 格式的 messages 数组""" image_data_urls = [] pil_images_cache = [] # 保留 PIL Image 用于总体积重新编码 @@ -177,49 +262,152 @@ class UniversalLLMChat: 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 以内") - 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}] content_parts = [] + + # 图片 for url in image_data_urls: content_parts.append({ "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({ "type": "text", - "text": prompt + "text": prompt, }) 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( self, 模型: str, 提示词: str, 图片: Optional[torch.Tensor] = None, + 视频=None, + 文件: Optional[FileList] = None, + node_id: str = "", ) -> Tuple[str]: start_time = time.time() try: self._ensure_config() - # 构建 messages - messages = self._build_messages(提示词, 图片) + # 构建 input + input_data = self._build_input(提示词, 图片, "", 文件, 视频) 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: 输入 = {input_desc}") - # 构建请求体 + # 构建请求体（chat/completions 格式） request_body = { "model": 模型, - "messages": messages, - "stream": False, + "messages": input_data, + "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 事件循环冲突） import aiohttp import asyncio @@ -236,9 +424,9 @@ class UniversalLLMChat: async with aiohttp.ClientSession(timeout=timeout) as session: async with session.post(url, headers=headers, json=request_body) as resp: status = resp.status - body = await resp.text() if status != 200: + body = await resp.text() try: err_data = json.loads(body) err_msg = err_data.get("error", {}).get("message", body[:200]) @@ -256,7 +444,30 @@ class UniversalLLMChat: else: 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(): loop = asyncio.new_event_loop() @@ -266,24 +477,10 @@ class UniversalLLMChat: loop.close() with ThreadPoolExecutor(max_workers=1) as pool: - response_data = 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) + reply = pool.submit(_run_in_thread).result() elapsed = time.time() - start_time print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)") - print(f"全能LLM: Token 用量 — 输入: {prompt_tokens}, 输出: {completion_tokens}, 合计: {total_tokens}") if reply: preview = reply[:100] + "..." if len(reply) > 100 else reply print(f"全能LLM: 回复预览: {preview}") diff --git a/nodes/video_preview.py b/nodes/video_preview.py index d660598..4280246 100644 --- a/nodes/video_preview.py +++ b/nodes/video_preview.py @@ -1,9 +1,10 @@ """ -通用视频预览节点 -ComfyUI 自定义节点，接收视频文件路径并在前端展示预览 +视频预览节点 +接收 VIDEO 类型，在前端内嵌播放器预览 """ import os +import io try: import folder_paths @@ -11,8 +12,6 @@ try: except ImportError: FOLDER_PATHS_AVAILABLE = False -SUPPORTED_VIDEO_EXTENSIONS = {".mp4", ".webm", ".mov", ".avi", ".mkv", ".flv", ".wmv", ".3gp"} - def _get_output_dir() -> str: if FOLDER_PATHS_AVAILABLE: @@ -22,86 +21,70 @@ def _get_output_dir() -> str: class VideoPreview: - """ - 通用视频预览节点 - - 功能： - - 接收视频文件路径（STRING） - - 在 ComfyUI 前端节点上内嵌