feat(metadata): collect generation params from Krea two/three stage samplers

py/metadata_collector/metadata_processor.py

@@ -214,6 +214,24 @@ class MetadataProcessor:
                     max_denoise = denoise
                     primary_sampler = sampler_info
                     primary_sampler_id = node_id
+
+        # Last resort: any registered sampler. Samplers without a denoise or
+        # add_noise parameter (e.g. multi-stage samplers like KreaTwoStageSampler)
+        # are not caught by the criteria above. Prefer execution order so the
+        # first executed sampler wins, matching the downstream_id branch.
+        if primary_sampler is None:
+            sampler_ids = [
+                node_id
+                for node_id, sampler_info in metadata.get(SAMPLING, {}).items()
+                if sampler_info.get(IS_SAMPLER, False)
+            ]
+            if sampler_ids:
+                if downstream_id and "execution_order" in metadata:
+                    for node_id in metadata["execution_order"]:
+                        if node_id in sampler_ids:
+                            return node_id, metadata[SAMPLING][node_id]
+                primary_sampler_id = sampler_ids[0]
+                primary_sampler = metadata[SAMPLING][sampler_ids[0]]
                 
         return primary_sampler_id, primary_sampler
     

py/metadata_collector/node_extractors.py

@@ -861,6 +861,65 @@ class TSCKSamplerAdvancedExtractor(KSamplerAdvancedExtractor, TSCSamplerBaseExtr
 
     # Update method is inherited from TSCSamplerBaseExtractor
 
+class KreaTwoStageSamplerExtractor(BaseSamplerExtractor):
+    """Extractor for Krea Two/Three Stage Samplers (Auryg/Krea-2-Two-Stage-Sampler).
+
+    The node samples in two (or three) stages with per-stage settings
+    (stage1_steps/stage2_steps, stage1_cfg/stage2_cfg, ...). The canonical
+    metadata fields consumed by ``extract_generation_params`` (steps, cfg,
+    sampler_name, scheduler) are derived from the base stage (stage 1; the
+    three-stage variant reuses stage 1 settings for stage 3), while the full
+    per-stage breakdown is preserved in the raw parameters.
+    """
+
+    # All per-stage parameter keys present on both node variants.
+    _STAGE_PARAM_KEYS = (
+        "stage1_steps", "stage1_cfg", "stage1_sampler_name", "stage1_scheduler",
+        "stage2_steps", "stage2_cfg", "stage2_sampler_name", "stage2_scheduler",
+    )
+
+    @staticmethod
+    def extract(node_id, inputs, outputs, metadata):
+        if not inputs:
+            return
+
+        BaseSamplerExtractor.extract_sampling_params(
+            node_id,
+            inputs,
+            metadata,
+            ("seed", "handoff_percent", "stage3_handoff_percent")
+            + KreaTwoStageSamplerExtractor._STAGE_PARAM_KEYS,
+        )
+
+        # Derive the canonical fields expected by extract_generation_params.
+        sampling_params = metadata[SAMPLING][node_id]["parameters"]
+        if "stage1_steps" in sampling_params or "stage2_steps" in sampling_params:
+            sampling_params["steps"] = (
+                (sampling_params.get("stage1_steps") or 0)
+                + (sampling_params.get("stage2_steps") or 0)
+            )
+        if "stage1_cfg" in sampling_params:
+            sampling_params["cfg"] = sampling_params["stage1_cfg"]
+        if "stage1_sampler_name" in sampling_params:
+            sampling_params["sampler_name"] = sampling_params["stage1_sampler_name"]
+        if "stage1_scheduler" in sampling_params:
+            sampling_params["scheduler"] = sampling_params["stage1_scheduler"]
+
+        BaseSamplerExtractor.extract_conditioning(node_id, inputs, metadata)
+
+        # Prefer the final generation resolution; latent dims are the fallback.
+        BaseSamplerExtractor.extract_latent_dimensions(node_id, inputs, metadata)
+        final_width = inputs.get("final_width")
+        final_height = inputs.get("final_height")
+        if final_width and final_height:
+            if SIZE not in metadata:
+                metadata[SIZE] = {}
+            metadata[SIZE][node_id] = {
+                "width": final_width,
+                "height": final_height,
+                "node_id": node_id,
+            }
+
 class LoraLoaderExtractor(NodeMetadataExtractor):
     @staticmethod
     def extract(node_id, inputs, outputs, metadata):
@@ -901,6 +960,37 @@ class ImageSizeExtractor(NodeMetadataExtractor):
             "node_id": node_id
         }
 
