From b8ceeea5ee02468b8a6b7bc7cb4ce0a52dbc2c1a Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Tue, 29 Apr 2025 11:53:10 -0700 Subject: [PATCH] Add utils to map from pydantic model fields to comfy node inputs (#30) --- comfy_api_nodes/README.md | 3 +- comfy_api_nodes/mapper_utils.py | 117 +++++++ .../comfy_api_nodes_test/mapper_utils_test.py | 297 ++++++++++++++++++ 3 files changed, 415 insertions(+), 2 deletions(-) create mode 100644 comfy_api_nodes/mapper_utils.py create mode 100644 tests-unit/comfy_api_nodes_test/mapper_utils_test.py diff --git a/comfy_api_nodes/README.md b/comfy_api_nodes/README.md index 86ba379f7..0ebecae63 100644 --- a/comfy_api_nodes/README.md +++ b/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 ``` - diff --git a/comfy_api_nodes/mapper_utils.py b/comfy_api_nodes/mapper_utils.py new file mode 100644 index 000000000..f8fd5632e --- /dev/null +++ b/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 diff --git a/tests-unit/comfy_api_nodes_test/mapper_utils_test.py b/tests-unit/comfy_api_nodes_test/mapper_utils_test.py new file mode 100644 index 000000000..69488f691 --- /dev/null +++ b/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]