app.py

import gradio as gr
from PIL import Image
from PIL.ExifTags import TAGS
import re
import json
import tempfile
import os
import math


def get_aspect_ratio(w, h):
    """Return the closest standard aspect ratio string for given dimensions."""
    if w == 0 or h == 0:
        return "N/A"
    ratio = w / h
    tolerance = 0.02
    ratios = {
        "1:1": 1.0,
        "5:4": 5 / 4,
        "4:3": 4 / 3,
        "3:2": 3 / 2,
        "16:9": 16 / 9,
        "16:10": 16 / 10,
        "21:9": 21 / 9,
        "2:3": 2 / 3,
        "3:4": 3 / 4,
        "4:5": 4 / 5,
        "9:16": 9 / 16,
    }
    for label, val in ratios.items():
        if abs(ratio - val) < tolerance:
            return label
    gcd = math.gcd(w, h)
    return f"{w // gcd}:{h // gcd}"


def get_image_metadata(img):
    """Extract metadata based on image format."""
    metadata = {}
    if img.format == "PNG":
        metadata = img.info
    elif img.format in ["JPEG", "TIFF", "WEBP"]:
        exif_data = img._getexif()
        if exif_data:
            for tag, value in exif_data.items():
                tag_name = TAGS.get(tag, tag)
                metadata[tag_name] = value
    return metadata


def parse_comfy_prompt(prompt_data):
    """Parse ComfyUI prompt JSON using graph tracing for correct pos/neg identification."""
    positive_prompt = "N/A"
    negative_prompt = "N/A"
    seed = "N/A"

    try:
        if isinstance(prompt_data, str):
            prompt_data = json.loads(prompt_data)
    except (json.JSONDecodeError, TypeError):
        return positive_prompt, negative_prompt, seed, []

    if not isinstance(prompt_data, dict):
        return positive_prompt, negative_prompt, seed, []

    clip_text_nodes = {}
    lora_names = []

    for node_id, node in prompt_data.items():
        if not isinstance(node, dict):
            continue
        class_type = node.get("class_type", "")
        inputs = node.get("inputs", {})
        widgets = node.get("widgets_values", [])

        if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL"):
            text = inputs.get("text", "")
            if isinstance(text, list):
                text = ""
            if not text or not text.strip():
                continue
            clip_text_nodes[node_id] = {
                "text": text.strip(),
                "links": node.get("inputs", []),
            }

        if "LoraLoader" in class_type or "LoraLoaderModelOnly" in class_type:
            lora_name = inputs.get("lora_name", inputs.get("lora", ""))
            if isinstance(lora_name, str) and lora_name.strip():
                lora_names.append(lora_name.strip())
            for wv in widgets:
                if isinstance(wv, str) and wv.endswith(
                    (".safetensors", ".ckpt", ".pt")
                ):
                    if wv not in lora_names:
                        lora_names.append(wv)

        if class_type in (
            "KSampler",
            "KSamplerAdvanced",
            "SamplerCustom",
            "SamplerCustomAdvanced",
        ):
            if "seed" in inputs:
                seed = str(inputs["seed"])
            if "noise_seed" in inputs:
                seed = str(inputs["noise_seed"])

    positive_ids = set()
    negative_ids = set()

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

        is_sampler = class_type in (
            "KSampler",
            "KSamplerAdvanced",
            "SamplerCustom",
            "SamplerCustomAdvanced",
        )
        if is_sampler:
            for input_name in ("positive", "latent_image"):
                link = inputs.get(input_name)
                if isinstance(link, list) and len(link) >= 1:
                    positive_ids.add(str(link[0]))
            for input_name in ("negative",):
                link = inputs.get(input_name)
                if isinstance(link, list) and len(link) >= 1:
                    negative_ids.add(str(link[0]))

        if class_type in ("SamplerCustom", "SamplerCustomAdvanced"):
            for input_name in ("guider", "sampler", "sigmas"):
                link = inputs.get(input_name)
                if isinstance(link, list) and len(link) >= 1:
                    pass

    def trace_positive(node_id, visited=None):
        if visited is None:
            visited = set()
        if node_id in visited:
            return []
        visited.add(node_id)
        node = prompt_data.get(node_id, {})
        if not isinstance(node, dict):
            return []
        class_type = node.get("class_type", "")
        inputs = node.get("inputs", {})
        texts = []
        if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL"):
            text = inputs.get("text", "")
            if isinstance(text, list):
                text = ""
            if text and text.strip():
                texts.append(text.strip())
        for input_name, link in inputs.items():
            if isinstance(link, list) and len(link) >= 1:
                texts.extend(trace_positive(str(link[0]), visited))
        return texts

    def trace_negative(node_id, visited=None):
        if visited is None:
            visited = set()
        if node_id in visited:
            return []
        visited.add(node_id)
        node = prompt_data.get(node_id, {})
        if not isinstance(node, dict):
            return []
        class_type = node.get("class_type", "")
        inputs = node.get("inputs", {})
        texts = []
        if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL"):
            text = inputs.get("text", "")
            if isinstance(text, list):
                text = ""
            if text and text.strip():
                texts.append(text.strip())
        for input_name, link in inputs.items():
            if isinstance(link, list) and len(link) >= 1:
                texts.extend(trace_negative(str(link[0]), visited))
        return texts

    found_pos = False
    found_neg = False
    for nid in sorted(positive_ids):
        if nid in clip_text_nodes:
            positive_prompt = clip_text_nodes[nid]["text"]
            found_pos = True
            break
    if not found_pos:
        for nid in sorted(positive_ids):
            texts = trace_positive(nid)
            if texts:
                positive_prompt = texts[0]
                found_pos = True
                break

    for nid in sorted(negative_ids):
        if nid in clip_text_nodes:
            negative_prompt = clip_text_nodes[nid]["text"]
            found_neg = True
            break
    if not found_neg:
        for nid in sorted(negative_ids):
            texts = trace_negative(nid)
            if texts:
                negative_prompt = texts[0]
                found_neg = True
                break

    if not found_pos and not found_neg:
        texts = [n["text"] for n in clip_text_nodes.values()]
        if len(texts) >= 1:
            positive_prompt = texts[0]
        if len(texts) >= 2:
            negative_prompt = texts[1]

    return positive_prompt, negative_prompt, seed, lora_names


