comfy_builder_core/
node.rs

1use crate::registry::Registerable;
2use indexmap::IndexMap;
3use pyo3::types::PyDict;
4use pyo3::{Bound, PyClass, PyResult, Python};
5use std::fmt::Display;
6
7pub enum DataType {
8    Int,
9    Float,
10    String,
11    Boolean,
12    Image,
13    Latent,
14    Mask,
15    Audio,
16    Noise,
17    Sampler,
18    Sigmas,
19    Guider,
20    Model,
21    Clip,
22    Vae,
23    Combo,
24    Conditioning,
25    Custom(&'static str),
26}
27
28impl From<&str> for DataType {
29    fn from(value: &str) -> Self {
30        match value {
31            // Primitive
32            "u8" | "u16" | "u32" | "u128" | "u64" | "usize" => DataType::Int,
33            "i8" | "i16" | "i32" | "i128" | "i64" | "isize" => DataType::Int,
34            "f32" | "f64" => DataType::Float,
35            "bool" => DataType::Boolean,
36            "String" => DataType::String,
37
38            // Tensors
39            "Tensor" => DataType::Image,
40            "Mask" => DataType::Mask,
41
42            // Hidden Inputs
43            "UniqueId" => DataType::String,
44            "Prompt" => DataType::String,
45            "DynPrompt" => DataType::String,
46            "ExtraPngInfo" => DataType::String,
47
48            kind => todo!("handle more types {:?}", kind),
49        }
50    }
51}
52
53impl Display for DataType {
54    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        let value = match self {
56            DataType::Int => "INT",
57            DataType::Combo => "COMBO",
58            DataType::Float => "FLOAT",
59            DataType::String => "STRING",
60            DataType::Boolean => "BOOLEAN",
61            DataType::Image => "IMAGE",
62            DataType::Latent => "LATENT",
63            DataType::Mask => "MASK",
64            DataType::Audio => "AUDIO",
65            DataType::Noise => "NOISE",
66            DataType::Sampler => "SAMPLER",
67            DataType::Sigmas => "SIGMAS",
68            DataType::Guider => "GUIDER",
69            DataType::Model => "MODEL",
70            DataType::Clip => "CLIP",
71            DataType::Vae => "VAE",
72            DataType::Conditioning => "CONDITIONING",
73            DataType::Custom(name) => name,
74        };
75
76        write!(formatter, "{}", value)
77    }
78}
79
80pub trait InputPort<'a>: From<Option<&'a Bound<'a, PyDict>>> {
81    fn get_inputs(py: Python<'a>) -> PyResult<Bound<'a, PyDict>>;
82}
83
84pub trait OutputPort<'a> {
85    fn get_outputs() -> IndexMap<&'static str, DataType>;
86
87    fn values() -> Vec<String> {
88        Self::get_outputs()
89            .into_values()
90            .map(|value| value.to_string())
91            .collect()
92    }
93
94    fn keys() -> Vec<&'static str> {
95        Self::get_outputs().into_keys().collect()
96    }
97}
98
99pub type NodeResult<'a, T> = Result<<T as Node<'a>>::Out, Box<dyn std::error::Error>>;
100
101pub trait Node<'a>: PyClass + Default + Sync + Send {
102    type In: InputPort<'a>;
103    type Out: OutputPort<'a>;
104
105    const CATEGORY: &'static str;
106    const DESCRIPTION: &'static str;
107
108    fn initialize_input(&'a self, kwargs: Option<&'a Bound<'a, PyDict>>) -> Self::In {
109        Self::In::from(kwargs)
110    }
111
112    fn new() -> Self {
113        Self::default()
114    }
115
116    fn execute(&self, input: Self::In) -> NodeResult<'a, Self>;
117}
118
119impl<'a, T: Node<'a>> Registerable for T {}
120
121pub trait EnumVariants: From<String> {
122    fn variants() -> Vec<&'static str>;
123}
comfy_builder_core/node.rs

comfy_builder_core/
node.rs
