369 lines
11 KiB
Python
369 lines
11 KiB
Python
"""
|
||
图像处理工具模块
|
||
提供 ComfyUI Tensor 与 PIL Image 之间的转换功能
|
||
"""
|
||
|
||
import base64
|
||
from io import BytesIO
|
||
import json
|
||
from typing import 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 = []
|
||
|
||
# 转换为 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')
|
||
|
||
|
||
_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
|