def extract_metadata(image_file):
    """Extract and parse metadata from the uploaded image."""
    try:
        img = Image.open(image_file)
    except Exception:
        return ("Error: Unable to open image.", *(["N/A"] * 7))

    metadata = get_image_metadata(img)
    width, height = img.size
    aspect_ratio = get_aspect_ratio(width, height)
    dimensions = f"{width} × {height} ({aspect_ratio})"

    metadata_str = "\n".join([f"{key}: {value}" for key, value in metadata.items()])

    prompt = "N/A"
    negative_prompt = "N/A"
    seed_number = "N/A"
    comfy_workflow = "N/A"
    lora_names = []

    comfy_prompt_data = metadata.get("prompt", None)
    if comfy_prompt_data is not None:
        pos, neg, s, loras = parse_comfy_prompt(comfy_prompt_data)
        if pos != "N/A":
            prompt = pos
        if neg != "N/A":
            negative_prompt = neg
        if s != "N/A":
            seed_number = s
        lora_names = loras

    elif "parameters" in metadata:
        params = metadata["parameters"]
        prompt_match = re.search(
            r"(.*?)Negative prompt:", params, re.DOTALL | re.IGNORECASE
        )
        if prompt_match:
            prompt = prompt_match.group(1).strip()
        else:
            prompt_match = re.search(
                r"(.*?)(Steps:|$)", params, re.DOTALL | re.IGNORECASE
            )
            if prompt_match:
                prompt = prompt_match.group(1).strip()

        neg_match = re.search(
            r"Negative prompt:(.*?)(Steps:|$)", params, re.DOTALL | re.IGNORECASE
        )
        if neg_match:
            negative_prompt = neg_match.group(1).strip()

        seed_match = re.search(r"Seed:\s*(\d+)", params, re.IGNORECASE)
        if seed_match:
            seed_number = seed_match.group(1).strip()

        lora_pattern = re.compile(r"<lora:([^:>]+)", re.IGNORECASE)
        lora_names = lora_pattern.findall(params)

    workflow_data = metadata.get("workflow", None)
    if workflow_data is not None:
        if isinstance(workflow_data, str):
            try:
                json.loads(workflow_data)
                comfy_workflow = workflow_data
            except json.JSONDecodeError:
                comfy_workflow = workflow_data
        elif isinstance(workflow_data, dict):
            comfy_workflow = json.dumps(workflow_data, indent=2)
        else:
            comfy_workflow = str(workflow_data)

    return (
        dimensions,
        prompt,
        negative_prompt,
        seed_number,
        metadata_str,
        comfy_workflow,
        lora_names,
        "N/A",
    )