+class KreaDualResolutionSelectorExtractor(NodeMetadataExtractor):
+    """Extract base resolution from Krea Dual Resolution Selector outputs
+    (Auryg/Krea-2-Two-Stage-Sampler).
+
+    The node computes base/final dimensions at runtime from aspect ratio and
+    megapixel settings, so the values are only available in the update phase
+    (outputs: base_width, base_height, final_width, final_height, seed).
+    """
+
+    @staticmethod
+    def extract(node_id, inputs, outputs, metadata):
+        # Dimensions are computed at runtime; nothing to do here.
+        pass
+
+    @staticmethod
+    def update(node_id, outputs, metadata):
+        output_tuple = _first_output_tuple(outputs)
+        if not output_tuple or len(output_tuple) < 2:
+            return
+        width, height = output_tuple[0], output_tuple[1]
+        if not isinstance(width, int) or not isinstance(height, int):
+            return
+
+        if SIZE not in metadata:
+            metadata[SIZE] = {}
+        metadata[SIZE][node_id] = {
+            "width": width,
+            "height": height,
+            "node_id": node_id,
+        }
+
 class RgthreePowerLoraLoaderExtractor(NodeMetadataExtractor):
     """Extract LoRA metadata from rgthree Power Lora Loader.
 
@@ -1302,6 +1392,8 @@ NODE_EXTRACTORS = {
     "ClownsharKSampler_Beta": SamplerExtractor,
     "TSC_KSampler": TSCKSamplerExtractor,   # Efficient Nodes
     "TSC_KSamplerAdvanced": TSCKSamplerAdvancedExtractor,  # Efficient Nodes
+    "KreaTwoStageSampler": KreaTwoStageSamplerExtractor,  # Auryg/Krea-2-Two-Stage-Sampler
+    "KreaThreeStageSampler": KreaTwoStageSamplerExtractor,  # Auryg/Krea-2-Two-Stage-Sampler
     "KSamplerBasicPipe": KSamplerBasicPipeExtractor,    # comfyui-impact-pack
     "KSamplerAdvancedBasicPipe": KSamplerAdvancedBasicPipeExtractor,    # comfyui-impact-pack
     "KSampler_inspire_pipe": KSamplerBasicPipeExtractor,    # comfyui-inspire-pack
@@ -1353,6 +1445,7 @@ NODE_EXTRACTORS = {
     "GetNode": GetNodeExtractor,
     # Latent
     "EmptyLatentImage": ImageSizeExtractor,
+    "KreaDualResolutionSelector": KreaDualResolutionSelectorExtractor,  # Auryg/Krea-2-Two-Stage-Sampler
     # Flux
     "FluxGuidance": FluxGuidanceExtractor,      # Add FluxGuidance
     "CFGGuider": CFGGuiderExtractor,            # Add CFGGuider

tests/metadata_collector/test_metadata_collector.py

@@ -1613,3 +1613,213 @@ def test_fill_missing_metadata_fills_overwrite_for_muted_node(metadata_registry)
     assert "ow-1" not in metadata.get(OVERWRITE, {})
 
     metadata_registry.clear_metadata()
+
+
+def test_krea_two_stage_sampler_prompt_and_params_collected(
+    metadata_registry, monkeypatch
+):
+    """KreaTwoStageSampler should be recognized as the primary sampler and
+    contribute the prompt, canonical sampling params, and final resolution."""
+    prompt_graph = {
+        "encode_pos": {
+            "class_type": "PromptLM",
+            "inputs": {"text": "krea masterpiece", "clip": ["clip", 0]},
+        },
+        "encode_neg": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {"text": "low quality", "clip": ["clip", 0]},
+        },
+        "sampler": {
+            "class_type": "KreaTwoStageSampler",
+            "inputs": {
+                "seed": 42,
+                "handoff_percent": 16.67,
+                "stage1_steps": 52,
+                "stage1_cfg": 4.0,
+                "stage1_sampler_name": "euler",
+                "stage1_scheduler": "simple",
+                "stage2_steps": 12,
+                "stage2_cfg": 1.0,
+                "stage2_sampler_name": "euler",
+                "stage2_scheduler": "simple",
+                "final_width": 2048,
+                "final_height": 2048,
+                "upscale_method": "bislerp",
+                "positive": ["encode_pos", 0],
+                "negative": ["encode_neg", 0],
+                "latent_image": {
+                    "samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
+                },
+            },
+        },
+    }
+    prompt = SimpleNamespace(original_prompt=prompt_graph)
+
+    pos_conditioning = object()
+    neg_conditioning = object()
+
+    monkeypatch.setattr(metadata_processor, "standalone_mode", False)
+
+    metadata_registry.start_collection("krea-two-stage")
+    metadata_registry.set_current_prompt(prompt)
+
+    metadata_registry.record_node_execution(
+        "encode_pos", "PromptLM", {"text": "krea masterpiece"}, None
+    )
+    metadata_registry.update_node_execution(
+        "encode_pos", "PromptLM", [(pos_conditioning, "krea masterpiece")]
+    )
+    metadata_registry.record_node_execution(
+        "encode_neg", "CLIPTextEncode", {"text": "low quality"}, None
+    )
+    metadata_registry.update_node_execution(
+        "encode_neg", "CLIPTextEncode", [(neg_conditioning,)]
+    )
+    metadata_registry.record_node_execution(
+        "sampler",
+        "KreaTwoStageSampler",
+        {
+            "seed": 42,
+            "handoff_percent": 16.67,
+            "stage1_steps": 52,
+            "stage1_cfg": 4.0,
+            "stage1_sampler_name": "euler",
+            "stage1_scheduler": "simple",
+            "stage2_steps": 12,
+            "stage2_cfg": 1.0,
+            "stage2_sampler_name": "euler",
+            "stage2_scheduler": "simple",
+            "final_width": 2048,
+            "final_height": 2048,
+            "upscale_method": "bislerp",
+            "positive": pos_conditioning,
+            "negative": neg_conditioning,
+            "latent_image": {
+                "samples": types.SimpleNamespace(shape=(1, 4, 16, 16))
+            },
+        },
+        None,
+    )
+
+    metadata = metadata_registry.get_metadata("krea-two-stage")
+
+    sampler_data = metadata[SAMPLING]["sampler"]
+    assert sampler_data["is_sampler"] is True
+    parameters = sampler_data["parameters"]
+    assert parameters["seed"] == 42
+    assert parameters["steps"] == 64
+    assert parameters["cfg"] == 4.0
+    assert parameters["sampler_name"] == "euler"
+    assert parameters["scheduler"] == "simple"
+    assert parameters["stage1_steps"] == 52
+    assert parameters["stage2_cfg"] == 1.0
+
+    assert metadata[SIZE]["sampler"] == {
+        "width": 2048,
+        "height": 2048,
+        "node_id": "sampler",
+    }
+
+    prompt_results = MetadataProcessor.match_conditioning_to_prompts(
+        metadata, "sampler"
+    )
+    assert prompt_results["prompt"] == "krea masterpiece"
+    assert prompt_results["negative_prompt"] == "low quality"
+
+    params = MetadataProcessor.extract_generation_params(metadata)
+    assert params["prompt"] == "krea masterpiece"
+    assert params["negative_prompt"] == "low quality"
+    assert params["seed"] == 42
+    assert params["steps"] == 64
+    assert params["cfg_scale"] == 4.0
+    assert params["sampler"] == "euler"
+    assert params["scheduler"] == "simple"
+    assert params["size"] == "2048x2048"
+
+
+def test_krea_three_stage_sampler_uses_stage1_canonical_fields(metadata_registry):
+    """KreaThreeStageSampler reuses stage 1 settings for stage 3, so canonical
+    fields map from stage 1 and the total counts both sampling stages."""
+    metadata_registry.start_collection("krea-three-stage")
+    metadata_registry.set_current_prompt(SimpleNamespace(original_prompt={}))
+
+    metadata_registry.record_node_execution(
+        "sampler",
+        "KreaThreeStageSampler",
+        {
+            "seed": 7,
+            "handoff_percent": 16.67,
+            "stage3_handoff_percent": 83.33,
+            "stage1_steps": 52,
+            "stage1_cfg": 4.0,
+            "stage1_sampler_name": "euler",
+            "stage1_scheduler": "simple",
+            "stage2_steps": 12,
+            "stage2_cfg": 1.0,
+            "stage2_sampler_name": "euler",
+            "stage2_scheduler": "simple",
+            "final_width": 1024,
+            "final_height": 2048,
+            "upscale_method": "bislerp",
+            "positive": object(),
+            "negative": object(),
+            "latent_image": {"samples": types.SimpleNamespace(shape=(1, 4, 8, 16))},
+        },
+        None,
+    )
+
+    metadata = metadata_registry.get_metadata("krea-three-stage")
+
+    sampler_data = metadata[SAMPLING]["sampler"]
+    assert sampler_data["is_sampler"] is True
+    parameters = sampler_data["parameters"]
+    assert parameters["seed"] == 7
+    assert parameters["stage3_handoff_percent"] == 83.33
+    assert parameters["steps"] == 64
+    assert parameters["cfg"] == 4.0
+    assert parameters["sampler_name"] == "euler"
+    assert parameters["scheduler"] == "simple"
+
+    # Final resolution takes precedence over the latent dimensions (64x128).
+    assert metadata[SIZE]["sampler"] == {
+        "width": 1024,
+        "height": 2048,
+        "node_id": "sampler",
+    }
+
+
+def test_krea_dual_resolution_selector_extracts_size_from_outputs(
+    metadata_registry,
+):
+    """KreaDualResolutionSelector computes dimensions at runtime, so the base
+    resolution is recorded from its outputs in the update phase."""
+    metadata_registry.start_collection("krea-selector")
+    metadata_registry.set_current_prompt(SimpleNamespace(original_prompt={}))
+
+    metadata_registry.record_node_execution(
+        "selector",
+        "KreaDualResolutionSelector",
+        {
+            "aspect_ratio": "1:1",
+            "base_megapixels": 1.0,
+            "final_megapixels": 2.0,
+            "multiple": 16,
+            "random_seed": 123,
+        },
+        None,
+        return_types=("INT", "INT", "INT", "INT", "INT"),
+    )
+    metadata_registry.update_node_execution(
+        "selector",
+        "KreaDualResolutionSelector",
+        [(1024, 1024, 2048, 2048, 123)],
+        return_types=("INT", "INT", "INT", "INT", "INT"),
+    )
+
+    metadata = metadata_registry.get_metadata("krea-selector")
+
+    assert metadata[SIZE]["selector"] == {
+        "width": 1024,
+        "height": 1024,
+        "node_id": "selector",
+    }