import os import io import base64 import locale import json import asyncio import aiohttp import numpy import PIL.Image import torch from openai import AsyncOpenAI from volcenginesdkarkruntime import Ark GLOBAL_CATEGORY = "JimengAI" DEFAULT_DOWNLOAD_TIMEOUT = 60 # 下载超时时间（秒） DEFAULT_DOWNLOAD_RETRIES = 3 # 下载重试次数 jimeng_api_dir = os.path.dirname(os.path.abspath(__file__)) API_KEYS_FILE = os.path.join(jimeng_api_dir, "api_keys.json") API_KEYS_CONFIG = [] CURRENT_LANG = "en" LOG_TRANSLATIONS = { "zh": { "api_file_not_found": "[JimengAI] 提示: 未找到 API 密钥文件。请将 'api_keys.json.example' 重命名为 'api_keys.json' 并填入您的密钥。", "api_file_empty": "[JimengAI] 警告: 'api_keys.json' 为空或格式不正确。", "api_load_error": "[JimengAI] 错误: 加载 'api_keys.json' 失败: {e}", "api_key_not_found": "[JimengAI] 错误: 未找到 '{key_name}' 的 API Key，请检查配置文件。", "est_fallback": "默认兜底", "est_history": "历史均值", "est_regression": "线性回归", "est_recent": "近期负载调整", "task_submitted_est": "[JimengAI] 任务已提交。预估生成时间: {time}s (预估方式: {method})", "task_info_simple": "[JimengAI] 任务 ID: {task_id} | 模型: {model}", "batch_submit_start": "[JimengAI] 正在提交 {count} 个任务 (模型: {model})...", "batch_submit_result": "[JimengAI] 提交完成。成功: {created}，失败: {failed}。", "polling_single": "[JimengAI] 任务 {task_id}: 已耗时 {elapsed}s / 预估 {max}s ...", "polling_batch": "[JimengAI] 批量进度: 完成 {done}/{total}。等待中: {pending}... (耗时 {elapsed}s / {max}s)", "interrupted": "\n[JimengAI] 用户中断了处理。正在尝试取消挂起任务...", "task_finished_single": "[JimengAI] 任务执行成功。", "batch_finished_stats": "[JimengAI] 批量处理完成。成功: {success}，失败: {failed}。", "batch_handling": "[JimengAI] 正在处理 {count} 个成功任务，按 Seed 排序并下载...", "batch_copying": "[JimengAI] 正在保存文件到输出目录: {path}", "err_download_url": "[JimengAI] 异步下载或转换失败，URL: {url}，错误: {e}", "err_task_create": "[JimengAI] 任务创建失败: {e}", "err_task_check": "[JimengAI] 检查任务状态失败 ({tid}): {e}", "err_task_fail_msg": "[JimengAI] 任务 {tid} 失败: {msg}", "err_batch_fail_all": "[JimengAI] 批量任务失败: 无任务成功。", "err_copy_fail": "[JimengAI] 复制文件失败: {path}，错误: {e}", "err_convert_tensor": "[JimengAI] 转换图片 Tensor 失败: {e}", "err_check_status_batch": "[JimengAI] 检查非阻塞任务状态失败: {e}", "download_retry": "[JimengAI] 警告: 下载失败 (尝试 {attempt}/{total})。{delay}秒后重试... 错误: {e}", # 前端弹窗提示 "popup_req_failed": "[JimengAI] 请求失败: {msg}", "popup_task_failed": "[JimengAI] 任务 {task_id} 失败: {msg}", "popup_batch_pending": "[JimengAI] 批量任务 ({count} 个) 处理中。请再次运行以检查结果。", "popup_task_pending": "[JimengAI] 任务 {task_id} 状态为 {status}。请再次运行以检查结果。", "popup_param_not_allowed": "[JimengAI] 参数错误: 提示词中不允许包含参数 '--{param}'。请使用节点的组件进行设置。", "popup_first_frame_missing": "[JimengAI] 参数错误: 使用尾帧图片时必须提供首帧图片。", "popup_ref_missing": "[JimengAI] 参数错误: 必须提供至少一张参考图。", "popup_prepare_failed": "[JimengAI] 任务准备失败: {e}", # 图片节点提示 "err_pixels_range": "[JimengAI] 参数错误: 总像素数必须在 {min} ({min_desc}) 和 {max} ({max_desc}) 之间。当前值: {current}", "err_aspect_ratio": "[JimengAI] 参数错误: 宽高比必须在 {min} 和 {max} 之间。当前值: {current}", "err_download_img": "[JimengAI] 错误: 下载生成的图像失败。", "err_gen_model": "[JimengAI] 模型 {model} 生成失败: {e}", "err_img_limit_10": "[JimengAI] 参数错误: 输入图像数量不能超过 10 张。", "err_img_limit_15": "[JimengAI] 参数错误: 输入图像数 ({n}) 与最大生成数 ({max}) 之和不能超过 15。", # API Key 相关 "combo_no_keys": "未找到密钥 (请配置 api_keys.json)", "popup_key_valid_err": "[JimengAI] 配置错误: 所选密钥 '{key}' 无效或未找到。请检查 api_keys.json。", }, "en": { "api_file_not_found": "[JimengAI] Info: API keys file not found. Please rename 'api_keys.json.example' to 'api_keys.json' and fill in your keys.", "api_file_empty": "[JimengAI] Warning: 'api_keys.json' is empty or not formatted correctly.", "api_load_error": "[JimengAI] Error: Failed to load 'api_keys.json': {e}", "api_key_not_found": "[JimengAI] Error: API Key for '{key_name}' not found.", "est_fallback": "Fallback Default", "est_history": "History Average", "est_regression": "Linear Regression", "est_recent": "Recent Load Adjustment", "task_submitted_est": "[JimengAI] Task submitted. Est. time: {time}s (Method: {method})", "task_info_simple": "[JimengAI] Task ID: {task_id} | Model: {model}", "batch_submit_start": "[JimengAI] Submitting batch of {count} tasks (Model: {model})...", "batch_submit_result": "[JimengAI] Submission complete. Created: {created}, Failed: {failed}.", "polling_single": "[JimengAI] Task {task_id}: {elapsed}s / {max}s elapsed...", "polling_batch": "[JimengAI] Batch Progress: {done}/{total} done. {pending} pending... ({elapsed}s / {max}s Est.)", "interrupted": "\n[JimengAI] Processing interrupted by user. Cancelling pending tasks...", "task_finished_single": "[JimengAI] Task completed successfully.", "batch_finished_stats": "[JimengAI] Batch finished. Success: {success}, Failed: {failed}.", "batch_handling": "[JimengAI] Handling {count} successful tasks. Sorting by seed and downloading...", "batch_copying": "[JimengAI] Copying files to output directory: {path}", "err_download_url": "[JimengAI] Async download failed, URL: {url}, Error: {e}", "err_task_create": "[JimengAI] Task creation failed: {e}", "err_task_check": "[JimengAI] Failed to check status for {tid}: {e}", "err_task_fail_msg": "[JimengAI] Task {tid} failed: {msg}", "err_batch_fail_all": "[JimengAI] Batch failed: No tasks succeeded.", "err_copy_fail": "[JimengAI] Failed to copy file: {path}. Error: {e}", "err_convert_tensor": "[JimengAI] Failed to convert frame to tensor: {e}", "err_check_status_batch": "[JimengAI] API Error checking batch status: {e}", "download_retry": "[JimengAI] Warning: Download failed (Attempt {attempt}/{total}). Retrying in {delay}s... Error: {e}", # Popup Messages "popup_req_failed": "[JimengAI] Request failed: {msg}", "popup_task_failed": "[JimengAI] Task {task_id} failed: {msg}", "popup_batch_pending": "[JimengAI] Batch ({count} tasks) is pending. Run again to check results.", "popup_task_pending": "[JimengAI] Task {task_id} is {status}. Run again to check results.", "popup_param_not_allowed": "[JimengAI] Parameter Error: Parameter '--{param}' is not allowed in the prompt. Please use the node's widget for this value.", "popup_first_frame_missing": "[JimengAI] Parameter Error: A first frame image must be provided when using a last frame image.", "popup_ref_missing": "[JimengAI] Parameter Error: At least one reference image must be provided.", "popup_prepare_failed": "[JimengAI] Failed to prepare task: {e}", # Image Nodes "err_pixels_range": "[JimengAI] Parameter Error: Total pixels must be between {min} ({min_desc}) and {max} ({max_desc}). Your current: {current}", "err_aspect_ratio": "[JimengAI] Parameter Error: Aspect ratio must be between {min} and {max}. Your current: {current}", "err_download_img": "[JimengAI] Error: Failed to download the generated image.", "err_gen_model": "[JimengAI] Failed to generate image with model {model}: {e}", "err_img_limit_10": "[JimengAI] Parameter Error: The number of input images cannot exceed 10.", "err_img_limit_15": "[JimengAI] Parameter Error: The sum of input images ({n}) and max generated images ({max}) cannot exceed 15.", # API Key "combo_no_keys": "No Keys Found (Check api_keys.json)", "popup_key_valid_err": "[JimengAI] Config Error: Selected key '{key}' is invalid or not found. Please check api_keys.json.", } } def detect_system_language(): try: lang_code, _ = locale.getdefaultlocale() if lang_code and lang_code.startswith('zh'): return 'zh' except: pass return 'en' def get_text(key): global CURRENT_LANG mapping = LOG_TRANSLATIONS.get(CURRENT_LANG, LOG_TRANSLATIONS["en"]) return mapping.get(key, LOG_TRANSLATIONS["en"].get(key, key)) def log_msg(key, default_msg="", **kwargs): global CURRENT_LANG mapping = LOG_TRANSLATIONS.get(CURRENT_LANG, LOG_TRANSLATIONS["en"]) msg = mapping.get(key, None) if msg is None: msg = LOG_TRANSLATIONS["en"].get(key, default_msg) if msg: try: print(msg.format(**kwargs)) except Exception: print(msg) def load_api_keys(): global API_KEYS_CONFIG, CURRENT_LANG API_KEYS_CONFIG = [] CURRENT_LANG = detect_system_language() if not os.path.exists(API_KEYS_FILE): log_msg("api_file_not_found") return try: with open(API_KEYS_FILE, 'r', encoding='utf-8') as f: keys_data = json.load(f) if isinstance(keys_data, list): for item in keys_data: if "customName" in item and "apiKey" in item: API_KEYS_CONFIG.append(item) if not API_KEYS_CONFIG: log_msg("api_file_empty") except Exception as e: log_msg("api_load_error", e=e) async def _fetch_data_from_url_async(session: aiohttp.ClientSession, url: str, timeout: int = DEFAULT_DOWNLOAD_TIMEOUT, retries: int = DEFAULT_DOWNLOAD_RETRIES) -> bytes: for attempt in range(1, retries + 2): try: client_timeout = aiohttp.ClientTimeout(total=timeout) async with session.get(url, timeout=client_timeout) as response: response.raise_for_status() return await response.read() except (aiohttp.ClientError, asyncio.TimeoutError) as e: if attempt > retries: raise e retry_delay = 2 log_msg("download_retry", attempt=attempt, total=retries+1, delay=retry_delay, e=e) await asyncio.sleep(retry_delay) return b"" def _tensor2images(tensor: torch.Tensor) -> list: np_imgs = numpy.clip(tensor.cpu().numpy() * 255.0, 0, 255.0).astype(numpy.uint8) return [PIL.Image.fromarray(np_img) for np_img in np_imgs] def _image_to_base64(image: torch.Tensor) -> str: if image is None: return None with io.BytesIO() as bytes_io: _tensor2images(image)[0].save(bytes_io, format='JPEG') data_bytes = bytes_io.getvalue() return base64.b64encode(data_bytes).decode("utf-8") async def _download_url_to_image_tensor_async(session: aiohttp.ClientSession, url: str) -> torch.Tensor | None: if not url: return None try: image_data = await _fetch_data_from_url_async(session, url) i = PIL.Image.open(io.BytesIO(image_data)) image = i.convert("RGB") image = numpy.array(image).astype(numpy.float32) / 255.0 return torch.from_numpy(image)[None,] except Exception as e: log_msg("err_download_url", url=url, e=e) return None class JimengClients: def __init__(self, openai_client, ark_client): self.openai = openai_client self.ark = ark_client class JimengAPIClient: # 即梦 API 客户端节点 @classmethod def INPUT_TYPES(s): load_api_keys() key_names = [key["customName"] for key in API_KEYS_CONFIG] if not key_names: # 汉化：列表为空时的占位符 key_names = [get_text("combo_no_keys")] return { "required": { "key_name": (key_names,), } } RETURN_TYPES = ("JIMENG_CLIENT",) RETURN_NAMES = ("client",) FUNCTION = "create_clients" CATEGORY = GLOBAL_CATEGORY def create_clients(self, key_name): api_key = None # 检查是否选择了占位符 if key_name == get_text("combo_no_keys"): # 抛出中文弹窗，提示去配置 Keys raise RuntimeError(get_text("api_file_not_found")) for key_info in API_KEYS_CONFIG: if key_info["customName"] == key_name: api_key = key_info["apiKey"] break if not api_key: log_msg("api_key_not_found", key_name=key_name) # 抛出中文弹窗，提示 Key 无效 raise RuntimeError(get_text("popup_key_valid_err").format(key=key_name)) openai_client = AsyncOpenAI(api_key=api_key, base_url="https://ark.cn-beijing.volces.com/api/v3") ark_client = Ark(api_key=api_key) clients = JimengClients(openai_client, ark_client) return (clients,) NODE_CLASS_MAPPINGS = { "JimengAPIClient": JimengAPIClient, } NODE_DISPLAY_NAME_MAPPINGS = { "JimengAPIClient": "Jimeng API Client", }