alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-21 21:22:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat: Enhance VAEDecodeExtractor to improve image caching and metadata handling

Browse Source
This commit is contained in:
Will Miao
2025-04-18 10:03:26 +08:00
parent 91b4827c1d
commit 0734252e98
2 changed files with 29 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
py/metadata_collector/metadata_registry.py
View File

@@ -249,4 +249,27 @@ class MetadataRegistry:
# If it's a tensor, return as is for processing in the route handler
return image_data
# If no image is found in the current metadata, try to find it in the cache
# This handles the case where VAEDecode was cached by ComfyUI and not executed
prompt_obj = metadata.get("current_prompt")
if prompt_obj and hasattr(prompt_obj, "original_prompt"):
original_prompt = prompt_obj.original_prompt
for node_id, node_data in original_prompt.items():
class_type = node_data.get("class_type")
if class_type and class_type in NODE_CLASS_MAPPINGS:
class_obj = NODE_CLASS_MAPPINGS[class_type]
class_name = class_obj.__name__
# Check if this is a VAEDecode node
if class_name == "VAEDecode":
# Try to find this node in the cache
cache_key = f"{node_id}:{class_name}"
if cache_key in self.node_cache:
cached_data = self.node_cache[cache_key]
if IMAGES in cached_data and node_id in cached_data[IMAGES]:
image_data = cached_data[IMAGES][node_id]["image"]
# Handle different image formats
if isinstance(image_data, (list, tuple)) and len(image_data) > 0:
return image_data[0]
return image_data
return None

12
py/metadata_collector/node_extractors.py
View File

@@ -243,19 +243,19 @@ class VAEDecodeExtractor(NodeMetadataExtractor):
@staticmethod
def update(node_id, outputs, metadata):
# Check if we already have a first VAEDecode result
if IMAGES in metadata and "first_decode" in metadata[IMAGES]:
return
# Ensure IMAGES category exists
if IMAGES not in metadata:
metadata[IMAGES] = {}
# Save reference to the first VAEDecode result
metadata[IMAGES]["first_decode"] = {
# Save image data under node ID index to be captured by caching mechanism
metadata[IMAGES][node_id] = {
"node_id": node_id,
"image": outputs
}
# Only set first_decode if it hasn't been recorded yet
if "first_decode" not in metadata[IMAGES]:
metadata[IMAGES]["first_decode"] = metadata[IMAGES][node_id]
# Registry of node-specific extractors
NODE_EXTRACTORS = {