comfy-builder-core 0.0.2

A system for building custom nodes for ComfyUI in Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use crate::registry::Registerable;
use indexmap::IndexMap;
use pyo3::types::PyDict;
use pyo3::{Bound, PyClass, PyResult, Python};
use std::fmt::Display;

pub enum DataType {
    Int,
    Float,
    String,
    Boolean,
    Image,
    Latent,
    Mask,
    Audio,
    Noise,
    Sampler,
    Sigmas,
    Guider,
    Model,
    Clip,
    Vae,
    Combo,
    Conditioning,
    Custom(&'static str),
}

impl From<&str> for DataType {
    fn from(value: &str) -> Self {
        match value {
            // Primitive
            "u8" | "u16" | "u32" | "u128" | "u64" | "usize" => DataType::Int,
            "i8" | "i16" | "i32" | "i128" | "i64" | "isize" => DataType::Int,
            "f32" | "f64" => DataType::Float,
            "bool" => DataType::Boolean,
            "String" => DataType::String,

            // Tensors
            "Tensor" => DataType::Image,
            "Mask" => DataType::Mask,

            // Hidden Inputs
            "UniqueId" => DataType::String,
            "Prompt" => DataType::String,
            "DynPrompt" => DataType::String,
            "ExtraPngInfo" => DataType::String,

            kind => todo!("handle more types {:?}", kind),
        }
    }
}

impl Display for DataType {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let value = match self {
            DataType::Int => "INT",
            DataType::Combo => "COMBO",
            DataType::Float => "FLOAT",
            DataType::String => "STRING",
            DataType::Boolean => "BOOLEAN",
            DataType::Image => "IMAGE",
            DataType::Latent => "LATENT",
            DataType::Mask => "MASK",
            DataType::Audio => "AUDIO",
            DataType::Noise => "NOISE",
            DataType::Sampler => "SAMPLER",
            DataType::Sigmas => "SIGMAS",
            DataType::Guider => "GUIDER",
            DataType::Model => "MODEL",
            DataType::Clip => "CLIP",
            DataType::Vae => "VAE",
            DataType::Conditioning => "CONDITIONING",
            DataType::Custom(name) => name,
        };

        write!(formatter, "{}", value)
    }
}

pub trait InputPort<'a>: From<Option<&'a Bound<'a, PyDict>>> {
    fn get_inputs(py: Python<'a>) -> PyResult<Bound<'a, PyDict>>;
}

pub trait OutputPort<'a> {
    fn get_outputs() -> IndexMap<&'static str, DataType>;

    fn values() -> Vec<String> {
        Self::get_outputs()
            .into_values()
            .map(|value| value.to_string())
            .collect()
    }

    fn keys() -> Vec<&'static str> {
        Self::get_outputs().into_keys().collect()
    }
}

pub type NodeResult<'a, T> = Result<<T as Node<'a>>::Out, Box<dyn std::error::Error>>;

pub trait Node<'a>: PyClass + Default + Sync + Send {
    type In: InputPort<'a>;
    type Out: OutputPort<'a>;

    const CATEGORY: &'static str;
    const DESCRIPTION: &'static str;

    fn initialize_input(&'a self, kwargs: Option<&'a Bound<'a, PyDict>>) -> Self::In {
        Self::In::from(kwargs)
    }

    fn new() -> Self {
        Self::default()
    }

    fn execute(&self, input: Self::In) -> NodeResult<'a, Self>;
}

impl<'a, T: Node<'a>> Registerable for T {}

pub trait EnumVariants: From<String> {
    fn variants() -> Vec<&'static str>;
}