Publish current ComfyUI O1Key code baseline

Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
2026-09-24 19:56:48 +08:00
parent 3e337722ab
commit ba920f2b66
183 changed files with 49496 additions and 9934 deletions
@@ -6,7 +6,7 @@
 import base64
 from io import BytesIO
 import json
-from typing import Callable, List, Tuple
+from typing import Any, Callable, List, Tuple

 import numpy as np
 import torch
@@ -29,6 +29,7 @@ def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
        ...     img.save(f"output_{i}.png")
    """
    images = []
+    source_metadata = getattr(tensor, "_o1key_source_metadata", None)
    
    # 转换为 numpy 数组
    np_images = tensor.cpu().numpy()
@@ -42,6 +43,22 @@ def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
        
        # 创建 PIL Image
        img = Image.fromarray(img_array)
+        if isinstance(source_metadata, list) and i < len(source_metadata):
+            metadata = source_metadata[i]
+            if isinstance(metadata, dict):
+                source_format = metadata.get("format")
+                if source_format:
+                    img.format = source_format
+                    setattr(img, "_o1key_original_format", source_format)
+                source_path = metadata.get("path")
+                if source_path:
+                    setattr(img, "_o1key_original_path", source_path)
+                source_filename = metadata.get("filename")
+                if source_filename:
+                    setattr(img, "_o1key_original_filename", source_filename)
+                source_bytes = metadata.get("bytes")
+                if isinstance(source_bytes, bytes):
+                    setattr(img, "_o1key_original_bytes", source_bytes)
        images.append(img)
    
    return images
@@ -63,8 +80,16 @@ def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
        >>> print(tensor.shape)  # [1, H, W, 3]
    """
    tensors = []
+    source_metadata = []

    for img in images:
+        source_metadata.append({
+            "format": getattr(img, "_o1key_original_format", None) or img.format,
+            "path": getattr(img, "_o1key_original_path", None),
+            "filename": getattr(img, "_o1key_original_filename", None),
+            "bytes": getattr(img, "_o1key_original_bytes", None),
+            "modified": bool(getattr(img, "_o1key_pixels_modified", False)),
+        })
        # 确保是 RGB 模式
        if img.mode != 'RGB':
            img = img.convert('RGB')
@@ -81,7 +106,9 @@ def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
    batch_tensor = np.stack(tensors, axis=0)

    # 转换为 torch tensor
-    return torch.from_numpy(batch_tensor)
+    tensor = torch.from_numpy(batch_tensor)
+    tensor._o1key_source_metadata = source_metadata
+    return tensor


 def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
@@ -366,3 +393,106 @@ def parse_batch_prompts(prompt: str) -> List[str]:
        raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
    
    return filtered_prompts
+
+
+def expand_batch_prompt_tasks(prompt: str, images_per_prompt: int) -> List[str]:
+    """Expand one or more prompts into prompt-major generation tasks."""
+    count = int(images_per_prompt)
+    if count < 1:
+        raise ValueError("每条提示词的生图数量必须大于 0")
+    prompts = parse_batch_prompts(prompt) or [prompt.strip()]
+    return [task_prompt for task_prompt in prompts for _ in range(count)]
+
+
+IMAGE_BATCH_MODE_GROUP_TO_MODELS = "一组搭配+多模特"
+IMAGE_BATCH_MODE_CARTESIAN = "全部搭配×全部模特"
+IMAGE_BATCH_MODE_SINGLE_REFERENCES = "单图素材批量"
+IMAGE_BATCH_MODES = (
+    IMAGE_BATCH_MODE_GROUP_TO_MODELS,
+    IMAGE_BATCH_MODE_CARTESIAN,
+    IMAGE_BATCH_MODE_SINGLE_REFERENCES,
+)
+
+
+def expand_image_generation_tasks(
+    prompt: str,
+    images_per_pair: int,
+    *,
+    batch_enabled: bool = False,
+    batch_mode: str = IMAGE_BATCH_MODE_GROUP_TO_MODELS,
+    reference_count: int = 0,
+    model_reference_count: int = 0,
+) -> list[dict[str, Any]]:
+    """Expand prompts and reference pairing into stable prompt-major tasks.
+
+    Normal mode keeps the historical ``prompt x image_count`` ordering. In
+    batch mode, ``reference_indices`` point into the outfit/reference list and
+    ``model_reference_indices`` point into the separately uploaded model list.
+    Single-reference batch mode creates one task per source image and ignores
+    the separately uploaded target list. Keeping indexes instead of file data
+    makes the plan safe to serialize and lets direct and background execution
+    share the exact same ordering.
+    """
+    count = int(images_per_pair)
+    if count < 1:
+        raise ValueError("每个组合的生图数量必须大于 0")
+    prompts = parse_batch_prompts(prompt) or [prompt.strip()]
+    if not batch_enabled:
+        pairings = [{
+            "reference_indices": tuple(range(max(0, int(reference_count)))),
+            "model_reference_indices": (),
+            "outfit_index": None,
+            "model_index": None,
+        }]
+    else:
+        if batch_mode not in IMAGE_BATCH_MODES:
+            raise ValueError("批量模式无效")
+        outfit_total = max(0, int(reference_count))
+        model_total = max(0, int(model_reference_count))
+        if batch_mode == IMAGE_BATCH_MODE_SINGLE_REFERENCES:
+            if outfit_total < 1:
+                raise ValueError("单图批量至少需要上传1张素材图")
+            pairings = [
+                {
+                    "reference_indices": (outfit_index,),
+                    "model_reference_indices": (),
+                    "outfit_index": outfit_index,
+                    "model_index": None,
+                }
+                for outfit_index in range(outfit_total)
+            ]
+        elif model_total < 1:
+            raise ValueError("批量出图至少需要上传1张目标图")
+        elif outfit_total < 1:
+            raise ValueError("批量出图至少需要上传1张素材图")
+        elif batch_mode == IMAGE_BATCH_MODE_GROUP_TO_MODELS:
+            pairings = [
+                {
+                    "reference_indices": tuple(range(outfit_total)),
+                    "model_reference_indices": (model_index,),
+                    "outfit_index": None,
+                    "model_index": model_index,
+                }
+                for model_index in range(model_total)
+            ]
+        else:
+            pairings = [
+                {
+                    "reference_indices": (outfit_index,),
+                    "model_reference_indices": (model_index,),
+                    "outfit_index": outfit_index,
+                    "model_index": model_index,
+                }
+                for outfit_index in range(outfit_total)
+                for model_index in range(model_total)
+            ]
+
+    return [
+        {
+            "prompt": task_prompt,
+            **pairing,
+        }
+        for task_prompt in prompts
+        for pairing in pairings
+        for _ in range(count)
+    ]