diff --git a/__init__.py b/__init__.py index 80bb9b4..2a2aa73 100644 --- a/__init__.py +++ b/__init__.py @@ -12,7 +12,7 @@ Comfyui_o1key - ComfyUI 自定义节点集合 import ssl -from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, KVideo +from .nodes import NanoBananaPro, NanoBananaProAsync, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, KVideo from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck # 报错弹框友好文案（不修改原节点代码，仅在外层统一处理） @@ -46,11 +46,13 @@ def _wrap_generate_for_error_display(cls, attr="generate"): setattr(cls, attr, wrapped) _wrap_generate_for_error_display(NanoBananaPro) +_wrap_generate_for_error_display(NanoBananaProAsync) _wrap_generate_for_error_display(BatchNanoBananaPro) # ComfyUI 节点注册 NODE_CLASS_MAPPINGS = { "NanoBananaPro": NanoBananaPro, + "NanoBananaProAsync": NanoBananaProAsync, "BatchNanoBananaPro": BatchNanoBananaPro, "GoogleGemini": GoogleGemini, "LoadFile": LoadFile, @@ -81,6 +83,7 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "NanoBananaPro": "Nano Banana", + "NanoBananaProAsync": "Nano Banana（异步）", "BatchNanoBananaPro": "批量 Nano Banana", "GoogleGemini": "Google Gemini", "LoadFile": "加载文件", diff --git a/clients/gemini_client.py b/clients/gemini_client.py index d703d38..fa7d8af 100644 --- a/clients/gemini_client.py +++ b/clients/gemini_client.py @@ -230,7 +230,7 @@ class GeminiAPIClient(BaseAPIClient): parts + img_parts, { "generationConfig": { - "responseModalities": ["TEXT", "IMAGE"], + "responseModalities": ["IMAGE"], "imageConfig": { "aspectRatio": aspect_ratio, "imageSize": resolution @@ -274,7 +274,7 @@ class GeminiAPIClient(BaseAPIClient): } ], "generationConfig": { - "responseModalities": ["TEXT", "IMAGE"], + "responseModalities": ["IMAGE"], "imageConfig": { "aspectRatio": aspect_ratio, "imageSize": resolution diff --git a/nodes/__init__.py b/nodes/__init__.py index cdf1317..248d9ff 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -5,6 +5,7 @@ from .stream_preview import StreamPreview from .nano_banana_pro import NanoBananaPro +from .nano_banana_pro_async import NanoBananaProAsync from .batch_nano_banana_pro import BatchNanoBananaPro from .google_gemini import GoogleGemini from .load_file import LoadFile @@ -25,4 +26,4 @@ from .K3_video import K3Video from .K3_video_firstlast import K3VideoFirstLast from .K3_motion_control import K3MotionControl, K3MotionVideoCheck -__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'KVideo', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck'] +__all__ = ['NanoBananaPro', 'NanoBananaProAsync', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'KVideo', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck'] diff --git a/nodes/nano_banana_pro.py b/nodes/nano_banana_pro.py index cdb4a98..5a2717e 100644 --- a/nodes/nano_banana_pro.py +++ b/nodes/nano_banana_pro.py @@ -52,10 +52,10 @@ except ImportError: # ============================================================================ # 是否启用调试日志（打印完整的 API 响应内容） # 设置为 True 以启用调试日志，False 以禁用 -DEBUG_LOG_ENABLED = False +DEBUG_LOG_ENABLED = True # 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断） # 设置为 True 以启用请求体日志，False 以禁用 -REQUEST_LOG_ENABLED = False +REQUEST_LOG_ENABLED = True # ============================================================================ _NODE = "Nano Banana Pro" diff --git a/nodes/nano_banana_pro_async.py b/nodes/nano_banana_pro_async.py new file mode 100644 index 0000000..6fb1db4 --- /dev/null +++ b/nodes/nano_banana_pro_async.py @@ -0,0 +1,669 @@ +""" +Nano Banana Pro（异步）节点 +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 ..utils.config import get_async_api_base_url +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 + print("⚠️ NanoBananaProAsync: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + print("⚠️ NanoBananaProAsync: psutil 不可用，内存监控功能禁用") + +DEBUG_LOG_ENABLED = True +REQUEST_LOG_ENABLED = True + +_NODE = "Nano Banana Pro（异步）" +_POLL_INTERVAL = 4 # 轮询间隔（秒） +_MAX_WAIT_TIME = 300 # 最大等待时间（秒） + + +def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: + if not images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + return pil_to_tensor([placeholder]) + + base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size + 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}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str})，" + f"仅输出最大尺寸 {base_size[0]}×{base_size[1]} 的 {len(matched)} 张" + ) + + return pil_to_tensor(matched) + + +class NanoBananaProAsync: + """ + Nano Banana Pro（异步）节点 + + 功能： + - 异步提交任务到 cf-api.o1key.com + - 轮询任务状态直到完成 + - 支持批量并发生成 + """ + + 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 = ["512px", "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() + 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["代理端口"] = ("STRING", { + "default": "", + "multiline": False, + "placeholder": "本地代理端口，如 7897（Clash Verge）或 10808（v2rayN），留空不使用" + }) + + 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 + }), + "谷歌搜索（联网）": (["关闭", "打开"], { + "default": "关闭" + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }) + }, + "optional": optional_inputs + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + + try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True + except ImportError: + FOLDER_PATHS_AVAILABLE = False + + FUNCTION = "generate" + CATEGORY = "image/generation" + + async def _submit_task_async( + self, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[Image.Image], + enable_grounding: bool = False, + ) -> str: + """提交异步任务，返回 task_id""" + endpoint = self.client.get_endpoint(model=model, resolution=resolution, image_format="url") + async_endpoint = f"/async{endpoint.split('?')[0]}" + if "?" in endpoint: + async_endpoint += "?" + endpoint.split("?")[1] + + request_body = self.client.build_request_body( + prompt=prompt, + images=images if images else None, + aspect_ratio=aspect_ratio, + resolution=resolution, + enable_grounding=enable_grounding, + enable_image_search=False, + ) + + url = f"{get_async_api_base_url()}{async_endpoint}" + headers = { + "Authorization": f"Bearer {self.client.api_key}", + "Content-Type": "application/json" + } + + if REQUEST_LOG_ENABLED: + import json + import copy + debug_body = copy.deepcopy(request_body) + for content in debug_body.get("contents", []): + for part in content.get("parts", []): + if "inline_data" in part and "data" in part["inline_data"]: + data_str = part["inline_data"]["data"] + part["inline_data"]["data"] = f"{data_str[:50]}...[截断]" if len(data_str) > 50 else data_str + print(f"\n{'='*60}") + print(f"[异步提交] URL: {url}") + print(f"[异步提交] 请求体:\n{json.dumps(debug_body, indent=2, ensure_ascii=False)}") + print(f"{'='*60}\n") + + async with session.post(url, json=request_body, headers=headers, proxy=self.client.proxy_url) as response: + if response.status != 200: + error_text = await response.text() + raise RuntimeError(f"提交任务失败 ({response.status}): {error_text}") + + data = await response.json() + + if DEBUG_LOG_ENABLED: + import json + print(f"\n{'='*60}") + print(f"[异步提交] 响应:\n{json.dumps(data, indent=2, ensure_ascii=False)}") + print(f"{'='*60}\n") + + task_id = data.get("task_id") + if not task_id: + raise RuntimeError(f"提交响应中未找到 task_id: {data}") + + return task_id + + async def _poll_task_async( + self, + session: aiohttp.ClientSession, + task_id: str, + ) -> dict: + """轮询任务状态直到完成""" + url = f"{get_async_api_base_url()}/async/v1/tasks/{task_id}" + headers = { + "Authorization": f"Bearer {self.client.api_key}", + "Content-Type": "application/json" + } + + start_time = time.time() + poll_count = 0 + + while True: + elapsed = time.time() - start_time + if elapsed > _MAX_WAIT_TIME: + raise RuntimeError(f"任务 {task_id} 超时（{_MAX_WAIT_TIME}秒），请稍后手动查询") + + poll_count += 1 + + async with session.get(url, headers=headers, proxy=self.client.proxy_url) as response: + if response.status != 200: + error_text = await response.text() + raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}") + + result = await response.json() + status = result.get("status") + + if DEBUG_LOG_ENABLED: + import json + print(f"\n{'='*60}") + print(f"[轮询 #{poll_count}] task_id: {task_id}") + print(f"[轮询 #{poll_count}] 响应:\n{json.dumps(result, indent=2, ensure_ascii=False)}") + print(f"{'='*60}\n") + + if status == "SUCCESS": + return result.get("data", {}) + elif status == "FAILURE": + error_msg = result.get("error", "未知错误") + raise RuntimeError(f"任务失败: {error_msg}") + elif status in ["SUBMITTED", "IN_PROGRESS"]: + await asyncio.sleep(_POLL_INTERVAL) + else: + raise RuntimeError(f"未知任务状态: {status}") + + 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, + 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 + } + + try: + task_id = await self._submit_task_async( + session=session, + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images, + enable_grounding=enable_grounding, + ) + + response_data = await self._poll_task_async(session=session, task_id=task_id) + + images_list, _ = await self.client.parse_response_async(response_data, session=session) + + 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( + 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, + save_to_disk: bool = True, + ) -> 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) + max_concurrent = 50 + num_batches = math.ceil(total_tasks / max_concurrent) + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 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): + 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, + save_to_disk=save_to_disk, + ) + ) + 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) + + import gc + gc.collect() + + await asyncio.sleep(0.1) + + return all_results + + def generate( + self, + prompt: str, + 模型: str, + 宽高比: str, + 分辨率: str, + 生图数量: int, + seed: int, + **kwargs + ) -> Tuple[torch.Tensor]: + """生成图像（异步模式）""" + start_time = time.time() + + enable_grounding: bool = (kwargs.pop("谷歌搜索（联网）", "关闭") == "打开") + proxy_port: str = kwargs.pop("代理端口", "") + + 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"{_NODE}: 初始内存使用: {initial_memory:.1f} MB") + + if self.client is None: + try: + self.client = GeminiAPIClient() + except ValueError as e: + raise ValueError(f"初始化失败: {str(e)}") + + self.client.proxy_url = GeminiAPIClient.build_proxy_url(proxy_port) + if self.client.proxy_url: + print(f"{_NODE}: 已启用代理加速 → {self.client.proxy_url}") + + 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): + 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 张，请减少输入图像数量" + ) + + batch_prompts = parse_batch_prompts(prompt) + + grounding_str = "" + if 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 + + if batch_prompts: + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + + if pbar is not None: + pbar = ProgressBar(total_images) + + 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="", + pbar=pbar, + enable_grounding=enable_grounding, + save_to_disk=False, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=900) + except TimeoutError: + 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) + + 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}") + + output_images = [] + 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) + + import gc + gc.collect() + return (output_tensor,) + else: + if pbar is not None: + pbar = ProgressBar(生图数量) + + 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=[prompt], + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder="", + pbar=pbar, + enable_grounding=enable_grounding, + save_to_disk=False, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=900) + except TimeoutError: + 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) + + 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}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}") + + output_images = [] + 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) + + import gc + gc.collect() + return (output_tensor,) + + except ValueError as e: + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + raise ValueError(str(e)) from None + + except RuntimeError as e: + raise RuntimeError(str(e)) from None + + except Exception as e: + raise type(e)(str(e)) 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 + + import gc + gc.collect() diff --git a/nodes/universal_llm.py b/nodes/universal_llm.py index bc32266..f38bd46 100644 --- a/nodes/universal_llm.py +++ b/nodes/universal_llm.py @@ -26,14 +26,11 @@ from ..utils.file_types import FileList SUPPORTED_MODELS = [ "gpt-5.4", - "gemini-3-flash-preview", "gemini-3.1-flash-lite-preview", "gemini-3.1-pro-preview", "deepseek-v3.2", - "kimi-k2.5", + "deepseek-v4-pro", "doubao-seed-2-0-pro-260215", - "qwen3.5-plus-2026-02-15", - "qwen3.5-plus", ] # 图片缩放最大尺寸 diff --git a/utils/config.py b/utils/config.py index 90bf003..c9c9e60 100644 --- a/utils/config.py +++ b/utils/config.py @@ -17,6 +17,10 @@ CONFIG_FILE = os.path.join(PLUGIN_ROOT, ".config") # 可通过环境变量 O1KEY_API_BASE_URL 覆盖 DEFAULT_API_BASE_URL = "https://api.o1key.com" +# 异步 API 基础 URL（用于异步提交+轮询模式） +# 可通过环境变量 O1KEY_ASYNC_API_BASE_URL 覆盖 +DEFAULT_ASYNC_API_BASE_URL = "https://cf-api.o1key.com" + def load_config(config_path: Optional[str] = None) -> Dict[str, str]: """ @@ -115,3 +119,20 @@ def get_api_base_url() -> str: return base_url.rstrip('/') return DEFAULT_API_BASE_URL + + +def get_async_api_base_url() -> str: + """ + 获取异步 API 基础 URL + 从 .config 文件读取，如果未配置则使用默认值 + + Returns: + 异步 API 基础 URL 字符串 + """ + config = load_config() + base_url = config.get("O1KEY_ASYNC_API_BASE_URL") + + if base_url: + return base_url.rstrip('/') + + return DEFAULT_ASYNC_API_BASE_URL