Refactor: rewrite ComfyUI parsing with graph tracing, add LoRA/dimensions/strip metadata, update for Gradio 6.x

.gitignore

@@ -1,2 +1,4 @@
 /venv
+venv/
+__pycache__/
 .DS_Store

README.md

@@ -4,7 +4,7 @@ emoji: 🐨
 colorFrom: blue
 colorTo: blue
 sdk: gradio
-sdk_version: 4.44.0
+sdk_version: 5.34.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -1,18 +1,45 @@
 import gradio as gr
 from PIL import Image
 from PIL.ExifTags import TAGS
-import numpy as np
 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':
+    if img.format == "PNG":
         metadata = img.info
-    elif img.format in ['JPEG', 'TIFF']:
+    elif img.format in ["JPEG", "TIFF", "WEBP"]:
         exif_data = img._getexif()
         if exif_data:
             for tag, value in exif_data.items():
@@ -20,180 +47,384 @@ def get_image_metadata(img):
                 metadata[tag_name] = value
     return metadata
 
-def parse_parameters(parameters):
-    """Parse the 'parameters' field to extract prompts and seed number."""
-    prompt = 'N/A'
-    negative_prompt = 'N/A'
-    adetailer_prompt_1 = 'N/A'
-    adetailer_prompt_2 = 'N/A'
-    seed_number = -1
-
-    # Extract the prompt
-    prompt_match = re.search(r'(.*?)Negative prompt:', parameters, re.DOTALL | re.IGNORECASE)
-    if prompt_match:
-        prompt = prompt_match.group(1).strip()
-    else:
-        prompt_match = re.search(r'(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
-        if prompt_match:
-            prompt = prompt_match.group(1).strip()
 
-    # Extract the negative prompt
-    negative_prompt_match = re.search(r'Negative prompt:(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
-    if negative_prompt_match:
-        negative_prompt = negative_prompt_match.group(1).strip()
+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"
 
-    # Extract ADetailer prompts
-    adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters, re.IGNORECASE)
-    if adetailer_prompt_1_match:
-        adetailer_prompt_1 = adetailer_prompt_1_match.group(1).strip()
+    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", [])
 
-    adetailer_prompt_2_match = re.search(r'ADetailer negative prompt 2nd:\s*"(.*?)"', parameters, re.IGNORECASE)
-    if adetailer_prompt_2_match:
-        adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()
+        if class_type in ("CLIPTextEncode", "CLIPTextEncodeSDXL"):
+            text = inputs.get("text", "")
+            if not text or not text.strip():
+                continue
+            clip_text_nodes[node_id] = {
+                "text": text.strip(),
+                "links": node.get("inputs", []),
+            }
 
-    # Extract Seed Number
-    seed_match = re.search(r'Seed:\s*(\d+)', parameters, re.IGNORECASE)
-    if seed_match:
-        seed_number = seed_match.group(1).strip()
+        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 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 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
 
-    return prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number
 
 def extract_metadata(image_file):
     """Extract and parse metadata from the uploaded image."""
     try:
         img = Image.open(image_file)
-    except Exception as e:
-        # Return placeholders in case of error
-        return (None, 'Error: Unable to open image.', *(['N/A'] * 7))
+    except Exception:
+        return ("Error: Unable to open image.", *(["N/A"] * 6), [], "N/A")
 
     metadata = get_image_metadata(img)
-    img_display = np.array(img)
+    width, height = img.size
+    aspect_ratio = get_aspect_ratio(width, height)
+    dimensions = f"{width} × {height} ({aspect_ratio})"
 
-    # Convert metadata to string
     metadata_str = "\n".join([f"{key}: {value}" for key, value in metadata.items()])
 
-    # Initialize default values
-    prompt = negative_prompt = adetailer_prompt_1 = adetailer_prompt_2 = 'N/A'
-    seed_number = -1
-    comfy_prompt = comfy_workflow = 'N/A'
+    prompt = "N/A"
+    negative_prompt = "N/A"
+    seed_number = "N/A"
+    comfy_workflow = "N/A"
+    lora_names = []
 
-    # Parse parameters if available
-    if 'parameters' in metadata:
-        parsed = parse_parameters(metadata['parameters'])
-        prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number = parsed
+    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
 
-    # Direct extraction with validation for Comfy Prompt
-    if 'prompt' in metadata and isinstance(metadata['prompt'], str):
-        comfy_prompt = metadata['prompt'].strip()
-    else:
-        comfy_prompt = 'N/A'  # Fallback if prompt is missing or not a string
+    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()
 
-    # Direct extraction with validation for Comfy Workflow
-    if 'workflow' in metadata and isinstance(metadata['workflow'], str):
-        comfy_workflow = metadata['workflow'].strip()
-    else:
-        comfy_workflow = 'N/A'  # Fallback if workflow is missing or not a string
+        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 (
-        img_display, 
-        prompt, 
-        negative_prompt, 
-        seed_number, 
-        adetailer_prompt_1, 
-        adetailer_prompt_2, 
-        metadata_str, 
-        comfy_prompt, 
-        comfy_workflow
+        dimensions,
+        prompt,
+        negative_prompt,
+        seed_number,
+        metadata_str,
+        comfy_workflow,
+        lora_names,
+        "N/A",
     )
 
+
 def export_workflow(workflow_text):
-    """
-    Converts the workflow text into JSON format and returns it as a downloadable file.
-    """
-    print("Export workflow function called.")
+    """Convert the workflow text into a downloadable JSON file."""
     if workflow_text == "N/A" or not workflow_text.strip():
-        print("No workflow data to export.")
         return None, "No workflow data to export."
 
-    # Structure the JSON data
-    workflow_data = {
-        "comfy_workflow": workflow_text
-    }
+    workflow_data = {"comfy_workflow": workflow_text}
 
-    # Create a temporary JSON file
     try:
-        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as tmp_file:
+        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
-        print(f"Temporary JSON file created at: {tmp_file_path}")
-    except Exception as e:
-        print(f"Error creating temporary file: {e}")
+    except Exception:
         return None, "Failed to create workflow JSON file."
 
-    # Check if the file was created successfully
     if os.path.exists(tmp_file_path):
-        message = "Workflow exported successfully."
-        print(message)
-        return tmp_file_path, message
+        return tmp_file_path, "Workflow exported successfully."
     else:
-        message = "Failed to export workflow."
-        print(message)
-        return None, message
+        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():
-    # Create Gradio User Interface
     with gr.Blocks() as iface:
-        gr.Markdown("<h1>Automatic 1111/ Comfy Metadata Reader</h1>")
-        gr.Markdown("<p>Upload an image (PNG, JPEG) to extract its metadata and parse it for prompts.</p>")
+        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():  # First column for image upload and preview
-                image_input = gr.File(label="Upload Image", type="filepath")
-                image_output = gr.Image(label="Image Preview")
-                original_metadata_output = gr.Textbox(label="Original Metadata", lines=20, interactive=False)
-
-            with gr.Column():  # Second column for metadata outputs
-                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)
-                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True, interactive=False)
-                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True, interactive=False)
-
-            with gr.Column():  # Third column for additional metadata outputs
-                with gr.Row():
-                    comfy_prompt_output = gr.Textbox(label="Comfy Prompt", lines=4, value="N/A", interactive=False)
+            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, buttons=["copy"], interactive=False
+                )
+                negative_prompt_output = gr.Textbox(
+                    label="Negative Prompt",
+                    lines=4,
+                    buttons=["copy"],
+                    interactive=False,
+                )
+                seed_output = gr.Textbox(
+                    label="Seed Number", lines=1, buttons=["copy"], interactive=False
+                )
+                lora_output = gr.Textbox(
+                    label="LoRA(s)", lines=2, buttons=["copy"], 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():
-                    comfy_workflow_output = gr.Textbox(label="Comfy Workflow", lines=20, value="N/A", interactive=False)
                     export_button = gr.Button("Export Workflow as JSON")
-                    workflow_file = gr.File(label="Download Workflow JSON", visible=False)
-                with gr.Row():
-                    export_message = gr.Textbox(label="Export Status", lines=1, interactive=False, visible=False)
+                    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
+                )
 
-        # Set up the input-output relationships
         image_input.change(
             fn=extract_metadata,
             inputs=image_input,
             outputs=[
-                image_output,
+                dimensions_output,
                 prompt_output,
                 negative_prompt_output,
                 seed_output,
-                adetailer_prompt_1_output,
-                adetailer_prompt_2_output,
                 original_metadata_output,
-                comfy_prompt_output, 
                 comfy_workflow_output,
-            ]
+                lora_output,
+                strip_message,
+            ],
         )
 
-        # Connect the export button to the export function
         export_button.click(
             fn=export_workflow,
             inputs=comfy_workflow_output,
             outputs=[workflow_file, export_message],
-            queue=False
+            queue=False,
+        )
+
+        strip_button.click(
+            fn=strip_metadata,
+            inputs=image_input,
+            outputs=[strip_file, strip_message],
+            queue=False,
         )
 
     iface.launch()
 
+
 if __name__ == "__main__":
-    main()
+    main()

requirements.txt

@@ -1,2 +1,3 @@
 gradio
-Pillow
+Pillow
+numpy