def export_workflow(workflow_text):
    """Convert the workflow text into a downloadable JSON file."""
    if workflow_text == "N/A" or not workflow_text.strip():
        return None, "No workflow data to export."

    workflow_data = {"comfy_workflow": workflow_text}

    try:
        with tempfile.NamedTemporaryFile(
            mode="w", delete=False, suffix=".json"
        ) as tmp_file:
            json.dump(workflow_data, tmp_file, indent=4)
            tmp_file_path = tmp_file.name
    except Exception:
        return None, "Failed to create workflow JSON file."

    if os.path.exists(tmp_file_path):
        return tmp_file_path, "Workflow exported successfully."
    else:
        return None, "Failed to export workflow."


def strip_metadata(image_file):
    """Strip all metadata from an image and return a clean file."""
    if image_file is None:
        return None, "No image provided."

    try:
        img = Image.open(image_file)
    except Exception:
        return None, "Error: Unable to open image."

    clean_img = Image.new(img.mode, img.size)
    clean_img.putdata(list(img.getdata()))

    try:
        with tempfile.NamedTemporaryFile(
            mode="wb", delete=False, suffix=".png"
        ) as tmp_file:
            clean_img.save(tmp_file, format="PNG")
            tmp_file_path = tmp_file.name
    except Exception:
        return None, "Error saving stripped image."

    return tmp_file_path, "Metadata stripped successfully."


def main():
    with gr.Blocks() as iface:
        gr.Markdown("<h1>Comfy / A1111 Metadata Reader</h1>")
        gr.Markdown(
            "<p>Upload an image (PNG, JPEG, WebP) to extract its metadata and parse it for prompts.</p>"
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(label="Drop Image Here", type="filepath")

            with gr.Column(scale=2):
                dimensions_output = gr.Textbox(
                    label="Dimensions", lines=1, interactive=False
                )
                prompt_output = gr.Textbox(
                    label="Prompt", lines=4, show_copy_button=True, interactive=False
                )
                negative_prompt_output = gr.Textbox(
                    label="Negative Prompt",
                    lines=4,
                    show_copy_button=True,
                    interactive=False,
                )
                seed_output = gr.Textbox(
                    label="Seed Number",
                    lines=1,
                    show_copy_button=True,
                    interactive=False,
                )
                lora_output = gr.Textbox(
                    label="LoRA(s)", lines=2, show_copy_button=True, interactive=False
                )
                original_metadata_output = gr.Textbox(
                    label="Original Metadata", lines=15, interactive=False
                )

            with gr.Column(scale=2):
                comfy_workflow_output = gr.Textbox(
                    label="Comfy Workflow", lines=20, value="N/A", interactive=False
                )
                with gr.Row():
                    export_button = gr.Button("Export Workflow as JSON")
                    strip_button = gr.Button("Strip Metadata", variant="stop")
                workflow_file = gr.File(label="Download Workflow JSON", visible=False)
                export_message = gr.Textbox(
                    label="Export Status", lines=1, interactive=False
                )
                strip_file = gr.File(label="Download Stripped Image", visible=False)
                strip_message = gr.Textbox(
                    label="Strip Status", lines=1, interactive=False
                )

        image_input.change(
            fn=extract_metadata,
            inputs=image_input,
            outputs=[
                dimensions_output,
                prompt_output,
                negative_prompt_output,
                seed_output,
                original_metadata_output,
                comfy_workflow_output,
                lora_output,
                strip_message,
            ],
        )

        export_button.click(
            fn=export_workflow,
            inputs=comfy_workflow_output,
            outputs=[workflow_file, export_message],
            queue=False,
        )

        strip_button.click(
            fn=strip_metadata,
            inputs=image_input,
            outputs=[strip_file, strip_message],
            queue=False,
        )

    iface.launch()


if __name__ == "__main__":
    main()