mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-06-12 14:32:08 +08:00
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
![]() |
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
|