Files
rookiestar28-ComfyUI-Sigmax/scripts/generate_wan_workflow_fixtures.py
rookiestar28 076c5a2a63 feat(wan): add source-qualified sigma schedules
Add clean-room Wan 2.1 and Wan 2.2 profiles, explicit scheduler node, deterministic workflow/golden evidence, and host contract seams. Keep A14B routing caller-owned and defer derivatives and sampler semantics.
2026-08-03 03:04:38 +08:00

258 lines
9.4 KiB
Python

"""Add deterministic Wan 2.1/2.2 workflow fixtures and host schema evidence."""
from __future__ import annotations
import argparse
import json
from collections.abc import Sequence
from pathlib import Path
from typing import Any, cast
from comfyui_sigmax.adapters.registration import builtin_node_registry
from comfyui_sigmax.core import (
CapabilityDimension,
CompatibilityDecision,
CompatibilityLevel,
CompatibilityReason,
WorkflowArtifactReference,
WorkflowHostRequirement,
WorkflowMetadata,
WorkflowRequirement,
attach_workflow_metadata,
construction_fingerprint,
numerical_fingerprint,
)
from comfyui_sigmax.nodes.wan_sigma_scheduler import build_wan_sigma_schedule
def _json(value: object) -> str:
return json.dumps(value, allow_nan=False, ensure_ascii=False, indent=2) + "\n"
def _fixture(
*,
identifier: str,
variant: str,
generation: str,
task: str,
source: str,
resolution: str,
steps: int,
profile_id: str,
) -> dict[str, object]:
result = build_wan_sigma_schedule(
generation=generation,
task=task,
source=source,
resolution=resolution,
steps=steps,
strict_source=True,
start_step=0,
end_step=-1,
already_shifted=False,
)
numerical = numerical_fingerprint(result.sigmas, domain=result.domain, precision="float64")
info = json.loads(result.schedule_info_json)
construction = {
"base_grid": {
"id": "flowmatch.reciprocal_step",
"parameters": {"training_timesteps": 1000},
},
"effective": {"steps": steps},
"engine": {"version": "1.0.0"},
"evidence": {"level": info["profile"]["evidence"]},
"numerical_fingerprint": numerical,
"overrides": [],
"ownership": "external_sigmas",
"requested": {"steps": steps},
"schema": "sigmax.schedule-artifact/1",
"slicing": {"end_step": None, "start_step": 0},
"source": {
"profile_id": profile_id,
"revision": info["profile"]["revision"],
},
"terminal": {"policy": "append_zero"},
"transforms": [{"id": "direct_ratio.shift", "ratio": str(info["shift"]["ratio"])}],
"warnings": list(info["warnings"]),
}
node_type = "Sigmax.WanSigmaScheduler"
node = WorkflowRequirement(identifier=node_type, version="1")
profile = WorkflowRequirement(identifier=profile_id, version="1")
metadata = WorkflowMetadata(
package=WorkflowRequirement(identifier="comfyui-sigmax", version="1.0.0"),
nodes=(node,),
host=WorkflowHostRequirement(
identifier="comfyui", version="0.29.0", api_version="legacy_v1"
),
profile=profile,
compatibility=CompatibilityDecision(
level=CompatibilityLevel.ALLOW,
considered=tuple(CapabilityDimension),
reasons=(CompatibilityReason.COMPATIBLE,),
),
artifact=WorkflowArtifactReference(
construction_fingerprint=construction_fingerprint(construction),
numerical_fingerprint=numerical,
),
)
workflow = attach_workflow_metadata(
{
"version": 0.4,
"last_node_id": 1,
"last_link_id": 0,
"nodes": [
{
"id": 1,
"type": node_type,
"pos": [40, 80],
"size": [440, 330],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "sigmas", "type": "SIGMAS", "links": None, "slot_index": 0},
{"name": "boundary_step", "type": "INT", "links": None, "slot_index": 1},
{"name": "schedule_info", "type": "STRING", "links": None, "slot_index": 2},
],
"properties": {"Node name for S&R": node_type, "cnr_id": "comfyui-sigmax"},
"widgets_values": [
generation,
task,
source,
resolution,
steps,
True,
0,
-1,
False,
],
}
],
"links": [],
},
metadata,
)
return {
"host": {"api_version": "legacy_v1", "id": "comfyui", "version": "0.29.0"},
"id": identifier,
"node_contracts": [
{
"node_id": 1,
"node_type": node_type,
"widget_slots": [
{"fixed": generation, "name": "generation", "type": "COMBO"},
{"fixed": task, "name": "task", "type": "COMBO"},
{"fixed": source, "name": "source", "type": "COMBO"},
{"fixed": resolution, "name": "resolution", "type": "COMBO"},
{"name": "steps", "type": "INT"},
{"name": "strict_source", "type": "BOOLEAN"},
{"name": "start_step", "type": "INT"},
{"name": "end_step", "type": "INT"},
{"name": "already_shifted", "type": "BOOLEAN"},
],
}
],
"nodes": [node.projection()],
"package": {"id": "comfyui-sigmax", "version": "1.0.0"},
"profile": profile.projection(),
"variant": variant,
"workflow": workflow,
}
def main(arguments: Sequence[str] | None = None) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--fixtures", type=Path, required=True)
parser.add_argument("--host-baseline", type=Path, required=True)
parsed = parser.parse_args(arguments)
bundle = cast(dict[str, Any], json.loads(parsed.fixtures.read_text(encoding="utf-8")))
retained = [item for item in bundle["fixtures"] if not item["id"].startswith("wan")]
retained.extend(
(
_fixture(
identifier="wan21-t2v-official-50",
variant="Wan 2.1 T2V",
generation="Wan 2.1",
task="T2V",
source="Official native",
resolution="None",
steps=50,
profile_id="wan2.1.t2v.official-native",
),
_fixture(
identifier="wan21-i2v-480p-official-40",
variant="Wan 2.1 I2V 480P",
generation="Wan 2.1",
task="I2V",
source="Official native",
resolution="480P",
steps=40,
profile_id="wan2.1.i2v.480p.official-native",
),
_fixture(
identifier="wan22-ti2v-5b-native-50",
variant="Wan 2.2 TI2V 5B",
generation="Wan 2.2",
task="TI2V",
source="ComfyUI native",
resolution="None",
steps=50,
profile_id="wan2.2.ti2v.5b.comfy-native",
),
_fixture(
identifier="wan22-t2v-a14b-native-40",
variant="Wan 2.2 T2V A14B",
generation="Wan 2.2",
task="T2V A14B",
source="Official native",
resolution="None",
steps=40,
profile_id="wan2.2.t2v-a14b.official-native",
),
)
)
bundle["fixtures"] = sorted(retained, key=lambda item: item["id"])
parsed.fixtures.write_text(_json(bundle), encoding="utf-8", newline="\n")
baseline = cast(dict[str, Any], json.loads(parsed.host_baseline.read_text(encoding="utf-8")))
registry = builtin_node_registry()
required_ids = {
"Sigmax.AnimaSigmaScheduler",
"Sigmax.AuraFlowSigmaScheduler",
"Sigmax.Flux1SchnellSigmaScheduler",
"Sigmax.HunyuanImage21SigmaScheduler",
"Sigmax.Krea2ConditioningRebalance",
"Sigmax.Krea2SigmaScheduler",
"Sigmax.Lumina2SigmaScheduler",
"Sigmax.QwenImageSigmaScheduler",
"Sigmax.RawWorkflowOutput",
"Sigmax.SD3SigmaScheduler",
"Sigmax.ScheduleInspector",
"Sigmax.TurboWorkflowOutput",
"Sigmax.WanSigmaScheduler",
"Sigmax.ZImageSigmaScheduler",
}
object_info = registry.object_info_projection()
node_definition_v2 = registry.node_definition_v2_projection()
classes = registry.class_mappings()
for node_id in required_ids:
node_class = cast(Any, classes[node_id])
input_order = tuple(node_class.INPUT_TYPES().get("required", {}))
legacy_input = object_info[node_id].get("input", {})
if isinstance(legacy_input, dict) and isinstance(legacy_input.get("required"), dict):
required = legacy_input["required"]
legacy_input["required"] = {name: required[name] for name in input_order}
v2_inputs = node_definition_v2[node_id].get("inputs", {})
if isinstance(v2_inputs, dict):
node_definition_v2[node_id]["inputs"] = {
name: v2_inputs[name] for name in input_order if name in v2_inputs
}
baseline["object_info"] = {key: object_info[key] for key in sorted(required_ids)}
baseline["node_definition_v2"] = {key: node_definition_v2[key] for key in sorted(required_ids)}
parsed.host_baseline.write_text(_json(baseline), encoding="utf-8", newline="\n")
return 0
if __name__ == "__main__":
raise SystemExit(main())