107 lines
4.6 KiB
Python
107 lines
4.6 KiB
Python
import torch
|
|
import os
|
|
import json
|
|
from PIL import Image, ImageOps
|
|
import folder_paths
|
|
import numpy as np
|
|
|
|
class ImageMetadataExtractor:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
return {"required":
|
|
{"image": (sorted(files), {"image_upload": True})},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "positive_prompt", "width", "height", "filename")
|
|
FUNCTION = "extract_metadata"
|
|
CATEGORY = "utils"
|
|
|
|
def extract_metadata(self, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
img = Image.open(image_path)
|
|
|
|
output_image = ImageOps.exif_transpose(img)
|
|
output_image = output_image.convert("RGB")
|
|
output_image = np.array(output_image).astype(np.float32) / 255.0
|
|
output_image = torch.from_numpy(output_image)[None,]
|
|
|
|
positive_prompt = ""
|
|
width = 0
|
|
height = 0
|
|
|
|
# Extract from 'prompt' (API format) which is what ComfyUI uses for execution
|
|
if 'prompt' in img.info:
|
|
try:
|
|
prompt = json.loads(img.info['prompt'])
|
|
|
|
# 1. Find Positive Prompt
|
|
# Strategy: Find KSampler -> positive input -> CLIPTextEncode -> text
|
|
ksampler_nodes = []
|
|
for node_id, node in prompt.items():
|
|
class_type = node.get('class_type', '')
|
|
if 'KSampler' in class_type or 'SamplerCustom' in class_type:
|
|
ksampler_nodes.append(node)
|
|
|
|
for ksampler in ksampler_nodes:
|
|
inputs = ksampler.get('inputs', {})
|
|
if 'positive' in inputs:
|
|
positive_link = inputs['positive']
|
|
if isinstance(positive_link, list): # It's a link [node_id, slot_index]
|
|
positive_node_id = str(positive_link[0])
|
|
if positive_node_id in prompt:
|
|
positive_node = prompt[positive_node_id]
|
|
if positive_node.get('class_type') == 'CLIPTextEncode':
|
|
positive_prompt = positive_node.get('inputs', {}).get('text', "")
|
|
break # Found it
|
|
|
|
# Fallback: Look for any CLIPTextEncode with "positive" in title/meta if not found via KSampler
|
|
if not positive_prompt:
|
|
candidates = []
|
|
for node_id, node in prompt.items():
|
|
if node.get('class_type') == 'CLIPTextEncode':
|
|
title = node.get('_meta', {}).get('title', '').lower()
|
|
text = node.get('inputs', {}).get('text', "")
|
|
if 'positive' in title and 'negative' not in title:
|
|
candidates.append(text)
|
|
# Also consider just long text if no clear title match
|
|
elif len(text) > 50:
|
|
candidates.append(text)
|
|
|
|
# Pick the longest candidate if any found
|
|
if candidates:
|
|
positive_prompt = max(candidates, key=len)
|
|
|
|
# 2. Find Width/Height
|
|
# Strategy: Find EmptyLatentImage
|
|
for node_id, node in prompt.items():
|
|
if node.get('class_type') == 'EmptyLatentImage':
|
|
width = node.get('inputs', {}).get('width', 0)
|
|
height = node.get('inputs', {}).get('height', 0)
|
|
break
|
|
|
|
# Fallback: Look for width/height in any node if still 0
|
|
if width == 0 or height == 0:
|
|
for node_id, node in prompt.items():
|
|
inputs = node.get('inputs', {})
|
|
if 'width' in inputs and 'height' in inputs:
|
|
width = inputs['width']
|
|
height = inputs['height']
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"Error parsing metadata: {e}")
|
|
|
|
return (output_image, positive_prompt, width, height, image)
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageMetadataExtractor": ImageMetadataExtractor
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageMetadataExtractor": "Load Image ErosDiffusion"
|
|
}
|