diff --git a/py/metadata_collector/metadata_registry.py b/py/metadata_collector/metadata_registry.py index e287c5b1..6e33e806 100644 --- a/py/metadata_collector/metadata_registry.py +++ b/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 diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index cca73095..64dda557 100644 --- a/py/metadata_collector/node_extractors.py +++ b/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 - 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 = {