Initial commit: Comfyui_o1key v1.10.0

utils/__init__.py

@@ -0,0 +1,39 @@
+"""
+工具模块
+包含图像处理、配置管理、文件处理等通用工具函数
+"""
+
+from .image_utils import (
+    tensor_to_pil,
+    pil_to_tensor,
+    encode_image_to_base64,
+    decode_base64_to_pil
+)
+from .config import load_config, get_api_key
+from .file_utils import (
+    ImageInfo,
+    load_images_from_folder,
+    pair_images_indexed,
+    pair_images_cartesian,
+    generate_output_filename,
+    generate_batch_output_filenames,
+    save_image,
+    get_folder_image_count
+)
+
+__all__ = [
+    'tensor_to_pil',
+    'pil_to_tensor', 
+    'encode_image_to_base64',
+    'decode_base64_to_pil',
+    'load_config',
+    'get_api_key',
+    'ImageInfo',
+    'load_images_from_folder',
+    'pair_images_indexed',
+    'pair_images_cartesian',
+    'generate_output_filename',
+    'generate_batch_output_filenames',
+    'save_image',
+    'get_folder_image_count'
+]

utils/config.py

@@ -0,0 +1,114 @@
+"""
+配置管理模块
+处理环境变量和 API 密钥管理
+"""
+
+import os
+from typing import Dict, Optional
+
+
+# 获取插件根目录
+PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+CONFIG_FILE = os.path.join(PLUGIN_ROOT, ".config")
+
+
+def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
+    """
+    从配置文件加载所有配置项
+    
+    Args:
+        config_path: 配置文件路径，默认为插件目录下的 .config
+    
+    Returns:
+        配置字典 {key: value}
+    
+    Example:
+        >>> config = load_config()
+        >>> api_key = config.get('O1KEY_API_KEY')
+    """
+    if config_path is None:
+        config_path = CONFIG_FILE
+    
+    config = {}
+    
+    if not os.path.exists(config_path):
+        return config
+    
+    try:
+        with open(config_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                
+                # 跳过空行和注释
+                if not line or line.startswith('#'):
+                    continue
+                
+                # 解析 KEY=VALUE 格式
+                if '=' in line:
+                    key, value = line.split('=', 1)
+                    key = key.strip()
+                    value = value.strip().strip('"').strip("'")
+                    
+                    if key and value:
+                        config[key] = value
+    
+    except Exception as e:
+        print(f"⚠️ 读取配置文件失败: {e}")
+    
+    return config
+
+
+def get_api_key(key_name: str = "O1KEY_API_KEY") -> Optional[str]:
+    """
+    获取 API 密钥
+    优先级：环境变量（推荐） > .config 文件（向后兼容）
+    
+    Args:
+        key_name: 密钥名称，默认为 O1KEY_API_KEY
+    
+    Returns:
+        API 密钥字符串，如果未找到则返回 None
+    
+    Raises:
+        ValueError: 如果未找到 API 密钥
+    
+    Example:
+        >>> api_key = get_api_key()
+        >>> if api_key is None:
+        ...     raise ValueError("API key not found")
+    """
+    # 1. 优先从环境变量读取（推荐方式）
+    api_key = os.environ.get(key_name)
+    
+    if api_key:
+        return api_key
+    
+    # 2. 从 .config 文件读取（向后兼容，已弃用）
+    config = load_config()
+    api_key = config.get(key_name)
+    
+    if api_key:
+        return api_key
+    
+    return None
+
+
+def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
+    """
+    获取 API 密钥，如果未找到则抛出异常
+    
+    Args:
+        key_name: 密钥名称
+    
+    Returns:
+        API 密钥字符串
+    
+    Raises:
+        ValueError: 如果未找到 API 密钥
+    """
+    api_key = get_api_key(key_name)
+    
+    if not api_key:
+        raise ValueError("未授权！")
+    
+    return api_key

utils/file_utils.py

@@ -0,0 +1,328 @@
+"""
+文件处理工具模块
+提供文件夹图片加载、智能命名、图片配对等功能
+"""
+
+import os
+import uuid
+import time
+from itertools import product
+from pathlib import Path
+from typing import List, Tuple, Optional, NamedTuple
+
+from PIL import Image
+
+
+# 支持的图片格式
+SUPPORTED_IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
+
+
+class ImageInfo(NamedTuple):
+    """图片信息结构"""
+    image: Image.Image
+    filename: str  # 不含扩展名的文件名
+    extension: str  # 扩展名（如 .png）
+    source_path: str  # 原始文件路径
+
+
+def load_images_from_folder(
+    folder_path: str,
+    recursive: bool = False
+) -> List[ImageInfo]:
+    """
+    从文件夹加载所有图片
+    
+    Args:
+        folder_path: 文件夹路径
+        recursive: 是否递归加载子文件夹
+    
+    Returns:
+        ImageInfo 列表，包含图片和元数据
+    
+    Raises:
+        ValueError: 文件夹不存在或为空
+    
+    Example:
+        >>> images = load_images_from_folder("D:/images")
+        >>> for info in images:
+        ...     print(f"{info.filename}: {info.image.size}")
+    """
+    folder_path = folder_path.strip()
+    
+    if not folder_path:
+        return []
+    
+    path = Path(folder_path)
+    
+    if not path.exists():
+        raise ValueError(f"文件夹不存在: {folder_path}")
+    
+    if not path.is_dir():
+        raise ValueError(f"路径不是文件夹: {folder_path}")
+    
+    images = []
+    
+    # 获取文件列表
+    if recursive:
+        files = list(path.rglob("*"))
+    else:
+        files = list(path.iterdir())
+    
+    # 按文件名排序，确保顺序一致
+    files = sorted(files, key=lambda x: x.name.lower())
+    
+    for file_path in files:
+        if not file_path.is_file():
+            continue
+        
+        ext = file_path.suffix.lower()
+        if ext not in SUPPORTED_IMAGE_EXTENSIONS:
+            continue
+        
+        try:
+            img = Image.open(file_path)
+            img.load()  # 确保图片完全加载
+            
+            # 转换为 RGB 模式
+            if img.mode != 'RGB':
+                img = img.convert('RGB')
+            
+            images.append(ImageInfo(
+                image=img,
+                filename=file_path.stem,
+                extension=ext,
+                source_path=str(file_path)
+            ))
+        except Exception as e:
+            print(f"警告: 无法加载图片 {file_path}: {e}")
+            continue
+    
+    return images
+
+
+def pair_images_indexed(
+    *image_lists: List[ImageInfo]
+) -> List[Tuple[ImageInfo, ...]]:
+    """
+    1:1 索引配对
+    
+    按索引位置配对多个图片列表，以最短列表长度为准。
+    
+    Args:
+        *image_lists: 多个 ImageInfo 列表
+    
+    Returns:
+        配对后的元组列表
+    
+    Example:
+        >>> list_a = [a1, a2, a3]
+        >>> list_b = [b1, b2, b3]
+        >>> pairs = pair_images_indexed(list_a, list_b)
+        >>> # [(a1, b1), (a2, b2), (a3, b3)]
+    """
+    if not image_lists:
+        return []
+    
+    # 过滤空列表
+    non_empty_lists = [lst for lst in image_lists if lst]
+    
+    if not non_empty_lists:
+        return []
+    
+    # 使用 zip 进行索引配对（以最短列表为准）
+    return list(zip(*non_empty_lists))
+
+
+def pair_images_cartesian(
+    *image_lists: List[ImageInfo]
+) -> List[Tuple[ImageInfo, ...]]:
+    """
+    笛卡尔积配对
+    
+    生成多个图片列表的所有组合。
+    
+    Args:
+        *image_lists: 多个 ImageInfo 列表
+    
+    Returns:
+        配对后的元组列表
+    
+    Example:
+        >>> list_a = [a1, a2]
+        >>> list_b = [b1, b2]
+        >>> pairs = pair_images_cartesian(list_a, list_b)
+        >>> # [(a1, b1), (a1, b2), (a2, b1), (a2, b2)]
+    """
+    if not image_lists:
+        return []
+    
+    # 过滤空列表
+    non_empty_lists = [lst for lst in image_lists if lst]
+    
+    if not non_empty_lists:
+        return []
+    
+    # 使用 itertools.product 生成笛卡尔积
+    return list(product(*non_empty_lists))
+
+
+def generate_output_filename(
+    source_images: List[ImageInfo],
+    batch_index: int,
+    output_folder: str,
+    extension: str = ".png",
+    task_id: Optional[str] = None
+) -> str:
+    """
+    生成智能输出文件名
+    
+    基于源图片文件名生成输出文件名，使用任务ID和时间戳确保并发安全。
+    
+    Args:
+        source_images: 源图片信息列表
+        batch_index: 批次索引（从 0 开始）
+        output_folder: 输出文件夹路径
+        extension: 输出文件扩展名
+        task_id: 任务唯一标识符（用于并发场景）
+    
+    Returns:
+        完整的输出文件路径
+    
+    Example:
+        >>> # 单图片: hello.png -> hello_task0_12345_000.png
+        >>> # 多图片: hello.png + ref.png -> hello_ref_task0_12345_000.png
+        >>> # 并发安全：每个任务有唯一的 task_id 和时间戳
+    """
+    # 构建基础文件名
+    if len(source_images) == 1:
+        base_name = source_images[0].filename
+    else:
+        # 多个源图片，组合文件名
+        names = [info.filename for info in source_images]
+        base_name = "_".join(names)
+    
+    # 确保输出文件夹存在
+    output_path = Path(output_folder)
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    # 生成唯一性标识
+    if task_id is None:
+        # 如果没有提供 task_id，使用 UUID 前8位
+        task_id = str(uuid.uuid4())[:8]
+    
+    # 使用时间戳（毫秒级）增加唯一性
+    timestamp = int(time.time() * 1000) % 100000  # 精确到毫秒的后5位
+    
+    # 生成文件名：基础名_任务ID_时间戳_批次索引
+    filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}{extension}"
+    full_path = output_path / filename
+    
+    # 极小概率的冲突处理
+    counter = 1
+    while full_path.exists():
+        filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}_{counter}{extension}"
+        full_path = output_path / filename
+        counter += 1
+    
+    return str(full_path)
+
+
+def generate_batch_output_filenames(
+    source_images: List[ImageInfo],
+    count: int,
+    output_folder: str,
+    extension: str = ".png",
+    task_id: Optional[str] = None
+) -> List[str]:
+    """
+    批量生成输出文件名
+    
+    Args:
+        source_images: 源图片信息列表
+        count: 需要生成的文件名数量
+        output_folder: 输出文件夹路径
+        extension: 输出文件扩展名
+        task_id: 任务唯一标识符（用于并发场景）
+    
+    Returns:
+        输出文件路径列表
+    """
+    filenames = []
+    
+    for i in range(count):
+        filename = generate_output_filename(
+            source_images=source_images,
+            batch_index=i,
+            output_folder=output_folder,
+            extension=extension,
+            task_id=task_id
+        )
+        filenames.append(filename)
+    
+    return filenames
+
+
+def save_image(
+    image: Image.Image,
+    output_path: str,
+    quality: int = 95
+) -> str:
+    """
+    保存图片到指定路径
+    
+    Args:
+        image: PIL Image 对象
+        output_path: 输出文件路径
+        quality: JPEG 质量（仅对 JPEG 格式有效）
+    
+    Returns:
+        实际保存的文件路径
+    """
+    # 确保目录存在
+    output_dir = Path(output_path).parent
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # 根据扩展名选择保存参数
+    ext = Path(output_path).suffix.lower()
+    
+    if ext in {'.jpg', '.jpeg'}:
+        # 转换为 RGB（JPEG 不支持 alpha 通道）
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        image.save(output_path, quality=quality)
+    elif ext == '.png':
+        image.save(output_path)
+    elif ext == '.webp':
+        image.save(output_path, quality=quality)
+    else:
+        image.save(output_path)
+    
+    return output_path
+
+
+def get_folder_image_count(folder_path: str) -> int:
+    """
+    获取文件夹中的图片数量（不加载图片）
+    
+    Args:
+        folder_path: 文件夹路径
+    
+    Returns:
+        图片数量
+    """
+    folder_path = folder_path.strip()
+    
+    if not folder_path:
+        return 0
+    
+    path = Path(folder_path)
+    
+    if not path.exists() or not path.is_dir():
+        return 0
+    
+    count = 0
+    for file_path in path.iterdir():
+        if file_path.is_file() and file_path.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS:
+            count += 1
+    
+    return count

