273 lines
11 KiB
Python
273 lines
11 KiB
Python
|
"""
|
||
|
|
Извлечение и разбор метаданных ComfyUI из изображений (PNG/WebP/JPEG).
|
||
|
|
"""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Optional, Dict, Any, Tuple
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
try:
|
||
|
|
import piexif
|
||
|
|
import piexif.helper
|
||
|
|
HAS_PIEXIF = True
|
||
|
|
except ImportError:
|
||
|
|
HAS_PIEXIF = False
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class MetadataParser:
|
||
|
|
@staticmethod
|
||
|
|
def extract_raw_metadata(filepath: str) -> Tuple[Optional[str], Optional[str]]:
|
||
|
|
"""Возвращает (prompt_json, workflow_json) из файла изображения."""
|
||
|
|
path = Path(filepath)
|
||
|
|
if not path.exists():
|
||
|
|
return None, None
|
||
|
|
suffix = path.suffix.lower()
|
||
|
|
if suffix in ('.png', '.webp'):
|
||
|
|
return MetadataParser._extract_from_png_webp(path)
|
||
|
|
elif suffix in ('.jpg', '.jpeg'):
|
||
|
|
return MetadataParser._extract_from_jpeg(path)
|
||
|
|
return None, None
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _extract_from_png_webp(path: Path) -> Tuple[Optional[str], Optional[str]]:
|
||
|
|
try:
|
||
|
|
with Image.open(path) as img:
|
||
|
|
info = img.info
|
||
|
|
prompt = info.get("prompt")
|
||
|
|
workflow = info.get("workflow")
|
||
|
|
if not prompt and "comment" in info:
|
||
|
|
prompt = info.get("comment")
|
||
|
|
return prompt, workflow
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Ошибка чтения PNG/WebP метаданных {path.name}: {e}")
|
||
|
|
return None, None
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _extract_from_jpeg(path: Path) -> Tuple[Optional[str], Optional[str]]:
|
||
|
|
if not HAS_PIEXIF:
|
||
|
|
logger.warning("piexif не установлен — JPEG-метаданные недоступны.")
|
||
|
|
return None, None
|
||
|
|
try:
|
||
|
|
with Image.open(path) as img:
|
||
|
|
if "exif" not in img.info:
|
||
|
|
return None, None
|
||
|
|
exif_dict = piexif.load(img.info["exif"])
|
||
|
|
user_comment_bytes = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
|
||
|
|
if not user_comment_bytes:
|
||
|
|
return None, None
|
||
|
|
try:
|
||
|
|
comment_str = piexif.helper.UserComment.load(user_comment_bytes)
|
||
|
|
except ValueError:
|
||
|
|
comment_str = user_comment_bytes.decode('utf-8', errors='ignore')
|
||
|
|
if comment_str.startswith("{"):
|
||
|
|
try:
|
||
|
|
data = json.loads(comment_str)
|
||
|
|
prompt = data.get("prompt")
|
||
|
|
workflow = data.get("workflow")
|
||
|
|
if isinstance(prompt, dict):
|
||
|
|
prompt = json.dumps(prompt)
|
||
|
|
if isinstance(workflow, dict):
|
||
|
|
workflow = json.dumps(workflow)
|
||
|
|
return prompt, workflow
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return comment_str, None
|
||
|
|
return None, None
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Ошибка чтения JPEG EXIF {path.name}: {e}")
|
||
|
|
return None, None
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def parse_comfy_parameters(cls, prompt_json_str: Optional[str]) -> Dict[str, Any]:
|
||
|
|
"""Разбирает prompt-граф ComfyUI и вытаскивает основные параметры генерации."""
|
||
|
|
result = {
|
||
|
|
"positive_prompt": None, "negative_prompt": None, "seed": None,
|
||
|
|
"model_name": None, "sampler": None, "steps": None, "cfg": None
|
||
|
|
}
|
||
|
|
if not prompt_json_str:
|
||
|
|
return result
|
||
|
|
try:
|
||
|
|
prompt_graph = json.loads(prompt_json_str)
|
||
|
|
if not isinstance(prompt_graph, dict):
|
||
|
|
return result
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return result
|
||
|
|
|
||
|
|
sampler_node = None
|
||
|
|
for node_id, node in prompt_graph.items():
|
||
|
|
if not isinstance(node, dict):
|
||
|
|
continue
|
||
|
|
class_type = node.get("class_type", "")
|
||
|
|
if "KSampler" in class_type:
|
||
|
|
sampler_node = node
|
||
|
|
break
|
||
|
|
|
||
|
|
if sampler_node:
|
||
|
|
inputs = sampler_node.get("inputs", {})
|
||
|
|
result["seed"] = inputs.get("seed") or inputs.get("noise_seed")
|
||
|
|
result["steps"] = inputs.get("steps")
|
||
|
|
result["cfg"] = inputs.get("cfg")
|
||
|
|
result["sampler"] = inputs.get("sampler_name")
|
||
|
|
result["positive_prompt"] = cls._trace_conditioning(inputs.get("positive"), prompt_graph)
|
||
|
|
result["negative_prompt"] = cls._trace_conditioning(inputs.get("negative"), prompt_graph)
|
||
|
|
result["model_name"] = cls._trace_model(inputs.get("model"), prompt_graph)
|
||
|
|
else:
|
||
|
|
# Fallback: собираем все CLIPTextEncode-ноды
|
||
|
|
positives = []
|
||
|
|
for node in prompt_graph.values():
|
||
|
|
if isinstance(node, dict) and node.get("class_type") == "CLIPTextEncode":
|
||
|
|
text = node.get("inputs", {}).get("text", "")
|
||
|
|
if isinstance(text, str) and text.strip():
|
||
|
|
positives.append(text.strip())
|
||
|
|
if positives:
|
||
|
|
result["positive_prompt"] = "\n---\n".join(positives)
|
||
|
|
|
||
|
|
def clean_string(val) -> Optional[str]:
|
||
|
|
if val is None:
|
||
|
|
return None
|
||
|
|
if isinstance(val, list):
|
||
|
|
if all(isinstance(x, str) for x in val):
|
||
|
|
return "\n".join(val)
|
||
|
|
return json.dumps(val)
|
||
|
|
if isinstance(val, dict):
|
||
|
|
return json.dumps(val)
|
||
|
|
return str(val)
|
||
|
|
|
||
|
|
def clean_int(val) -> Optional[int]:
|
||
|
|
if val is None or isinstance(val, (list, dict)):
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
return int(val)
|
||
|
|
except (ValueError, TypeError):
|
||
|
|
return None
|
||
|
|
|
||
|
|
def clean_float(val) -> Optional[float]:
|
||
|
|
if val is None or isinstance(val, (list, dict)):
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
return float(val)
|
||
|
|
except (ValueError, TypeError):
|
||
|
|
return None
|
||
|
|
|
||
|
|
result["positive_prompt"] = clean_string(result["positive_prompt"])
|
||
|
|
result["negative_prompt"] = clean_string(result["negative_prompt"])
|
||
|
|
result["model_name"] = clean_string(result["model_name"])
|
||
|
|
result["sampler"] = clean_string(result["sampler"])
|
||
|
|
result["seed"] = clean_int(result["seed"])
|
||
|
|
result["steps"] = clean_int(result["steps"])
|
||
|
|
result["cfg"] = clean_float(result["cfg"])
|
||
|
|
return result
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def inject_parameters(cls, prompt_json_str: str,
|
||
|
|
positive: Optional[str] = None,
|
||
|
|
negative: Optional[str] = None,
|
||
|
|
seed: Optional[int] = None) -> str:
|
||
|
|
"""
|
||
|
|
Внедряет отредактированные параметры обратно в prompt-граф ComfyUI.
|
||
|
|
Возвращает обновлённый JSON (или исходный, если граф не разобрать).
|
||
|
|
|
||
|
|
Важно: без изменения графа ComfyUI отдаёт всё из кэша и не генерирует
|
||
|
|
("Prompt executed in 0.00 seconds"), поэтому при повторной генерации
|
||
|
|
нужно передавать новый seed.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
graph = json.loads(prompt_json_str)
|
||
|
|
if not isinstance(graph, dict):
|
||
|
|
return prompt_json_str
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return prompt_json_str
|
||
|
|
|
||
|
|
sampler_node = None
|
||
|
|
for node in graph.values():
|
||
|
|
if isinstance(node, dict) and "KSampler" in node.get("class_type", ""):
|
||
|
|
sampler_node = node
|
||
|
|
break
|
||
|
|
if not sampler_node:
|
||
|
|
return prompt_json_str
|
||
|
|
|
||
|
|
inputs = sampler_node.get("inputs", {})
|
||
|
|
|
||
|
|
if seed is not None:
|
||
|
|
# Ключ зависит от типа сэмплера: seed или noise_seed
|
||
|
|
if "noise_seed" in inputs:
|
||
|
|
inputs["noise_seed"] = seed
|
||
|
|
else:
|
||
|
|
inputs["seed"] = seed
|
||
|
|
|
||
|
|
if positive is not None:
|
||
|
|
text_node_id = cls._trace_text_node_id(inputs.get("positive"), graph)
|
||
|
|
if text_node_id:
|
||
|
|
graph[text_node_id]["inputs"]["text"] = positive
|
||
|
|
|
||
|
|
if negative is not None:
|
||
|
|
text_node_id = cls._trace_text_node_id(inputs.get("negative"), graph)
|
||
|
|
if text_node_id:
|
||
|
|
graph[text_node_id]["inputs"]["text"] = negative
|
||
|
|
|
||
|
|
return json.dumps(graph, ensure_ascii=False)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _trace_text_node_id(cls, link: Optional[list], graph: dict,
|
||
|
|
_depth: int = 0) -> Optional[str]:
|
||
|
|
"""Возвращает ID текстовой ноды (CLIPTextEncode), до которой ведёт ссылка."""
|
||
|
|
if _depth > 20 or not link or not isinstance(link, list) or len(link) < 1:
|
||
|
|
return None
|
||
|
|
node_id = str(link[0])
|
||
|
|
node = graph.get(node_id)
|
||
|
|
if not node:
|
||
|
|
return None
|
||
|
|
class_type = node.get("class_type", "")
|
||
|
|
if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL", "CLIPTextEncodeSequence"):
|
||
|
|
return node_id
|
||
|
|
if "Conditioning" in class_type:
|
||
|
|
for val in node.get("inputs", {}).values():
|
||
|
|
if isinstance(val, list) and len(val) >= 1:
|
||
|
|
found = cls._trace_text_node_id(val, graph, _depth + 1)
|
||
|
|
if found:
|
||
|
|
return found
|
||
|
|
return None
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _trace_conditioning(cls, link: Optional[list], graph: dict,
|
||
|
|
_depth: int = 0) -> Optional[str]:
|
||
|
|
"""Рекурсивно идёт по ссылкам графа до текстовой ноды промта."""
|
||
|
|
if _depth > 20 or not link or not isinstance(link, list) or len(link) < 1:
|
||
|
|
return None
|
||
|
|
node = graph.get(str(link[0]))
|
||
|
|
if not node:
|
||
|
|
return None
|
||
|
|
class_type = node.get("class_type", "")
|
||
|
|
inputs = node.get("inputs", {})
|
||
|
|
if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL", "CLIPTextEncodeSequence"):
|
||
|
|
return inputs.get("text") or inputs.get("text_g")
|
||
|
|
if "Conditioning" in class_type:
|
||
|
|
for val in inputs.values():
|
||
|
|
if isinstance(val, list) and len(val) >= 1:
|
||
|
|
text = cls._trace_conditioning(val, graph, _depth + 1)
|
||
|
|
if text:
|
||
|
|
return text
|
||
|
|
return None
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _trace_model(cls, link: Optional[list], graph: dict, _depth: int = 0) -> Optional[str]:
|
||
|
|
"""Рекурсивно идёт по ссылкам до ноды загрузки чекпоинта."""
|
||
|
|
if _depth > 20 or not link or not isinstance(link, list) or len(link) < 1:
|
||
|
|
return None
|
||
|
|
node = graph.get(str(link[0]))
|
||
|
|
if not node:
|
||
|
|
return None
|
||
|
|
class_type = node.get("class_type", "")
|
||
|
|
inputs = node.get("inputs", {})
|
||
|
|
if "CheckpointLoader" in class_type:
|
||
|
|
return inputs.get("ckpt_name")
|
||
|
|
elif "UNETLoader" in class_type:
|
||
|
|
return inputs.get("unet_name")
|
||
|
|
elif "LoraLoader" in class_type or "ModelMerge" in class_type:
|
||
|
|
return cls._trace_model(inputs.get("model"), graph, _depth + 1)
|
||
|
|
return None
|