Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
499 lines
16 KiB
Python
499 lines
16 KiB
Python
"""
|
||
图像处理工具模块
|
||
提供 ComfyUI Tensor 与 PIL Image 之间的转换功能
|
||
"""
|
||
|
||
import base64
|
||
from io import BytesIO
|
||
import json
|
||
from typing import Any, Callable, List, Tuple
|
||
|
||
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 = []
|
||
source_metadata = getattr(tensor, "_o1key_source_metadata", None)
|
||
|
||
# 转换为 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)
|
||
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
|
||
|
||
|
||
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 = []
|
||
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')
|
||
|
||
# 转换为 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
|
||
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:
|
||
"""
|
||
将 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')
|
||
|
||
|
||
_MAX_REQUEST_BODY_BYTES = 50 * 1024 * 1024 # 50MB 请求体上限
|
||
|
||
|
||
def _encode_image_to_base64_with_quality(image: Image.Image, quality: int) -> str:
|
||
buffered = BytesIO()
|
||
working = image
|
||
if working.mode != 'RGB':
|
||
working = working.convert('RGB')
|
||
|
||
working.save(
|
||
buffered,
|
||
format="JPEG",
|
||
quality=quality,
|
||
optimize=True,
|
||
subsampling=2,
|
||
)
|
||
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||
|
||
|
||
def encode_images_for_request_body_limit(
|
||
images: List[Image.Image],
|
||
build_body: Callable[[List[Tuple[str, str]]], dict],
|
||
max_body_bytes: int = _MAX_REQUEST_BODY_BYTES,
|
||
) -> List[Tuple[str, str]]:
|
||
"""
|
||
为请求体编码图片，并保证完整 JSON 请求体不超过 max_body_bytes。
|
||
|
||
策略：
|
||
- 先按原始 PNG 编码估算完整请求体；
|
||
- 若超过限制，改用 JPEG 质量压缩，逐步降低 quality；
|
||
- 全程不缩放图片尺寸。
|
||
|
||
Returns:
|
||
[(mime_type, base64), ...]
|
||
"""
|
||
encoded = [("image/png", encode_image_to_base64(img, format="PNG")) for img in images]
|
||
body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
|
||
if body_size <= max_body_bytes:
|
||
return encoded
|
||
|
||
for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
|
||
encoded = [
|
||
("image/jpeg", _encode_image_to_base64_with_quality(img, quality))
|
||
for img in images
|
||
]
|
||
body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
|
||
if body_size <= max_body_bytes:
|
||
print(
|
||
f"输入图片已通过 JPEG 质量压缩控制请求体积: "
|
||
f"quality={quality}, 请求体积={body_size / 1024 / 1024:.2f}MB "
|
||
f"(限制 {max_body_bytes / 1024 / 1024:.0f}MB)"
|
||
)
|
||
return encoded
|
||
|
||
raise ValueError(
|
||
f"请求体超过 {max_body_bytes / 1024 / 1024:.0f}MB，"
|
||
"即使压缩到最低图片质量仍无法满足限制；请减少参考图数量或输入图片内容复杂度"
|
||
)
|
||
|
||
|
||
_MAX_IMAGE_BYTES = 10 * 1024 * 1024 # 10MB 单张图片上限
|
||
|
||
|
||
def _encode_image_to_bytes(image: Image.Image, format: str = "PNG", quality: int = None) -> bytes:
|
||
buffered = BytesIO()
|
||
working = image
|
||
|
||
if format.upper() == "JPEG" and working.mode != 'RGB':
|
||
working = working.convert('RGB')
|
||
elif working.mode == 'RGBA':
|
||
working = working.convert('RGB')
|
||
|
||
save_kwargs = {"format": format}
|
||
if quality is not None:
|
||
save_kwargs.update({
|
||
"quality": quality,
|
||
"optimize": True,
|
||
"subsampling": 2,
|
||
})
|
||
|
||
working.save(buffered, **save_kwargs)
|
||
return buffered.getvalue()
|
||
|
||
|
||
def encode_images_for_image_size_limit(
|
||
images: List[Image.Image],
|
||
max_image_bytes: int = _MAX_IMAGE_BYTES,
|
||
) -> List[Tuple[str, str]]:
|
||
"""
|
||
将图片编码为 base64，并保证每张编码前的图片文件体积不超过 max_image_bytes。
|
||
|
||
策略：
|
||
- 先尝试 PNG 原图尺寸编码；
|
||
- 单张超过限制时，改用 JPEG 质量压缩；
|
||
- 全程不缩放图片尺寸。
|
||
|
||
Returns:
|
||
[(mime_type, base64), ...]
|
||
"""
|
||
encoded = []
|
||
|
||
for idx, img in enumerate(images, start=1):
|
||
png_bytes = _encode_image_to_bytes(img, format="PNG")
|
||
if len(png_bytes) <= max_image_bytes:
|
||
encoded.append(("image/png", base64.b64encode(png_bytes).decode('utf-8')))
|
||
continue
|
||
|
||
for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
|
||
jpg_bytes = _encode_image_to_bytes(img, format="JPEG", quality=quality)
|
||
if len(jpg_bytes) <= max_image_bytes:
|
||
print(
|
||
f"输入图片 {idx} 已通过 JPEG 质量压缩控制单图体积: "
|
||
f"quality={quality}, 图片体积={len(jpg_bytes) / 1024 / 1024:.2f}MB "
|
||
f"(限制 {max_image_bytes / 1024 / 1024:.0f}MB)，尺寸保持 {img.width}x{img.height}"
|
||
)
|
||
encoded.append(("image/jpeg", base64.b64encode(jpg_bytes).decode('utf-8')))
|
||
break
|
||
else:
|
||
raise ValueError(
|
||
f"输入图片 {idx} 超过 {max_image_bytes / 1024 / 1024:.0f}MB，"
|
||
"即使压缩到最低图片质量仍无法满足限制；请减少图片内容复杂度或手动处理图片"
|
||
)
|
||
|
||
return encoded
|
||
|
||
|
||
def encode_image_to_base64_limited(
|
||
image: Image.Image,
|
||
format: str = "PNG",
|
||
max_bytes: int = _MAX_IMAGE_BYTES,
|
||
) -> str:
|
||
"""
|
||
将 PIL Image 编码为 base64，若超过 max_bytes 则自动缩放直到满足限制。
|
||
|
||
策略：等比缩放，每轮缩小到上一轮的 80%，最多 10 轮。
|
||
|
||
Args:
|
||
image: PIL Image 对象
|
||
format: 图像格式，默认 PNG
|
||
max_bytes: base64 字符串最大字节数，默认 10MB
|
||
|
||
Returns:
|
||
base64 编码的字符串（保证 <= max_bytes）
|
||
"""
|
||
working = image
|
||
if working.mode == 'RGBA':
|
||
working = working.convert('RGB')
|
||
|
||
for attempt in range(10):
|
||
buffered = BytesIO()
|
||
working.save(buffered, format=format)
|
||
b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||
|
||
if len(b64) <= max_bytes:
|
||
if attempt > 0:
|
||
print(
|
||
f"图片已自动缩放: {image.width}x{image.height} → "
|
||
f"{working.width}x{working.height} "
|
||
f"({len(b64) / 1024 / 1024:.2f}MB)"
|
||
)
|
||
return b64
|
||
|
||
# 缩放到 80%
|
||
scale = 0.8
|
||
new_w = max(1, int(working.width * scale))
|
||
new_h = max(1, int(working.height * scale))
|
||
working = working.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||
|
||
# 兜底：返回最后一次编码结果
|
||
buffered = BytesIO()
|
||
working.save(buffered, format=format)
|
||
return base64.b64encode(buffered.getvalue()).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
|
||
|
||
|
||
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)
|
||
]
|