"""
全能生图 节点
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"全能生图: 最终内存清理完成")