Add utils to map from pydantic model fields to comfy node inputs (#30)

comfy_api_nodes/README.md

@@ -36,7 +36,6 @@ npm install -g @redocly/cli
 redocly bundle openapi.yaml --output filtered-openapi.yaml --config comfy_api_nodes/redocly-dev.yaml --remove-unused-components
 
 # Generate the pydantic datamodels for validation.
-datamodel-codegen --use-subclass-enum --input filtered-openapi.yaml --output comfy_api_nodes/apis/__init__.py --output-model-type pydantic_v2.BaseModel
+datamodel-codegen --use-subclass-enum --field-constraints --input filtered-openapi.yaml --output comfy_api_nodes/apis/__init__.py --output-model-type pydantic_v2.BaseModel
 
 ```
-

comfy_api_nodes/mapper_utils.py

@@ -0,0 +1,117 @@
+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
+
+    match input_type:
+        case IO.IMAGE:
+            result = _model_field_to_image_input(field_info, **kwargs)
+        case IO.STRING:
+            result = _model_field_to_string_input(field_info, **kwargs)
+        case IO.FLOAT:
+            result = _model_field_to_float_input(field_info, **kwargs)
+        case IO.INT:
+            result = _model_field_to_int_input(field_info, **kwargs)
+        case IO.COMBO:
+            result = _model_field_to_combo_input(field_info, **kwargs)
+        case _:
+            message = f"Invalid input type: {input_type}"
+            raise ValueError(message)
+
+    return result

tests-unit/comfy_api_nodes_test/mapper_utils_test.py

@@ -0,0 +1,297 @@
+from typing import Optional
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+from comfy.comfy_types.node_typing import IO
+from comfy_api_nodes.mapper_utils import model_field_to_node_input
+
+
+def test_model_field_to_float_input():
+    """Tests mapping a float field with constraints."""
+
+    class ModelWithFloatField(BaseModel):
+        cfg_scale: Optional[float] = Field(
+            default=0.5,
+            description="Flexibility in video generation",
+            ge=0.0,
+            le=1.0,
+            multiple_of=0.001,
+        )
+
+    expected_output = (
+        IO.FLOAT,
+        {
+            "default": 0.5,
+            "tooltip": "Flexibility in video generation",
+            "min": 0.0,
+            "max": 1.0,
+            "step": 0.001,
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.FLOAT, ModelWithFloatField, "cfg_scale"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_float_input_no_constraints():
+    """Tests mapping a float field with no constraints."""
+
+    class ModelWithFloatField(BaseModel):
+        cfg_scale: Optional[float] = Field(default=0.5)
+
+    expected_output = (
+        IO.FLOAT,
+        {
+            "default": 0.5,
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.FLOAT, ModelWithFloatField, "cfg_scale"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_int_input():
+    """Tests mapping an int field with constraints."""
+
+    class ModelWithIntField(BaseModel):
+        num_frames: Optional[int] = Field(
+            default=10,
+            description="Number of frames to generate",
+            ge=1,
+            le=100,
+            multiple_of=1,
+        )
+
+    expected_output = (
+        IO.INT,
+        {
+            "default": 10,
+            "tooltip": "Number of frames to generate",
+            "min": 1,
+            "max": 100,
+            "step": 1,
+        },
+    )
+
+    actual_output = model_field_to_node_input(IO.INT, ModelWithIntField, "num_frames")
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_string_input():
+    """Tests mapping a string field."""
+
+    class ModelWithStringField(BaseModel):
+        prompt: Optional[str] = Field(
+            default="A beautiful sunset over a calm ocean",
+            description="A prompt for the video generation",
+        )
+
+    expected_output = (
+        IO.STRING,
+        {
+            "default": "A beautiful sunset over a calm ocean",
+            "tooltip": "A prompt for the video generation",
+        },
+    )
+
+    actual_output = model_field_to_node_input(IO.STRING, ModelWithStringField, "prompt")
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_string_input_multiline():
+    """Tests mapping a string field."""
+
+    class ModelWithStringField(BaseModel):
+        prompt: Optional[str] = Field(
+            default="A beautiful sunset over a calm ocean",
+            description="A prompt for the video generation",
+        )
+
+    expected_output = (
+        IO.STRING,
+        {
+            "default": "A beautiful sunset over a calm ocean",
+            "tooltip": "A prompt for the video generation",
+            "multiline": True,
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.STRING, ModelWithStringField, "prompt", multiline=True
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_combo_input():
+    """Tests mapping a combo field."""
+
+    class MockEnum(str, Enum):
+        option_1 = "option 1"
+        option_2 = "option 2"
+        option_3 = "option 3"
+
+    class ModelWithComboField(BaseModel):
+        model_name: Optional[MockEnum] = Field("option 1", description="Model Name")
+
+    expected_output = (
+        IO.COMBO,
+        {
+            "options": ["option 1", "option 2", "option 3"],
+            "default": "option 1",
+            "tooltip": "Model Name",
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.COMBO, ModelWithComboField, "model_name", enum_type=MockEnum
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_combo_input_no_options():
+    """Tests mapping a combo field with no options."""
+
+    class ModelWithComboField(BaseModel):
+        model_name: Optional[str] = Field(description="Model Name")
+
+    expected_output = (
+        IO.COMBO,
+        {
+            "tooltip": "Model Name",
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.COMBO, ModelWithComboField, "model_name"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_image_input():
+    """Tests mapping an image field."""
+
+    class ModelWithImageField(BaseModel):
+        image: Optional[str] = Field(
+            default=None,
+            description="An image for the video generation",
+        )
+
+    expected_output = (
+        IO.IMAGE,
+        {
+            "default": None,
+            "tooltip": "An image for the video generation",
+        },
+    )
+
+    actual_output = model_field_to_node_input(IO.IMAGE, ModelWithImageField, "image")
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_node_input_no_description():
+    """Tests mapping a field with no description."""
+
+    class ModelWithNoDescriptionField(BaseModel):
+        field: Optional[str] = Field(default="default value")
+
+    expected_output = (
+        IO.STRING,
+        {
+            "default": "default value",
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.STRING, ModelWithNoDescriptionField, "field"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_node_input_no_default():
+    """Tests mapping a field with no default."""
+
+    class ModelWithNoDefaultField(BaseModel):
+        field: Optional[str] = Field(description="A field with no default")
+
+    expected_output = (
+        IO.STRING,
+        {
+            "tooltip": "A field with no default",
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.STRING, ModelWithNoDefaultField, "field"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_node_input_no_metadata():
+    """Tests mapping a field with no metadata or properties defined on the schema."""
+
+    class ModelWithNoMetadataField(BaseModel):
+        field: Optional[str] = Field()
+
+    expected_output = (
+        IO.STRING,
+        {},
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.STRING, ModelWithNoMetadataField, "field"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]
+
+
+def test_model_field_to_node_input_default_is_none():
+    """
+    Tests mapping a field with a default of `None`.
+    I.e., the default field should be included as the schema explicitly sets it to `None`.
+    """
+
+    class ModelWithNoneDefaultField(BaseModel):
+        field: Optional[str] = Field(
+            default=None, description="A field with a default of None"
+        )
+
+    expected_output = (
+        IO.STRING,
+        {
+            "default": None,
+            "tooltip": "A field with a default of None",
+        },
+    )
+
+    actual_output = model_field_to_node_input(
+        IO.STRING, ModelWithNoneDefaultField, "field"
+    )
+
+    assert actual_output[0] == expected_output[0]
+    assert actual_output[1] == expected_output[1]