From 7aa4ffdef2c26d573dd76e43ae91cbbc1e4bb988 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Wed, 25 Mar 2026 15:18:46 -0300 Subject: [PATCH] feat: enhance metadata handling with generation type and source image extraction --- .../metahub_save_video_node.cpython-310.pyc | Bin 8925 -> 8954 bytes .../video_metadata_utils.cpython-310.pyc | Bin 12307 -> 12336 bytes metadata_utils.py | 2 + metahub_save_node.py | 2 + tests/test_metadata_utils.py | 4 + tests/test_workflow_extractor.py | 56 +++++++++++++ workflow_extractor.py | 76 ++++++++++++++++++ 7 files changed, 140 insertions(+) diff --git a/__pycache__/metahub_save_video_node.cpython-310.pyc b/__pycache__/metahub_save_video_node.cpython-310.pyc index 8382beb64ad03ef77c5d6fb3ce35b8b79bf4282d..14ca55479bf0de111e09868d8af5fe37c3d9c971 100644 GIT binary patch delta 1785 zcmc(gO>7%Q6o6|iFSa*!()g!M8YfQFHg!-)B|=iAC8=vWuG_|T`~yvDX)}1o&XWDH zX4h7ei!1&?AW-Lq)B_Tr_6Rvw<$%PU14ttwB;bMs66YK$=B+K;1Po##glx9i~<7m;MuQ7Kxa-((K*mv zZoJ`6b)058PM^513HGD_{n>p*oB(TAk9jJmnUF8!Hu5FrOWiIN;!-}F&Sckfj6U$J zB$?Z2Sc;(|l?l6wEMxuAiI{^qtEx#BArRnbIk{5FXsC`iu!XKSU!6}FT5Z3anJ=Oi zt{gt~JaCB9;OzjMKm0;`;&MOFsniox_iL2xke=!$ss+f>5kcCbRGR_rp9&iQ_0Zejqr~bJ{rnttu z)m>cOYZ!`d@x6OBL&mCP?Kd#6}3_r~8Lp<#L@C7c$!n2&4;Sgs%92ts35En6;y3FXu;UG7E9$vlE+oZiW zdsCgw{CyvG`XmY~_EH+m0tV&)P*<`5xJXw|jkUM9$ixOQw1y?g$U@A6?^*;^Eraw# zEHl5Rn5LrFrJ90O*(574&w{({)@&1Gy2bekD7*p z$x<{-%G_{{ybNQn0B!>|0gF8>fpqEc;I!A#wugqn&@V~>ROqx$QS62OLuKIf?1YZj z2y!}}l3gN!96HC396}Tv&Qj;Z+zb76?0ovEz&AncMD4b>lY;gplXPcdFA42x*Y965 zcoRn70`zlv8`O6I?{dVPWRng@z7214W!qoJmef*YY)Fb;Gsq78J@T4x2b@(}nT)hG z5dXikK{bnXG(0KCq@ga#r=Gs$~kd7oo@sD+7AQ>whl4E=aN?9rQ3-;EDk znJ~@1Tsol3(WGk`0@HiZGi%JfUd|>;nOv5MxpFCAE=k#3GR549cqx_66}Fg|OcfJ_ zOdfp5C5ZM60Imdbow8^ov4h;n0}t3L0QjpSP$x!T qoSEZy%Fax1o{hu$QZiG_-->TZ+4x3E+KA`#nQXfF&qY3+x$+k}>8}+4 delta 1827 zcmc(fO>7%Q6o6}j7u#DqX&k4CV<#r0q)w^Q5PnN&aBRmoB(~)rpp(Kf?2hei_s5!D zTTLb83nwHl%mE20H!eLwE?oE#2yy3vWQhx$xFK;sia0QDY}roWz>Ssi<^9aOdGp>I zpJso~hU{Q4aDjhcu!jErOP__xk_Ry?l{i0_zLcl9i(Uduvj>lTGX;8Y!qhDee4iWl zeA)h@M1RpIz8y(ynr6TFw&W;S2YS+fXAViFa`B*4MZxTw)pANH6>|B)ZV|C#|IJn8 zvsFetu{844mR2Zv5o13JPbtF*`4W6J)6uEzBcs?1R8g3g^zO2fpJqr>Uryi zja~Wzm|g_j1dQ$AYH!lAsfi8!0a3*AUt!J5?62@%H_XM~WW$B=mVW%}edJrJ=BD8wOEbrO|3DCUL2*I*5Hcd+QJ? zB#I+yB8t&3!rA+aYtyq=8P6V~Tm`^9&^5p{c6fQdyT?T&x2SElTt&k!Mt*YNrC4=s zIuwb9LZK19>F+Ip%g{?b7cLx(?w4G)kzpTiw zuJebahz9I(!xfqX->ZN;Un2Y}XYX%B< z-vNwr!JxhiP&twlbca0^|6%quuAoqZxJpOYh^^>W-KHhxU%V^r!{XQ3(PF&&28jP( z?J6*pWBDOS=MeBV;4O}x)WQ^-4%NUyIMBrop?Y0c`Hfapv3wiJnQSFp&X-^f-GD%k z0N@vb-eP|y;*lDNqp=jyI!iwOQ{XOm_$B#af#vmik%k{^2Y>)T^of<1*LF}CPFB+U lsY*pTNR>+YLau_M!kEcdO8cqXN+ETS9s2%x=2y#Ge*?Rhq&WZp diff --git a/__pycache__/video_metadata_utils.cpython-310.pyc b/__pycache__/video_metadata_utils.cpython-310.pyc index ddf9644529b4fca84f97885b5661bf6748a071d3..c7c2e2b03faff0f16a71f6a75419ef212988e284 100644 GIT binary patch delta 1633 zcmZ{kO>A355P|HN_Zs4e+vTADwR)J+}dkF*U^f+z`us4bARl|seq#$MZxygz22 z?I!%R1h{bEvPcyIq)3#601dRJLPZ?f3L+$q9FX?JwHy%Q1d))Kc~094yu!H_2Rmc!FvsEx_ zWIPxr&7f%rFR8&(IN1om49Nnv3%ki`4`D_wH@$dql!iwTcLN9&&2!w4c@C(^&H?7@ zEuJ-)?k%}og-S-2JJpDW$sl3#JoAhae-<3e_53g)MNWGO2MK+I!wA(NnB|%rbemZm zk~KAFYv2$qL`9~+vqfA|VcP+VnJQ#CS#H7*8IO=oZHaMJhw!acB+G5b5Im1-`28^f z{rH14VNh;1H}6+5KUzJ5jxs7(7#bZ(M*PFVsQk6Lx!pGyL<1Gh(WagX`MU4~IXNnu zT1vf(G^DI01JFgPLAtSO9v>w0guLD|g?IXQODblQg+p-VU~6giWf}rO5IjPWP$HBO z$&d<>SD_+?RnT2QRXK)tPG~gGz3tzQPbEZ>P zF@IXs1@v!E3~#r z_>vT4s^`qy`!xN4@FC$N0_{KvWm%!Z!pmat5Tk`v1Xh3_XcGxOCVYax)?prlt2AGe zzw`+0MrEdVC2;Xml57w@BarLmcr|$JR?aK|pQCIJzM$p%M^KTS)%kkAqWtcGdDKx2 zpyvR~U&l7%Kuxr*rQt08=cD#Qbxl-LP?w!3qL$&<)!L>~z2}fRC>eer->GZVfSUvZ zu{EkqQadG(1&n*dDw_(o-a@Kf}x-p9FBgPa<4HX)^I& znS`(9U|(DN5wiV;Y-{bJrE0NRtEK1%wk6A zA+Jb}#q+)w6CKbK1r3A_!USQJ;D-~2+c>`0^e*!XVm_KPOpdRp@Kha1=!u10WU?l& z_P>zY!eKISaw@rBMRYyuWOUuXV~CPYoZuJNdpJ_^(e$xnej?jx8eDK3vl_{Bq&P{q SL|7->AlxF{A#BO@wD~U*_0sPE delta 1645 zcmZ{kO>h)N6o8%m%_jS!!H^J<{7@uWLN*}@QAh^UP)>CE)*PGxpx>-Fp| zS$`rw{)>uDm4an?=BBk?^k|iqRUAA!rCvPcVsW%M`QA(*9&`@hbiddA-s|^X_t#5m zDeGo3?f2n#Se2~U6W?V2(sslOylypPg^F7zR;P(8^=7=#5X+TTZ{)Rz*4_Fxw%VpK zRmceBwaoX(=d#~*vYhP#*UMa%h18`7O!ej->R}y*>lCe$Zd;CkN6;ZWieT#5+<9%b z^;51po_La^PpLm2d}D^iE$MK_kC~R?!Q-g#W40x|yaoqxax(&l@W&!f6*_SMTL_2L zXI$1j zMzFNxm8va3cX&m>Fv&+qo{~mcm^IX}gyTYm+bmvjO<}_Xt(l}X(K;8b4WT=0xSJ(S z;kcF*B~;>fdk*&C55@@N2tQr7OyTN9vgi3xECK5XMFyHSSgbrbGo4Qc#{@dCf4^D7 zLt3t*OM#g-WoR<6me{Zp$IU5Wzyb2JPfd2$hhL;2WiTIsLHw~8RWHoxW@6{nFWvK4 z->p4k9WT+2mkCTQ^wdwCq2VipvxIX5ov=jUhYt5Vjh;Ib5VxV`FGhJ@p| zg0-U%OQ7?rm5(viy{&I%j^r81i$ZU>uu`<$2I~l>;aWIF_U{!t4EdO{xBaXct$n?> z;yIQUj;w*uxil;btKUDNom2<=i&Io&KUs5D9Jk>#6Q1SS;=uZTl-`W8tih`uZF6g0 zwdVOTPn`Fd`nCT@%)nOzv)X?3$G`}VlYu_Mp#6kFJ6*xLEYXprtMWbh91scAl(Oq4-rfq9Xh?( zpy?{1NqCh&9(+tQH>IDDo{7Oj2Ps}hV5bCuc9Gys!bJpLF79!7pXQg;Kx&Ih4lij} zTjk-4k+ru-`ZnPm0-0}S!UN!)7Q7E84Q=M!MmKME@m+&-{ql)T`Q{yfYZS^0h~|_!JG4{X z+TFAAf4PE>Xz|AcN|PU#RolX=W;q@#lk^Bl+fQ@5CY}axAxI=#L1{kwU)h9D)WT?Q zA3Z$qDS}0vT7^aNriV=$1Q(|H#^@F8dTV6Q-|k%XRKc!2hWj(;y48T%5a`N+b2e!HL4 zb<;I;J*Y0EXkCU7w3fsXQzyqCn+;Ofhf_bc?7DUs%`>!cmS7Mr6Rr}j6TTpjC+G|s J>$ma6e*n?~vULCe diff --git a/metadata_utils.py b/metadata_utils.py index 2e933d9..4b6b739 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -396,6 +396,8 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict: "model_hash": params['model_hash'], "vae": params.get('vae_name', ''), "denoise": params.get('denoise', 1.0), + "generation_type": params.get('generation_type'), + "source_image": _sanitize(params.get('source_image')), "width": params['width'], "height": params['height'], diff --git a/metahub_save_node.py b/metahub_save_node.py index f3cc327..d1027a0 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -309,6 +309,8 @@ class MetaHubSaveNode: "project_name": project_name, "lora_list": lora_list, "lora_hashes": lora_hashes, + "generation_type": extracted.get("generation_type"), + "source_image": extracted.get("source_image"), # Performance metrics (Tier 1, 2, 3) "vram_peak_mb": vram_peak_mb if vram_peak_mb is not None else gpu_metrics.get("vram_peak_mb"), "gpu_device": gpu_device_override if gpu_device_override else gpu_metrics.get("gpu_device"), diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 3c559db..6218bf2 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -146,6 +146,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): "width": 512, "height": 768, "lora_list": [], + "generation_type": "img2img", + "source_image": {"fileName": "base.png", "relativePath": "inputs/base.png"}, } workflow_json = {"workflow": {"nodes": []}, "prompt": {"1": {"class_type": "KSampler"}}} @@ -154,6 +156,8 @@ def test_build_imh_metadata_includes_workflow_and_prompt(): assert imh["generator"] == "ComfyUI" assert imh["workflow"] == workflow_json["workflow"] assert imh["prompt_api"] == workflow_json["prompt"] + assert imh["generation_type"] == "img2img" + assert imh["source_image"]["fileName"] == "base.png" def test_build_imh_metadata_sanitizes_nan(): diff --git a/tests/test_workflow_extractor.py b/tests/test_workflow_extractor.py index 9a5470b..51e0827 100644 --- a/tests/test_workflow_extractor.py +++ b/tests/test_workflow_extractor.py @@ -68,3 +68,59 @@ def test_workflow_extractor_lora_detection(): assert "loras" not in missing assert data["lora_list"] == [{"name": "detail.safetensors", "weight": 0.8}] + + + +def test_workflow_extractor_detects_img2img_lineage(): + prompt = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "inputs/base.png"}, + }, + "2": { + "class_type": "VAEEncode", + "inputs": {"pixels": ["1", 0], "vae": ["5", 2]}, + }, + "3": { + "class_type": "KSampler", + "inputs": { + "seed": 123, + "steps": 20, + "cfg": 7.5, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 0.45, + "model": ["5", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["2", 0], + }, + }, + "4": { + "class_type": "VAEDecode", + "inputs": {"samples": ["3", 0], "vae": ["5", 2]}, + }, + "5": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "model.safetensors"}, + }, + "6": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "positive", "clip": ["5", 1]}, + }, + "7": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "negative", "clip": ["5", 1]}, + }, + "8": { + "class_type": "MetaHubSaveNode", + "inputs": {"images": ["4", 0]}, + }, + } + + extractor = WorkflowExtractor(prompt) + data, _missing = extractor.extract(save_node_id="8") + + assert data["generation_type"] == "img2img" + assert data["source_image"]["fileName"] == "base.png" + assert data["source_image"]["relativePath"] == "inputs/base.png" diff --git a/workflow_extractor.py b/workflow_extractor.py index 26d66f2..bbd12e8 100644 --- a/workflow_extractor.py +++ b/workflow_extractor.py @@ -54,6 +54,11 @@ class WorkflowExtractor: "VAEDecode (Tiled)", ] + LOAD_IMAGE_NODES = [ + "LoadImage", + "LoadImageMask", + ] + def __init__(self, prompt: Optional[Dict[str, Any]]): self.prompt: Dict[str, Any] = {} if isinstance(prompt, dict): @@ -116,6 +121,10 @@ class WorkflowExtractor: elif has_lora_nodes: missing.add("loras") + lineage = self.extract_lineage(sampler_node_id) + if lineage: + data.update(lineage) + return data, missing def find_sampler_for_save_node(self, save_node_id: Optional[str]) -> Optional[str]: @@ -236,6 +245,73 @@ class WorkflowExtractor: return loras, has_lora_nodes + def extract_lineage(self, sampler_node_id: Optional[str]) -> Dict[str, Any]: + if not sampler_node_id: + return {} + + node = self._get_node(sampler_node_id) + if not node: + return {} + + inputs = node.get("inputs", {}) + latent_conn = inputs.get("latent_image") or inputs.get("latent") or inputs.get("samples") + start_node_id = self._get_connection_node_id(latent_conn) + if not start_node_id: + return {} + + queue = [start_node_id] + visited: Set[str] = set() + source_image: Optional[Dict[str, Any]] = None + has_load_image = False + has_inpaint = False + has_outpaint = False + + while queue: + current_id = queue.pop(0) + if current_id in visited: + continue + visited.add(current_id) + + current = self._get_node(current_id) + if not current: + continue + + class_type = self._class_type(current) + class_lower = class_type.lower() + + if class_type in self.LOAD_IMAGE_NODES or class_lower in {"loadimage", "loadimagemask"}: + has_load_image = True + if not source_image: + image_value = current.get("inputs", {}).get("image") + if isinstance(image_value, str) and image_value.strip(): + normalized = image_value.replace("\\", "/") + source_image = { + "fileName": normalized.split("/")[-1], + "relativePath": normalized, + "nodeId": current_id, + "nodeType": class_type, + } + + if "outpaint" in class_lower or class_lower in {"padforoutpaint", "imagepadforoutpaint"}: + has_outpaint = True + + if "inpaint" in class_lower or class_lower == "setlatentnoisemask" or "mask" in class_lower: + has_inpaint = True + + for input_value in current.get("inputs", {}).values(): + upstream_id = self._get_connection_node_id(input_value) + if upstream_id and upstream_id not in visited: + queue.append(upstream_id) + + if not has_load_image: + return {} + + generation_type = "outpaint" if has_outpaint else "inpaint" if has_inpaint else "img2img" + return { + "generation_type": generation_type, + "source_image": source_image, + } + def _find_sampler_from_images_source(self, start_node_id: str) -> Optional[str]: start_node = self._get_node(start_node_id) if not start_node: