NAI-Tag-Viewer/comfyui_parser.py at main · namemechan/NAI-Tag-Viewer · 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import json
import ast


def _sanitize_info_field(raw):
    """
    img.info에서 꺼낸 필드를 안전한 JSON 문자열로 정제.
    WebP는 EXIF 경유로 bytes나 null bytes / garbage prefix가 붙어올 수 있음.
    """
    if raw is None:
        return "{}"
    if isinstance(raw, bytes):
        raw = raw.replace(b'\x00', b'').decode('utf-8', errors='ignore').strip()
    if isinstance(raw, str):
        raw = raw.replace('\x00', '').strip()
        idx = raw.find('{')
        if idx > 0:
            raw = raw[idx:]
    return raw or "{}"


def _parse_metadata_object(raw):
    sanitized = _sanitize_info_field(raw)
    try:
        return json.loads(sanitized)
    except Exception:
        try:
            value = ast.literal_eval(sanitized)
            if isinstance(value, (dict, list)):
                return value
        except Exception:
            pass
        raise


def get_comfyui_data(info, width, height):
    try:
        prompt_raw = info.get("prompt", "{}")

        # 이미 dict이면 그대로 사용 (PNG에서 직접 파싱된 경우)
        if isinstance(prompt_raw, dict):
            prompt_json = prompt_raw
        else:
            prompt_json = _parse_metadata_object(prompt_raw)

        positive_prompt = ""
        negative_prompt = ""
        options = {}

        for node in prompt_json.values():
            if not isinstance(node, dict):
                continue
            class_type = node.get("class_type", "")
            inputs = node.get("inputs", {})

            if class_type == "CLIPTextEncode":
                text = inputs.get("text", "")
                # list이면 노드 참조(연결)이므로 스킵
                if not isinstance(text, str):
                    continue
                if "negative" in text.lower() or "embedding:" in text.lower():
                    negative_prompt += text + "\n"
                else:
                    positive_prompt += text + "\n"

            elif "Sampler" in class_type:
                for param_key, param_value in inputs.items():
                    if param_key not in ["positive", "negative", "model", "latent_image",
                                         "control_net", "clip", "vae"]:
                        # 값이 노드 참조(list)이면 스킵
                        if not isinstance(param_value, list):
                            options[param_key] = param_value

        # prompt/negative 둘 다 비어있으면 ComfyUI 이미지가 아닌 것으로 간주
        if not positive_prompt and not negative_prompt:
            return None

        workflow_raw = info.get("workflow", "{}")
        if isinstance(workflow_raw, dict):
            workflow = workflow_raw
        else:
            try:
                workflow = _parse_metadata_object(workflow_raw)
            except Exception:
                workflow = {}

        # info의 모든 키를 etc에 포함 (prompt, workflow 제외한 나머지도 전부)
        # memo 등 커스텀 키가 있어도 누락되지 않도록
        extra_info = {}
        skip_keys = {"prompt", "workflow"}
        for k, v in info.items():
            if k in skip_keys:
                continue
            if isinstance(v, (bytes, bytearray)):
                try:
                    decoded = v.replace(b'\x00', b'').decode('utf-8', errors='ignore').strip()
                    try:
                        extra_info[k] = _parse_metadata_object(decoded)
                    except Exception:
                        extra_info[k] = decoded
                except Exception:
                    extra_info[k] = repr(v)
            elif isinstance(v, str):
                try:
                    extra_info[k] = _parse_metadata_object(v)
                except Exception:
                    extra_info[k] = v
            else:
                extra_info[k] = v

        return {
            "prompt": positive_prompt.strip(),
            "negative_prompt": negative_prompt.strip(),
            "option": options,
            "etc": {
                "prompt_nodes": prompt_json,
                "workflow": workflow,
                **extra_info
            }
        }
    except Exception as e:
        print(f"Error parsing ComfyUI data: {e}")
        return None