398 lines
12 KiB
Python
398 lines
12 KiB
Python
"""Executable publication contract for canonical Krea 2 RAW host workflows."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import struct
|
|
from dataclasses import FrozenInstanceError
|
|
from typing import Any, cast
|
|
|
|
import comfyui_sigmax
|
|
import pytest
|
|
from comfyui_sigmax.core import (
|
|
ExecutionStatus,
|
|
ScheduleContractError,
|
|
ScheduleOwnership,
|
|
deserialize_portable_execution_bundle,
|
|
)
|
|
from comfyui_sigmax.nodes.inspectors import build_schedule_inspection
|
|
from comfyui_sigmax.nodes.krea2_sigma_scheduler import (
|
|
bind_krea2_sigma_output_info,
|
|
build_krea2_sigma_schedule,
|
|
)
|
|
from comfyui_sigmax.nodes.raw_workflow_output import (
|
|
RAW_WORKFLOW_OUTPUT_NODE_ID,
|
|
RAW_WORKFLOW_OUTPUT_SCHEMA_ID,
|
|
RawWorkflowOutput,
|
|
RawWorkflowOutputResult,
|
|
build_raw_workflow_output,
|
|
)
|
|
|
|
CASES = (
|
|
(
|
|
52,
|
|
1024,
|
|
1024,
|
|
True,
|
|
1024,
|
|
1024,
|
|
4096,
|
|
0.90625,
|
|
"krea2.raw.official-full-52",
|
|
"official",
|
|
"sha256:5ff69c30df41c7f37eae14502155b31f23724d32427180f69118cabcd6a3ac61",
|
|
),
|
|
(
|
|
52,
|
|
1353,
|
|
761,
|
|
True,
|
|
1360,
|
|
768,
|
|
4080,
|
|
0.9045572916666667,
|
|
"krea2.raw.official-full-52",
|
|
"official",
|
|
"sha256:01352f42660bd3b31bbaf7548a9891273899afd375adeb68c7f7c93fd2a4f0d4",
|
|
),
|
|
(
|
|
28,
|
|
761,
|
|
1353,
|
|
False,
|
|
768,
|
|
1360,
|
|
4080,
|
|
0.9045572916666667,
|
|
"krea2.raw.diffusers-reference-28",
|
|
"framework_reference",
|
|
"sha256:52208c5fa3780c95cce399b1f842f3fea56503e76fdf5ef4abc3069cf3108f01",
|
|
),
|
|
)
|
|
|
|
|
|
def _upstream(
|
|
*,
|
|
steps: int = 52,
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
strict_official: bool = True,
|
|
) -> tuple[tuple[float, ...], str, str]:
|
|
schedule = build_krea2_sigma_schedule(
|
|
variant="RAW",
|
|
steps=steps,
|
|
width=width,
|
|
height=height,
|
|
strict_official=strict_official,
|
|
start_step=0,
|
|
end_step=-1,
|
|
)
|
|
inspection = build_schedule_inspection(
|
|
sigmas=schedule.sigmas,
|
|
schedule_info=schedule.schedule_info_json,
|
|
)
|
|
return schedule.sigmas, schedule.schedule_info_json, inspection.report_json
|
|
|
|
|
|
def _result() -> RawWorkflowOutputResult:
|
|
sigmas, schedule_info, schedule_report = _upstream()
|
|
return build_raw_workflow_output(
|
|
sigmas=sigmas,
|
|
schedule_info=schedule_info,
|
|
schedule_report=schedule_report,
|
|
)
|
|
|
|
|
|
def test_output_node_declares_stable_v1_schema() -> None:
|
|
inputs = RawWorkflowOutput.INPUT_TYPES()
|
|
|
|
assert RAW_WORKFLOW_OUTPUT_NODE_ID == "Sigmax.RawWorkflowOutput"
|
|
assert RAW_WORKFLOW_OUTPUT_SCHEMA_ID == "sigmax.raw-workflow-output/1"
|
|
assert RawWorkflowOutput.RETURN_TYPES == ()
|
|
assert RawWorkflowOutput.RETURN_NAMES == ()
|
|
assert RawWorkflowOutput.FUNCTION == "publish"
|
|
assert RawWorkflowOutput.CATEGORY == "Sigmax/workflows"
|
|
assert RawWorkflowOutput.OUTPUT_NODE is True
|
|
assert inputs == {
|
|
"required": {
|
|
"sigmas": ("SIGMAS",),
|
|
"schedule_info": ("STRING", {"default": "", "multiline": True}),
|
|
"schedule_report": ("STRING", {"default": "", "multiline": True}),
|
|
}
|
|
}
|
|
assert inputs == RawWorkflowOutput.INPUT_TYPES()
|
|
assert inputs is not RawWorkflowOutput.INPUT_TYPES()
|
|
|
|
|
|
def test_builtin_mapping_registers_the_raw_output_node() -> None:
|
|
assert comfyui_sigmax.NODE_CLASS_MAPPINGS[RAW_WORKFLOW_OUTPUT_NODE_ID] is RawWorkflowOutput
|
|
assert (
|
|
comfyui_sigmax.NODE_DISPLAY_NAME_MAPPINGS[RAW_WORKFLOW_OUTPUT_NODE_ID]
|
|
== "RAW Workflow Output"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"steps",
|
|
"width",
|
|
"height",
|
|
"strict_official",
|
|
"effective_width",
|
|
"effective_height",
|
|
"image_seq_len",
|
|
"mu",
|
|
"recipe_id",
|
|
"evidence",
|
|
"numerical_fingerprint",
|
|
),
|
|
CASES,
|
|
)
|
|
def test_output_builds_canonical_raw_bundle_for_every_published_case(
|
|
steps: int,
|
|
width: int,
|
|
height: int,
|
|
strict_official: bool,
|
|
effective_width: int,
|
|
effective_height: int,
|
|
image_seq_len: int,
|
|
mu: float,
|
|
recipe_id: str,
|
|
evidence: str,
|
|
numerical_fingerprint: str,
|
|
) -> None:
|
|
sigmas, schedule_info, schedule_report = _upstream(
|
|
steps=steps,
|
|
width=width,
|
|
height=height,
|
|
strict_official=strict_official,
|
|
)
|
|
result = build_raw_workflow_output(
|
|
sigmas=sigmas,
|
|
schedule_info=schedule_info,
|
|
schedule_report=schedule_report,
|
|
)
|
|
bundle = deserialize_portable_execution_bundle(result.bundle_json)
|
|
construction = bundle.artifact.construction_projection()
|
|
receipt = bundle.receipt.projection()
|
|
requested = cast(dict[str, Any], construction["requested"])
|
|
effective = cast(dict[str, Any], construction["effective"])
|
|
|
|
assert result.schema_id == RAW_WORKFLOW_OUTPUT_SCHEMA_ID
|
|
assert result.artifact_numerical_fingerprint == numerical_fingerprint
|
|
assert bundle.artifact.numerical_fingerprint == numerical_fingerprint
|
|
assert {key: requested[key] for key in ("height", "steps", "width")} == {
|
|
"height": height,
|
|
"steps": steps,
|
|
"width": width,
|
|
}
|
|
assert {key: effective[key] for key in ("height", "steps", "width")} == {
|
|
"height": effective_height,
|
|
"steps": steps,
|
|
"width": effective_width,
|
|
}
|
|
assert requested["profile"] == effective["profile"] == "krea2.raw.official"
|
|
assert cast(dict[str, object], construction["evidence"])["level"] == evidence
|
|
assert cast(dict[str, object], construction["source"])["id"] == (
|
|
"krea.krea2.official" if evidence == "official" else "diffusers.krea2.framework"
|
|
)
|
|
|
|
transforms = cast(list[dict[str, Any]], construction["transforms"])
|
|
assert [item["id"] for item in transforms] == [
|
|
"krea.exponential_mu",
|
|
"terminal.append_zero",
|
|
]
|
|
assert sum(item["id"] == "krea.exponential_mu" for item in transforms) == 1
|
|
assert transforms[0]["parameters"]["mu"] == {
|
|
"bits": struct.pack(">d", mu).hex(),
|
|
"precision": "float64",
|
|
}
|
|
base_grid = cast(dict[str, Any], construction["base_grid"])
|
|
assert base_grid["parameters"] == {
|
|
"image_seq_len": image_seq_len,
|
|
"recipe": recipe_id,
|
|
"steps": steps,
|
|
}
|
|
assert construction["ownership"] == {
|
|
"schedule": ScheduleOwnership.EXTERNAL_SIGMAS.value.casefold(),
|
|
"shift": "construction_pipeline",
|
|
}
|
|
assert receipt["execution"] == {
|
|
"reason_code": None,
|
|
"status": ExecutionStatus.NOT_EXECUTED.value,
|
|
}
|
|
assert receipt["counts"] == {
|
|
"effective_model_evaluations": 0,
|
|
"effective_transitions": 0,
|
|
"requested_model_evaluations": steps,
|
|
"requested_transitions": steps,
|
|
}
|
|
|
|
|
|
def test_output_is_deterministic_and_returns_only_bounded_history_payload() -> None:
|
|
sigmas, schedule_info, schedule_report = _upstream()
|
|
first = build_raw_workflow_output(
|
|
sigmas=sigmas,
|
|
schedule_info=schedule_info,
|
|
schedule_report=schedule_report,
|
|
)
|
|
second = build_raw_workflow_output(
|
|
sigmas=sigmas,
|
|
schedule_info=schedule_info,
|
|
schedule_report=schedule_report,
|
|
)
|
|
|
|
assert first == second
|
|
response = RawWorkflowOutput().publish(
|
|
sigmas=list(sigmas),
|
|
schedule_info=schedule_info,
|
|
schedule_report=schedule_report,
|
|
)
|
|
assert set(response) == {"result", "ui"}
|
|
assert response["result"] == ()
|
|
ui = cast(dict[str, object], response["ui"])
|
|
assert set(ui) == {"sigmax_execution_bundle"}
|
|
assert cast(list[object], ui["sigmax_execution_bundle"]) == [first.bundle_json]
|
|
|
|
|
|
@pytest.mark.parametrize("case", CASES)
|
|
def test_output_accepts_exact_host_float32_and_keeps_canonical_float64_artifact(
|
|
case: tuple[object, ...],
|
|
) -> None:
|
|
(
|
|
steps,
|
|
width,
|
|
height,
|
|
strict_official,
|
|
_effective_width,
|
|
_effective_height,
|
|
_image_seq_len,
|
|
_mu,
|
|
_recipe_id,
|
|
_evidence,
|
|
numerical_fingerprint,
|
|
) = case
|
|
schedule = build_krea2_sigma_schedule(
|
|
variant="RAW",
|
|
steps=steps,
|
|
width=width,
|
|
height=height,
|
|
strict_official=strict_official,
|
|
start_step=0,
|
|
end_step=-1,
|
|
)
|
|
host_values = tuple(
|
|
struct.unpack(">f", struct.pack(">f", value))[0] for value in schedule.sigmas
|
|
)
|
|
host_info = bind_krea2_sigma_output_info(schedule, output_sigmas=host_values)
|
|
host_report = build_schedule_inspection(
|
|
sigmas=host_values,
|
|
schedule_info=host_info,
|
|
).report_json
|
|
|
|
result = build_raw_workflow_output(
|
|
sigmas=host_values,
|
|
schedule_info=host_info,
|
|
schedule_report=host_report,
|
|
)
|
|
bundle = deserialize_portable_execution_bundle(result.bundle_json)
|
|
|
|
assert host_values != schedule.sigmas
|
|
assert bundle.artifact.numerical_fingerprint == numerical_fingerprint
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"mutation",
|
|
(
|
|
"sigma_value",
|
|
"schedule_info",
|
|
"schedule_report",
|
|
"turbo",
|
|
"partial_slice",
|
|
"wrong_mu",
|
|
"swapped_effective_dimensions",
|
|
),
|
|
)
|
|
def test_output_rejects_noncanonical_or_mismatched_upstream_evidence(mutation: str) -> None:
|
|
sigmas, schedule_info, schedule_report = _upstream(
|
|
width=1353,
|
|
height=761,
|
|
)
|
|
changed_sigmas = sigmas
|
|
changed_info = schedule_info
|
|
changed_report = schedule_report
|
|
|
|
if mutation == "sigma_value":
|
|
changed_sigmas = (sigmas[0] - 0.01, *sigmas[1:])
|
|
elif mutation in {
|
|
"schedule_info",
|
|
"wrong_mu",
|
|
"swapped_effective_dimensions",
|
|
}:
|
|
projection = cast(dict[str, Any], json.loads(schedule_info))
|
|
if mutation == "schedule_info":
|
|
projection["strict_official"] = False
|
|
elif mutation == "wrong_mu":
|
|
projection["shift"]["mu"] = cast(float, projection["shift"]["mu"]) + 0.01
|
|
else:
|
|
effective = projection["dimensions"]["effective"]
|
|
effective["width"], effective["height"] = effective["height"], effective["width"]
|
|
changed_info = json.dumps(projection, sort_keys=True, separators=(",", ":"))
|
|
elif mutation == "schedule_report":
|
|
projection = cast(dict[str, Any], json.loads(schedule_report))
|
|
projection["fingerprints"]["verified"] = False
|
|
changed_report = json.dumps(projection, sort_keys=True, separators=(",", ":"))
|
|
else:
|
|
alternate = build_krea2_sigma_schedule(
|
|
variant="Turbo" if mutation == "turbo" else "RAW",
|
|
steps=8 if mutation == "turbo" else 52,
|
|
width=1353,
|
|
height=761,
|
|
strict_official=True,
|
|
start_step=0 if mutation == "turbo" else 2,
|
|
end_step=-1 if mutation == "turbo" else 20,
|
|
)
|
|
changed_sigmas = alternate.sigmas
|
|
changed_info = alternate.schedule_info_json
|
|
changed_report = build_schedule_inspection(
|
|
sigmas=changed_sigmas,
|
|
schedule_info=changed_info,
|
|
).report_json
|
|
|
|
with pytest.raises(ScheduleContractError):
|
|
build_raw_workflow_output(
|
|
sigmas=changed_sigmas,
|
|
schedule_info=changed_info,
|
|
schedule_report=changed_report,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"changes",
|
|
(
|
|
{"sigmas": object()},
|
|
{"schedule_info": object()},
|
|
{"schedule_report": object()},
|
|
{"schedule_info": ""},
|
|
{"schedule_report": ""},
|
|
),
|
|
)
|
|
def test_output_rejects_wrong_host_boundary_types(changes: dict[str, object]) -> None:
|
|
sigmas, schedule_info, schedule_report = _upstream()
|
|
arguments: dict[str, object] = {
|
|
"sigmas": sigmas,
|
|
"schedule_info": schedule_info,
|
|
"schedule_report": schedule_report,
|
|
}
|
|
arguments.update(changes)
|
|
|
|
with pytest.raises(ScheduleContractError):
|
|
RawWorkflowOutput().publish(**arguments)
|
|
|
|
|
|
def test_output_result_is_immutable() -> None:
|
|
result = _result()
|
|
|
|
with pytest.raises(FrozenInstanceError):
|
|
result.bundle_json = "{}" # type: ignore[misc]
|