diff --git a/py/metadata_collector/metadata_processor.py b/py/metadata_collector/metadata_processor.py index a7675f7c..8b9a541f 100644 --- a/py/metadata_collector/metadata_processor.py +++ b/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 diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py index 9076fe1a..2eeb31bd 100644 --- a/py/metadata_collector/node_extractors.py +++ b/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 diff --git a/tests/metadata_collector/test_metadata_collector.py b/tests/metadata_collector/test_metadata_collector.py index f9efd403..b4d50443 100644 --- a/tests/metadata_collector/test_metadata_collector.py +++ b/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", + }