"""ComfyUI nodes for Apple's STARFlow image and STARFlow-V video models."""

from .nodes import (
    STARFlowCheckpointLoader,
    STARFlowT5TextEncode,
    STARFlowSampler,
    STARFlowVAEDecode,
)
from .starflow_v.nodes import (
    NODE_CLASS_MAPPINGS as STARFLOW_V_NODE_CLASS_MAPPINGS,
    NODE_DISPLAY_NAME_MAPPINGS as STARFLOW_V_NODE_DISPLAY_NAME_MAPPINGS,
)

NODE_CLASS_MAPPINGS = {
    "STARFlowCheckpointLoader": STARFlowCheckpointLoader,
    "STARFlowT5TextEncode": STARFlowT5TextEncode,
    "STARFlowSampler": STARFlowSampler,
    "STARFlowVAEDecode": STARFlowVAEDecode,
}
NODE_CLASS_MAPPINGS.update(STARFLOW_V_NODE_CLASS_MAPPINGS)

NODE_DISPLAY_NAME_MAPPINGS = {
    "STARFlowCheckpointLoader": "STARFlow Checkpoint Loader",
    "STARFlowT5TextEncode": "STARFlow T5 Text Encode",
    "STARFlowSampler": "STARFlow Sampler",
    "STARFlowVAEDecode": "STARFlow VAE Decode",
}
NODE_DISPLAY_NAME_MAPPINGS.update(STARFLOW_V_NODE_DISPLAY_NAME_MAPPINGS)

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]