feat: Enhance VAEDecodeExtractor to improve image caching and metadata handling

py/metadata_collector/metadata_registry.py

@@ -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

py/metadata_collector/node_extractors.py

@@ -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
+        # Save image data under node ID index to be captured by caching mechanism
-        metadata[IMAGES]["first_decode"] = {
+        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 = {