utils/image_utils.py

@@ -0,0 +1,193 @@
+"""
+图像处理工具模块
+提供 ComfyUI Tensor 与 PIL Image 之间的转换功能
+"""
+
+import base64
+from io import BytesIO
+from typing import List
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
+    """
+    将 ComfyUI 的 Tensor 转换为 PIL Image 列表
+    
+    Args:
+        tensor: 形状为 [B, H, W, C] 的张量，值范围 [0, 1]
+    
+    Returns:
+        PIL Image 列表
+    
+    Example:
+        >>> images = tensor_to_pil(input_tensor)
+        >>> for img in images:
+        ...     img.save(f"output_{i}.png")
+    """
+    images = []
+    
+    # 转换为 numpy 数组
+    np_images = tensor.cpu().numpy()
+    
+    # 处理每张图像
+    for i in range(np_images.shape[0]):
+        img_array = np_images[i]
+        
+        # 转换值范围从 [0, 1] 到 [0, 255]
+        img_array = (img_array * 255).astype(np.uint8)
+        
+        # 创建 PIL Image
+        img = Image.fromarray(img_array)
+        images.append(img)
+    
+    return images
+
+
+def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
+    """
+    将 PIL Image 列表转换为 ComfyUI 的 Tensor
+    
+    Args:
+        images: PIL Image 列表
+    
+    Returns:
+        形状为 [B, H, W, C] 的张量，值范围 [0, 1]
+    
+    Example:
+        >>> pil_images = [Image.open("test.png")]
+        >>> tensor = pil_to_tensor(pil_images)
+        >>> print(tensor.shape)  # [1, H, W, 3]
+    """
+    tensors = []
+    
+    for img in images:
+        # 确保是 RGB 模式
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+        
+        # 转换为 numpy 数组
+        img_array = np.array(img).astype(np.float32)
+        
+        # 转换值范围从 [0, 255] 到 [0, 1]
+        img_array = img_array / 255.0
+        
+        tensors.append(img_array)
+    
+    # 堆叠为批次
+    batch_tensor = np.stack(tensors, axis=0)
+    
+    # 转换为 torch tensor
+    return torch.from_numpy(batch_tensor)
+
+
+def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
+    """
+    将 PIL Image 编码为 base64 字符串
+    
+    Args:
+        image: PIL Image 对象
+        format: 图像格式，默认 PNG
+    
+    Returns:
+        base64 编码的字符串
+    
+    Example:
+        >>> img = Image.open("test.png")
+        >>> b64_str = encode_image_to_base64(img)
+    """
+    buffered = BytesIO()
+    
+    # 转换为 RGB 模式（如果是 RGBA）
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
+    
+    image.save(buffered, format=format)
+    img_bytes = buffered.getvalue()
+    
+    return base64.b64encode(img_bytes).decode('utf-8')
+
+
+def decode_base64_to_pil(base64_string: str) -> Image.Image:
+    """
+    将 base64 字符串解码为 PIL Image
+    
+    Args:
+        base64_string: base64 编码的图像字符串
+    
+    Returns:
+        PIL Image 对象
+    
+    Example:
+        >>> img = decode_base64_to_pil(b64_str)
+        >>> img.save("decoded.png")
+    """
+    img_bytes = base64.b64decode(base64_string)
+    img = Image.open(BytesIO(img_bytes))
+    
+    return img
+
+
+def parse_batch_prompts(prompt: str) -> List[str]:
+    """
+    解析批量提示词
+    
+    检测单独行的 --- 分隔符，分割提示词。
+    如果 --- 不是单独占据一行，则返回空列表（表示单提示词模式）。
+    
+    Args:
+        prompt: 用户输入的提示词文本
+    
+    Returns:
+        提示词列表。如果未检测到单独行的 ---，返回空列表（表示单提示词模式）
+    
+    Raises:
+        ValueError: 如果所有提示词都为空
+    
+    Example:
+        >>> prompts = parse_batch_prompts("a woman\\n---\\na man")
+        >>> print(prompts)  # ['a woman', 'a man']
+        
+        >>> prompts = parse_batch_prompts("a woman --- a man")
+        >>> print(prompts)  # []  (单提示词模式)
+    """
+    lines = prompt.split('\n')
+    
+    # 检查是否存在单独行的 ---
+    has_separator = False
+    for line in lines:
+        if line.strip() == '---':
+            has_separator = True
+            break
+    
+    # 如果没有单独行的 ---，返回空列表（单提示词模式）
+    if not has_separator:
+        return []
+    
+    # 按单独行的 --- 分割
+    # 先将所有单独行的 --- 替换为特殊标记
+    processed_lines = []
+    for line in lines:
+        if line.strip() == '---':
+            processed_lines.append('<<<SEPARATOR>>>')
+        else:
+            processed_lines.append(line)
+    
+    # 重新组合并分割
+    processed_text = '\n'.join(processed_lines)
+    raw_prompts = processed_text.split('<<<SEPARATOR>>>')
+    
+    # 过滤空提示词
+    filtered_prompts = []
+    for p in raw_prompts:
+        stripped = p.strip()
+        if stripped:
+            filtered_prompts.append(stripped)
+    
+    # 如果所有提示词都为空，抛出错误
+    if not filtered_prompts:
+        raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
+    
+    return filtered_prompts

