diff --git a/nodes/common/metadata_parser.py b/nodes/common/metadata_parser.py index 23eae3a..39aa357 100644 --- a/nodes/common/metadata_parser.py +++ b/nodes/common/metadata_parser.py @@ -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 diff --git a/nodes/image_metadata_extractor.py b/nodes/image_metadata_extractor.py index 895f536..e8d80dd 100644 --- a/nodes/image_metadata_extractor.py +++ b/nodes/image_metadata_extractor.py @@ -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 diff --git a/nodes/metadata_extract_by_seed.py b/nodes/metadata_extract_by_seed.py index 1b53c98..90fe71f 100644 --- a/nodes/metadata_extract_by_seed.py +++ b/nodes/metadata_extract_by_seed.py @@ -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