ComfyUI-EulerDiscreteScheduler/extract_metadata_node.py at master · econotintas/ComfyUI-EulerDiscreteScheduler · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
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"
}