296 lines
10 KiB
Python
296 lines
10 KiB
Python
"""
|
|
元数据处理模块
|
|
|
|
负责处理图像元数据的生成、格式化和缓存。
|
|
"""
|
|
|
|
import json
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime
|
|
|
|
|
|
class MetadataHandler:
|
|
"""元数据处理器"""
|
|
|
|
def __init__(self):
|
|
self.cache = {}
|
|
self.cache_enabled = True
|
|
|
|
def prepare_metadata(self, prompt: Dict, extra_pnginfo: Optional[Dict] = None,
|
|
positive_text: Optional[str] = None,
|
|
negative_text: Optional[str] = None,
|
|
metadata_mode: str = "comfy",
|
|
width: Optional[int] = None,
|
|
height: Optional[int] = None) -> Dict[str, Any]:
|
|
"""
|
|
准备元数据
|
|
|
|
Args:
|
|
prompt: ComfyUI提示字典
|
|
extra_pnginfo: 额外的PNG信息
|
|
positive_text: 正面提示文本
|
|
negative_text: 负面提示文本
|
|
|
|
Returns:
|
|
格式化的元数据字典
|
|
"""
|
|
# 生成缓存键
|
|
cache_key = self._generate_cache_key(prompt, extra_pnginfo, positive_text, negative_text, metadata_mode, width, height)
|
|
|
|
# 检查缓存
|
|
if self.cache_enabled and cache_key in self.cache:
|
|
return self.cache[cache_key]
|
|
|
|
metadata = {}
|
|
|
|
include_comfy = metadata_mode in ("comfy", "comfy+a1111")
|
|
include_a1111 = metadata_mode in ("a1111", "comfy+a1111")
|
|
|
|
if include_comfy:
|
|
# 添加基本提示信息
|
|
if prompt:
|
|
metadata['prompt'] = prompt
|
|
|
|
# 添加额外信息，例如 workflow
|
|
if extra_pnginfo:
|
|
metadata.update(extra_pnginfo)
|
|
|
|
# 添加文本提示（如果提供）
|
|
if positive_text:
|
|
metadata['positive_prompt'] = positive_text
|
|
|
|
if negative_text:
|
|
metadata['negative_prompt'] = negative_text
|
|
|
|
if include_a1111:
|
|
metadata['parameters'] = self.format_a1111_parameters(
|
|
prompt, positive_text, negative_text, width, height
|
|
)
|
|
|
|
# 添加生成时间戳
|
|
metadata['generated_at'] = datetime.now().isoformat()
|
|
|
|
# 添加插件信息
|
|
metadata['generator'] = {
|
|
'name': 'ComfyUI Save Image Pro',
|
|
'version': '1.1.0'
|
|
}
|
|
|
|
# 缓存结果
|
|
if self.cache_enabled:
|
|
self.cache[cache_key] = metadata
|
|
|
|
return metadata
|
|
|
|
def _generate_cache_key(self, prompt: Dict, extra_pnginfo: Optional[Dict],
|
|
positive_text: Optional[str], negative_text: Optional[str],
|
|
metadata_mode: str = "comfy", width: Optional[int] = None,
|
|
height: Optional[int] = None) -> str:
|
|
"""生成缓存键"""
|
|
# 使用JSON序列化生成稳定的键
|
|
key_data = {
|
|
'prompt_id': id(prompt),
|
|
'extra_id': id(extra_pnginfo) if extra_pnginfo else None,
|
|
'positive': positive_text,
|
|
'negative': negative_text,
|
|
'metadata_mode': metadata_mode,
|
|
'width': width,
|
|
'height': height
|
|
}
|
|
return json.dumps(key_data, sort_keys=True)
|
|
|
|
def format_for_png(self, metadata: Dict[str, Any]) -> Dict[str, str]:
|
|
"""
|
|
格式化为PNG元数据格式
|
|
|
|
Args:
|
|
metadata: 原始元数据
|
|
|
|
Returns:
|
|
PNG格式的元数据字典
|
|
"""
|
|
png_metadata = {}
|
|
|
|
for key, value in metadata.items():
|
|
try:
|
|
if isinstance(value, (dict, list)):
|
|
png_metadata[key] = json.dumps(value, ensure_ascii=False)
|
|
else:
|
|
png_metadata[key] = str(value)
|
|
except (TypeError, ValueError):
|
|
png_metadata[key] = str(value)
|
|
|
|
return png_metadata
|
|
|
|
def format_for_exif(self, metadata: Dict[str, Any]) -> bytes:
|
|
"""
|
|
格式化为EXIF元数据格式
|
|
|
|
Args:
|
|
metadata: 原始元数据
|
|
|
|
Returns:
|
|
EXIF字节数据
|
|
"""
|
|
try:
|
|
from PIL import Image
|
|
|
|
# 创建临时图像来生成EXIF
|
|
temp_img = Image.new('RGB', (1, 1))
|
|
exif = temp_img.getexif()
|
|
|
|
# 将关键信息编码到EXIF字段
|
|
if 'prompt' in metadata:
|
|
# 0x010f: Make字段存储prompt
|
|
exif[0x010f] = "Prompt: " + json.dumps(metadata['prompt'], ensure_ascii=False)
|
|
|
|
if 'workflow' in metadata:
|
|
# 0x010e: ImageDescription字段存储workflow
|
|
exif[0x010e] = "Workflow: " + json.dumps(metadata['workflow'], ensure_ascii=False)
|
|
|
|
# 添加生成器信息
|
|
if 'generator' in metadata:
|
|
# 0x0131: Software字段
|
|
generator = metadata['generator']
|
|
exif[0x0131] = f"{generator.get('name', 'Unknown')} v{generator.get('version', '1.0')}"
|
|
|
|
return exif.tobytes()
|
|
|
|
except Exception as e:
|
|
print(f"EXIF元数据生成错误: {e}")
|
|
return b''
|
|
|
|
|
|
def format_a1111_parameters(self, prompt: Dict, positive_text: Optional[str] = None,
|
|
negative_text: Optional[str] = None,
|
|
width: Optional[int] = None,
|
|
height: Optional[int] = None) -> str:
|
|
"""Build an Automatic1111-style parameters string for sharing sites."""
|
|
positive = positive_text or self._find_prompt_text(prompt, prefer_negative=False) or ""
|
|
negative = negative_text or self._find_prompt_text(prompt, prefer_negative=True) or ""
|
|
params = self._extract_common_parameters(prompt or {})
|
|
|
|
tail_parts = []
|
|
mapping = [
|
|
("steps", "Steps"),
|
|
("sampler_name", "Sampler"),
|
|
("scheduler", "Schedule type"),
|
|
("cfg", "CFG scale"),
|
|
("seed", "Seed"),
|
|
]
|
|
for source_key, label in mapping:
|
|
value = params.get(source_key)
|
|
if value:
|
|
tail_parts.append(f"{label}: {value}")
|
|
|
|
if width and height:
|
|
tail_parts.append(f"Size: {width}x{height}")
|
|
|
|
model_name = params.get("ckpt_name") or params.get("model")
|
|
if model_name:
|
|
model_name = str(model_name).rsplit("/", 1)[-1].rsplit("\\", 1)[-1]
|
|
for ext in (".safetensors", ".ckpt", ".pt", ".bin", ".pth"):
|
|
if model_name.endswith(ext):
|
|
model_name = model_name[:-len(ext)]
|
|
break
|
|
tail_parts.append(f"Model: {model_name}")
|
|
|
|
lines = [positive]
|
|
if negative:
|
|
lines.append(f"Negative prompt: {negative}")
|
|
if tail_parts:
|
|
lines.append(", ".join(tail_parts))
|
|
return "\n".join(lines).strip()
|
|
|
|
def _find_prompt_text(self, prompt: Dict, prefer_negative: bool = False) -> Optional[str]:
|
|
texts = []
|
|
if isinstance(prompt, dict):
|
|
for node in prompt.values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get('inputs', {})
|
|
class_type = str(node.get('class_type', '')).lower()
|
|
text = inputs.get('text') if isinstance(inputs, dict) else None
|
|
if text and ('cliptextencode' in class_type or 'text' in class_type):
|
|
texts.append(str(text))
|
|
if not texts:
|
|
return None
|
|
return texts[-1] if prefer_negative and len(texts) > 1 else texts[0]
|
|
|
|
def extract_parameters(self, metadata: Dict[str, Any]) -> Dict[str, str]:
|
|
"""
|
|
从元数据中提取关键参数
|
|
|
|
Args:
|
|
metadata: 元数据字典
|
|
|
|
Returns:
|
|
参数字典
|
|
"""
|
|
parameters = {}
|
|
|
|
if 'prompt' in metadata:
|
|
prompt = metadata['prompt']
|
|
|
|
# 递归提取常用参数
|
|
extracted = self._extract_common_parameters(prompt)
|
|
parameters.update(extracted)
|
|
|
|
# 添加文本提示
|
|
if 'positive_prompt' in metadata:
|
|
parameters['positive_prompt'] = metadata['positive_prompt']
|
|
|
|
if 'negative_prompt' in metadata:
|
|
parameters['negative_prompt'] = metadata['negative_prompt']
|
|
|
|
return parameters
|
|
|
|
def _extract_common_parameters(self, obj: Any, prefix: str = "") -> Dict[str, str]:
|
|
"""递归提取常用参数"""
|
|
parameters = {}
|
|
|
|
# 常用参数名列表
|
|
common_params = [
|
|
'sampler_name', 'scheduler', 'steps', 'cfg', 'seed',
|
|
'width', 'height', 'denoise', 'model', 'ckpt_name',
|
|
'vae_name', 'clip_skip', 'lora_name'
|
|
]
|
|
|
|
if isinstance(obj, dict):
|
|
for key, value in obj.items():
|
|
current_key = f"{prefix}.{key}" if prefix else key
|
|
|
|
# 检查是否是常用参数
|
|
if key.lower() in [p.lower() for p in common_params]:
|
|
parameters[key] = str(value)
|
|
|
|
# 递归搜索
|
|
if isinstance(value, (dict, list)):
|
|
nested = self._extract_common_parameters(value, current_key)
|
|
parameters.update(nested)
|
|
|
|
elif isinstance(obj, list):
|
|
for i, item in enumerate(obj):
|
|
current_key = f"{prefix}[{i}]" if prefix else f"[{i}]"
|
|
nested = self._extract_common_parameters(item, current_key)
|
|
parameters.update(nested)
|
|
|
|
return parameters
|
|
|
|
def clear_cache(self):
|
|
"""清空缓存"""
|
|
self.cache.clear()
|
|
|
|
def set_cache_enabled(self, enabled: bool):
|
|
"""启用或禁用缓存"""
|
|
self.cache_enabled = enabled
|
|
if not enabled:
|
|
self.cache.clear()
|
|
|
|
def get_cache_stats(self) -> Dict[str, int]:
|
|
"""获取缓存统计"""
|
|
return {
|
|
'cache_size': len(self.cache),
|
|
'cache_enabled': self.cache_enabled
|
|
}
|