nodes/leestuartx-ComfyUI-GG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Resolve Issues with metadata extract by seed

Browse Source
This commit is contained in:
leestuartx
2024-07-22 07:54:15 -06:00
parent cda23df1ab
commit 8643fbc358
3 changed files with 22 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/common/metadata_parser.py
+4 -4
View File
@@ -13,8 +13,9 @@ class MetadataParser:
prompt_info = self.parser_manager.parse(img)
if prompt_info:
metadata = self.format_metadata(prompt_info)
cfg_scale = self.find_value_in_dict(prompt_info.parameters, 'cfg')
steps = self.find_value_in_dict(prompt_info.parameters, 'steps')
cfg_scale = self.find_value_in_dict(prompt_info.parameters, 'cfg', default=0.0)
steps = self.find_value_in_dict(prompt_info.parameters, 'steps', default=0)
vaes = self.find_values_in_dict(prompt_info.parameters, 'vae_name')
models = [model['content'] if isinstance(model, dict) else model for model in self.find_values_in_dict(prompt_info.parameters, 'ckpt_name')]
sampler_name = self.find_value_in_dict(prompt_info.parameters, 'sampler_name')
@@ -26,11 +27,10 @@ class MetadataParser:
metadata_str_keys = self.convert_keys_to_strings(temp_metadata)
prompt_data = self.find_positive_prompt_data(metadata_str_keys)
positive_prompt = ""
positive_prompt = self.get_prompt_text(prompt_info.prompts)
if len(prompt_data) > 0:
positive_prompt = prompt_data[0]
negative_prompt = self.get_prompt_text(prompt_info.negative_prompts)
clip_skip = self.find_value_in_dict(prompt_info.parameters, 'clip')
return prompt_info, metadata, cfg_scale, steps, width, height, positive_prompt, negative_prompt, clip_skip, vaes, models, sampler_name, scheduler, denoise
nodes/image_metadata_extractor.py
+1 -53
View File
@@ -61,7 +61,7 @@ class ImageMetadataExtractor:
try:
prompt_info = parser_manager.parse(img)
metadata = self.format_metadata(prompt_info)
metadata = parser.format_metadata(prompt_info)
raw_metadata = str(prompt_info)
except Exception as e:
print('Error extracting metadata:', e)
@@ -74,58 +74,6 @@ class ImageMetadataExtractor:
filename, raw_metadata
)
def format_metadata(self, prompt_info):
metadata_parts = []
# Models
if hasattr(prompt_info, 'models') and prompt_info.models:
model_text = '\n'.join([model.name for model in prompt_info.models])
metadata_parts.append(f"Models:\n{model_text}")
# Samplers
if hasattr(prompt_info, 'samplers') and prompt_info.samplers:
sampler_names = '\n'.join([sampler.name for sampler in prompt_info.samplers])
metadata_parts.append(f"Samplers:\n{sampler_names}")
for sampler in prompt_info.samplers:
for param, value in sampler.parameters.items():
metadata_parts.append(f"{param.title()}: {value}")
# Prompts
if hasattr(prompt_info, 'prompts') and prompt_info.prompts:
prompt_text = '\n'.join([prompt.value for prompt in prompt_info.prompts])
metadata_parts.append(f"Prompts:\n{prompt_text}")
# Negative Prompts
if hasattr(prompt_info, 'negative_prompts') and prompt_info.negative_prompts:
negative_prompt_text = '\n'.join([prompt.value for prompt in prompt_info.negative_prompts])
metadata_parts.append(f"Negative Prompts:\n{negative_prompt_text}")
# Additional metadata
if hasattr(prompt_info, 'metadata') and prompt_info.metadata:
additional_metadata = '\n'.join([f"{k}: {v}" for k, v in prompt_info.metadata.items()])
metadata_parts.append(f"Additional Metadata:\n{additional_metadata}")
# Unmodified parameters
if hasattr(prompt_info, 'parameters') and prompt_info.parameters:
params_text = '\n'.join([f"{k}: {v}" for k, v in prompt_info.parameters.items()])
metadata_parts.append(f"Parameters:\n{params_text}")
return '\n\n'.join(metadata_parts)
def get_float(self, parameters, key):
try:
value = str(parameters.get(key, '0.0')).strip()
return float(value)
except (ValueError, TypeError):
return 0.0
def get_int(self, parameters, key):
try:
value = str(parameters.get(key, '0')).strip()
return int(value)
except (ValueError, TypeError):
return 0
NODE_CLASS_MAPPINGS = {
"ImageMetadataExtractor": ImageMetadataExtractor
nodes/metadata_extract_by_seed.py
+17 -147
View File
@@ -5,6 +5,7 @@ import torch
import numpy as np
from sd_parsers import ParserManager # , PromptInfo
from typing import Any, Dict, List
from common.metadata_parser import MetadataParser
# Initialize the parser manager
parser_manager = ParserManager()
@@ -61,40 +62,26 @@ class MetadataExtractorBySeed:
width, height = img.size
filename = os.path.basename(image_path)
parser = MetadataParser()
prompt_info, formatted_metadata, cfg_scale, steps, width, height, positive_prompt, negative_prompt, clip_skip, vaes, models, sampler_name, scheduler, denoise = parser.extract_metadata(image_path)
'''
model_name = ', '.join(models)
# Parse the image directly using the parser manager
prompt_info = parser_manager.parse(img)
if prompt_info:
metadata = self.format_metadata(prompt_info)
raw_metadata = self.format_raw_metadata(prompt_info)
cfg_scale = self.get_float(prompt_info.parameters, 'cfg')
steps = self.get_int(prompt_info.parameters, 'steps')
positive_prompt = self.get_prompt_text(prompt_info.prompts)
negative_prompt = self.get_prompt_text(prompt_info.negative_prompts)
clip_skip = self.get_int(prompt_info.parameters, 'clip')
model_name = self.get_model_name(prompt_info)
try:
prompt_info = parser_manager.parse(img)
metadata = parser.format_metadata(prompt_info)
raw_metadata = str(prompt_info)
except Exception as e:
print('Error extracting metadata:', e)
metadata = ''
raw_metadata = ''
# Try extracting the positive prompt using the new method
temp_metadata = self.extract_metadata_type2(image_path)
metadata_str_keys = self.convert_keys_to_strings(temp_metadata)
prompt_data = self.find_positive_prompt_data(metadata_str_keys)
if len(prompt_data) > 0:
positive_prompt = prompt_data[0]
return (
output_image, output_mask, metadata, cfg_scale, steps, int(width), int(height),
positive_prompt, negative_prompt, clip_skip, model_name,
filename, raw_metadata
)
return (
output_image, output_mask, metadata, cfg_scale, steps, int(width), int(height),
positive_prompt, negative_prompt, clip_skip, model_name,
filename, raw_metadata
)
else:
return (
output_image, output_mask, "No metadata found.", 0.0, 0, int(width), int(height),
"", "", 0, "Unknown",
os.path.basename(image_path), ""
)
'''
except Exception as e:
logging.exception("Error processing image")
return (
@@ -108,123 +95,6 @@ class MetadataExtractorBySeed:
path = os.path.abspath(path)
return path
def extract_metadata_type2(self, file_path: str) -> Dict[str, Any]:
try:
img = Image.open(file_path)
# Extract metadata using sd_parsers
prompt_info = parser_manager.parse(file_path)
if not prompt_info:
raise ValueError("No metadata found in image.")
metadata = {
"prompt": prompt_info.parameters,
"workflow": prompt_info.metadata
}
return metadata
except Exception as e:
raise ValueError(f"Error extracting metadata: {e}")
def convert_keys_to_strings(self, d):
if isinstance(d, dict):
return {str(k): self.convert_keys_to_strings(v) for k, v in d.items()}
elif isinstance(d, list):
return [self.convert_keys_to_strings(i) for i in d]
else:
return d
def find_positive_prompt_data(self, metadata: Dict[str, Any]) -> List[str]:
output_data = []
# Checking in both 'prompt' and 'workflow' sections for relevant nodes
for key, section in metadata.items():
if key == 'workflow' and isinstance(section, dict):
if "nodes" in section:
for node in section["nodes"]:
if node.get("type") == "ShowText|pysssss" and node.get("properties", {}).get("Node name for S&R") == "ShowText|pysssss":
widgets_values = node.get("widgets_values", [])
for widget in widgets_values:
if isinstance(widget, list):
text_value = widget[0]
output_data.append(text_value)
elif isinstance(section, dict):
for subkey, subsection in section.items():
if subkey == 'workflow' and isinstance(subsection, dict):
if "nodes" in subsection:
for node in subsection["nodes"]:
if node.get("type") == "ShowText|pysssss" and node.get("properties", {}).get("Node name for S&R") == "ShowText|pysssss":
widgets_values = node.get("widgets_values", [])
for widget in widgets_values:
if isinstance(widget, list):
text_value = widget[0]
output_data.append(text_value)
return output_data
def format_metadata(self, prompt_info):
metadata_parts = []
# Models
if prompt_info.models:
model_text = '\n'.join([model.name for model in prompt_info.models])
metadata_parts.append(f"Models:\n{model_text}")
# Samplers
if prompt_info.samplers:
sampler_names = '\n'.join([sampler.name for sampler in prompt_info.samplers])
metadata_parts.append(f"Samplers:\n{sampler_names}")
for sampler in prompt_info.samplers:
for param, value in sampler.parameters.items():
metadata_parts.append(f"{param.title()}: {value}")
# Prompts
if prompt_info.prompts:
prompt_text = '\n'.join([prompt.value for prompt in prompt_info.prompts])
metadata_parts.append(f"Prompts:\n{prompt_text}")
# Negative Prompts
if prompt_info.negative_prompts:
negative_prompt_text = '\n'.join([prompt.value for prompt in prompt_info.negative_prompts])
metadata_parts.append(f"Negative Prompts:\n{negative_prompt_text}")
# Additional metadata
if prompt_info.metadata:
additional_metadata = '\n'.join([f"{k}: {v}" for k, v in prompt_info.metadata.items()])
metadata_parts.append(f"Additional Metadata:\n{additional_metadata}")
# Unmodified parameters
if prompt_info.parameters:
params_text = '\n'.join([f"{k}: {v}" for k, v in prompt_info.parameters.items()])
metadata_parts.append(f"Parameters:\n{params_text}")
return '\n\n'.join(metadata_parts)
def format_raw_metadata(self, prompt_info):
return str(prompt_info)
def get_float(self, parameters, key):
try:
value = str(parameters.get(key, '0.0')).strip()
return float(value)
except (ValueError, TypeError):
return 0.0
def get_int(self, parameters, key):
try:
value = str(parameters.get(key, '0')).strip()
return int(value)
except (ValueError, TypeError):
return 0
def get_prompt_text(self, prompts):
if prompts:
return '\n'.join([prompt.value for prompt in prompts])
return ""
def get_model_name(self, prompt_info):
if isinstance(prompt_info.models, list) and len(prompt_info.models) > 0:
return prompt_info.models[0].name
return ""
NODE_CLASS_MAPPINGS = {
"MetadataExtractorBySeed": MetadataExtractorBySeed