comfy_api_nodes/mapper_utils.py

from enum import Enum

from pydantic.fields import FieldInfo
from pydantic import BaseModel
from pydantic_core import PydanticUndefined

from comfy.comfy_types.node_typing import IO, InputTypeOptions

NodeInput = tuple[IO, InputTypeOptions]


def _create_base_config(field_info: FieldInfo) -> InputTypeOptions:
    config = {}
    if hasattr(field_info, "default") and field_info.default is not PydanticUndefined:
        config["default"] = field_info.default
    if hasattr(field_info, "description") and field_info.description is not None:
        config["tooltip"] = field_info.description
    return config


def _get_number_constraints_config(field_info: FieldInfo) -> dict:
    config = {}
    if hasattr(field_info, "metadata"):
        metadata = field_info.metadata
        for constraint in metadata:
            if hasattr(constraint, "ge"):
                config["min"] = constraint.ge
            if hasattr(constraint, "le"):
                config["max"] = constraint.le
            if hasattr(constraint, "multiple_of"):
                config["step"] = constraint.multiple_of
    return config


def _model_field_to_image_input(field_info: FieldInfo, **kwargs) -> NodeInput:
    return IO.IMAGE, {
        **_create_base_config(field_info),
        **kwargs,
    }


def _model_field_to_string_input(field_info: FieldInfo, **kwargs) -> NodeInput:
    return IO.STRING, {
        **_create_base_config(field_info),
        **kwargs,
    }


def _model_field_to_float_input(field_info: FieldInfo, **kwargs) -> NodeInput:
    return IO.FLOAT, {
        **_create_base_config(field_info),
        **_get_number_constraints_config(field_info),
        **kwargs,
    }


def _model_field_to_int_input(field_info: FieldInfo, **kwargs) -> NodeInput:
    return IO.INT, {
        **_create_base_config(field_info),
        **_get_number_constraints_config(field_info),
        **kwargs,
    }


def _model_field_to_combo_input(
    field_info: FieldInfo, enum_type: type[Enum] = None, **kwargs
) -> NodeInput:
    combo_config = {}
    if enum_type is not None:
        combo_config["options"] = [option.value for option in enum_type]
    combo_config = {
        **combo_config,
        **_create_base_config(field_info),
        **kwargs,
    }
    return IO.COMBO, combo_config


def model_field_to_node_input(
    input_type: IO, base_model: type[BaseModel], field_name: str, **kwargs
) -> NodeInput:
    """
    Maps a field from a Pydantic model to a Comfy node input.

    Args:
        input_type: The type of the input.
        base_model: The Pydantic model to map the field from.
        field_name: The name of the field to map.
        **kwargs: Additional key/values to include in the input options.

    Note:
        For combo inputs, pass an `Enum` to the `enum_type` keyword argument to populate the options automatically.

    Example:
        >>> model_field_to_node_input(IO.STRING, MyModel, "my_field", multiline=True)
        >>> model_field_to_node_input(IO.COMBO, MyModel, "my_field", enum_type=MyEnum)
        >>> model_field_to_node_input(IO.FLOAT, MyModel, "my_field", slider=True)
    """
    field_info: FieldInfo = base_model.model_fields[field_name]
    result: NodeInput

    if input_type == IO.IMAGE:
        result = _model_field_to_image_input(field_info, **kwargs)
    elif input_type == IO.STRING:
        result = _model_field_to_string_input(field_info, **kwargs)
    elif input_type == IO.FLOAT:
        result = _model_field_to_float_input(field_info, **kwargs)
    elif input_type == IO.INT:
        result = _model_field_to_int_input(field_info, **kwargs)
    elif input_type == IO.COMBO:
        result = _model_field_to_combo_input(field_info, **kwargs)
    else:
        message = f"Invalid input type: {input_type}"
        raise ValueError(message)

    return result