From b61118ffc8d03a9e586a066fc492c1aa2f0a7c18 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Thu, 17 Jul 2025 09:02:23 -0400
Subject: [PATCH] detect comfyui metadata

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 CHANGELOG.md      |  2 ++
 modules/images.py | 25 +++++++++++++++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 11ea88bfb..06e861fb3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -95,6 +95,8 @@ Although upgrades and existing installations are tested and should work fine!
     while any other units will stay as-is and can be used as normal  
     when using controlnet/processor selected in xyz grid, behavior is forced as control-only  
     also freely selectable are control strength, start and end values  
+  - **Metadata** improved parsing and detect foreign metadata  
+    detect ComfyUI images  
 - **API**  
   - add `/sdapi/v1/lock-checkpoint` endpoint that can be used to lock/unlock model changes  
     if model is locked, it cannot be changed using normal load or unload methods  
diff --git a/modules/images.py b/modules/images.py
index d41270db0..9c9a75b1f 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -232,11 +232,24 @@ def safe_decode_string(s: bytes):
     return None
 
 
+def parse_comfy_metadata(data: str):
+    res = ''
+    try:
+        dct = json.loads(data)
+        version = dct['extra'].get('frontendVersion', None) if 'extra' in dct else ''
+        nodes = dct.get('nodes', []) if 'nodes' in dct else []
+        if 'ComfyUI' in data:
+            res = f"App: ComfyUI | Version: {version} | Nodes: {len(nodes)}"
+    except Exception as e:
+        shared.log.error(f'Error parsing ComfyUI metadata: {e}')
+    return res
+
+
 def read_info_from_image(image: Image, watermark: bool = False):
     if image is None:
         return '', {}
     items = image.info or {}
-    geninfo = items.pop('parameters', None) or items.pop('UserComment', None)
+    geninfo = items.pop('parameters', None) or items.pop('UserComment', None) or ''
     if geninfo is not None and len(geninfo) > 0:
         if 'UserComment' in geninfo:
             geninfo = geninfo['UserComment']
@@ -273,9 +286,10 @@ def read_info_from_image(image: Image, watermark: bool = False):
         if isinstance(val, bytes): # decode bytestring
             items[key] = safe_decode_string(val)
 
-    for key in ['exif', 'ExifOffset', 'JpegIFOffset', 'JpegIFByteCount', 'ExifVersion', 'icc_profile', 'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'adobe', 'photoshop', 'loop', 'duration', 'dpi']: # remove unwanted tags
-        items.pop(key, None)
-
+    if "prompt" in items:
+        geninfo += parse_comfy_metadata(items["prompt"])
+    if "workflow" in items:
+        geninfo += parse_comfy_metadata(items["workflow"])
     if items.get("Software", None) == "NovelAI":
         try:
             json_info = json.loads(items["Comment"])
@@ -286,6 +300,9 @@ Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}
         except Exception as e:
             errors.display(e, 'novelai image parser')
 
+    for key in ['exif', 'ExifOffset', 'JpegIFOffset', 'JpegIFByteCount', 'ExifVersion', 'icc_profile', 'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'adobe', 'photoshop', 'loop', 'duration', 'dpi']: # remove unwanted tags
+        items.pop(key, None)
+
     try:
         items['width'] = image.width
         items['height'] = image.height