utils/update_checker.py

@@ -0,0 +1,83 @@
+"""
+更新检查工具
+在插件加载时检查是否有新版本
+"""
+
+import os
+import subprocess
+from typing import Optional
+
+
+def get_current_version() -> Optional[str]:
+    """
+    获取当前版本号
+    
+    Returns:
+        版本号字符串，如果读取失败返回 None
+    """
+    version_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "version.txt")
+    try:
+        with open(version_file, 'r', encoding='utf-8') as f:
+            return f.read().strip()
+    except Exception:
+        return None
+
+
+def check_for_updates() -> bool:
+    """
+    检查是否有更新
+    
+    Returns:
+        True 如果有更新，False 如果已是最新或检查失败
+    """
+    try:
+        # 获取当前目录
+        plugin_dir = os.path.dirname(os.path.dirname(__file__))
+        
+        # 检查是否是 Git 仓库
+        git_dir = os.path.join(plugin_dir, '.git')
+        if not os.path.exists(git_dir):
+            return False
+        
+        # 执行 git fetch
+        subprocess.run(
+            ['git', 'fetch', 'origin'],
+            cwd=plugin_dir,
+            capture_output=True,
+            timeout=10
+        )
+        
+        # 检查本地和远程版本
+        local = subprocess.run(
+            ['git', 'rev-parse', '@'],
+            cwd=plugin_dir,
+            capture_output=True,
+            text=True
+        ).stdout.strip()
+        
+        remote = subprocess.run(
+            ['git', 'rev-parse', '@{u}'],
+            cwd=plugin_dir,
+            capture_output=True,
+            text=True
+        ).stdout.strip()
+        
+        return local != remote
+        
+    except Exception:
+        return False
+
+
+def notify_update_available():
+    """通知用户有更新可用"""
+    current_version = get_current_version()
+    version_str = f" (当前版本: {current_version})" if current_version else ""
+    
+    print("\n" + "="*60)
+    print(f"🎉 Comfyui_o1key 有新版本可用{version_str}")
+    print("="*60)
+    print("更新方法：")
+    print("  Windows: 双击运行 update.bat")
+    print("  Linux/Mac: 运行 ./update.sh")
+    print("或手动执行: git pull origin main")
+    print("="*60 + "\n")