Initial release: Comfy-Studio v1.0.0
Десктопная галерея и менеджер изображений ComfyUI (PyQt6 + SQLite): - мультипапочный мониторинг (watchdog), грид миниатюр, drag-and-drop - теги, избранное, рейтинг, поиск по промтам - парсер метаданных ComfyUI (PNG/WebP/JPEG), редактор workflow/prompt JSON - перевод выделенного текста через контекстное меню (настраиваемые языки) - отправка промта в очередь ComfyUI с внедрением правок и нового seed - полноэкранный просмотр, локализация ru/en, тёмная тема Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
dinlo
@@ -0,0 +1,272 @@
|
||||
"""
|
||||
Извлечение и разбор метаданных 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
|
||||
Reference in New Issue
Block a user