comfyui_api/comfy/setter.rs
1use anyhow::{anyhow, Context};
2
3use crate::{
4 comfy::getter::{guess_node_mut, GetExt, Getter},
5 models::*,
6};
7
8/// Extension methods for `Prompt` to use Setters to set values on nodes.
9pub trait SetterExt<T, N>
10where
11 N: Node + 'static,
12{
13 /// Uses a heuristic to find a `Node` and set the value on it.
14 ///
15 /// # Inputs
16 ///
17 /// * `value` - The value to set.
18 ///
19 /// # Returns
20 ///
21 /// `Ok(())` on success, or an error if the node could not be found.
22 fn set<S>(&mut self, value: T) -> anyhow::Result<()>
23 where
24 S: Setter<T, N>;
25
26 /// Finds a `Node` leading into the given `output_node` and sets the value on it.
27 ///
28 /// # Inputs
29 ///
30 /// * `output_node` - The id of the node to path from.
31 ///
32 /// # Returns
33 ///
34 /// `Ok(())` on success, or an error if the node could not be found.
35 fn set_from<S>(&mut self, output_node: &str, value: T) -> anyhow::Result<()>
36 where
37 S: Setter<T, N>;
38
39 /// Sets the value on the node with id `node`.
40 ///
41 /// # Inputs
42 ///
43 /// * `node` - The id of the node to set the value on.
44 /// * `value` - The value to set.
45 ///
46 /// # Returns
47 ///
48 /// `Ok(())` on success, or an error if the node could not be found.
49 fn set_node<S>(&mut self, node: &str, value: T) -> anyhow::Result<()>
50 where
51 S: Setter<T, N>;
52}
53
54/// Extension methods for `Prompt` to set values on nodes.
55pub trait SetExt<N>
56where
57 N: Node + 'static,
58{
59 /// Uses a heuristic to find a `Node` and set the value on it.
60 ///
61 /// # Inputs
62 ///
63 /// * `f` - A function that takes a mutable reference to the node and returns a `Result`.
64 ///
65 /// # Returns
66 ///
67 /// `Ok(())` on success, or an error if the node could not be found.
68 fn set_with<F>(&mut self, f: F) -> anyhow::Result<()>
69 where
70 F: FnOnce(&mut N) -> anyhow::Result<()>;
71
72 /// Sets the value on the node with id `node`.
73 ///
74 /// # Inputs
75 ///
76 /// * `node` - The id of the node to set the value on.
77 ///
78 /// # Returns
79 ///
80 /// `Ok(())` on success, or an error if the node could not be found.
81 fn set_node_with<F>(&mut self, node: &str, f: F) -> anyhow::Result<()>
82 where
83 F: FnOnce(&mut N) -> anyhow::Result<()>;
84}
85
86impl<T, N: Node + 'static> SetterExt<T, N> for crate::models::Prompt {
87 fn set<S>(&mut self, value: T) -> anyhow::Result<()>
88 where
89 S: Setter<T, N>,
90 {
91 S::default().set(self, value)
92 }
93
94 fn set_from<S>(&mut self, output_node: &str, value: T) -> anyhow::Result<()>
95 where
96 S: Setter<T, N>,
97 {
98 S::default().set_from(self, output_node, value)
99 }
100
101 fn set_node<S>(&mut self, node: &str, value: T) -> anyhow::Result<()>
102 where
103 S: Setter<T, N>,
104 {
105 S::default().set_node(self, node, value)
106 }
107}
108
109impl<N: Node + 'static> SetExt<N> for crate::models::Prompt {
110 fn set_with<F>(&mut self, f: F) -> anyhow::Result<()>
111 where
112 F: FnOnce(&mut N) -> anyhow::Result<()>,
113 {
114 if let Some(node) = guess_node_mut::<N>(self, None) {
115 f(as_node_mut::<N>(node).context("Failed to cast node")?)
116 } else {
117 Err(anyhow!("Failed to find node"))
118 }
119 }
120
121 fn set_node_with<F>(&mut self, node: &str, f: F) -> anyhow::Result<()>
122 where
123 F: FnOnce(&mut N) -> anyhow::Result<()>,
124 {
125 f(self.get_typed_node_mut(node)?)
126 }
127}
128
129/// A trait for setting values on nodes.
130///
131/// This trait is implemented for types that can be used to set values on nodes.
132/// Usually, this trait does not need to be implemented directly, as it is implemented
133/// for all types that implement `Getter`.
134pub trait Setter<T, N>
135where
136 N: Node + 'static,
137 Self: Getter<T, N>,
138{
139 /// Uses a heuristic to find a `Node` and set the value on it.
140 ///
141 /// # Inputs
142 ///
143 /// * `prompt` - A mutable reference to a `Prompt`.
144 ///
145 /// # Returns
146 ///
147 /// `Ok(())` on success, or an error if the node could not be found.
148 fn set(&self, prompt: &mut crate::models::Prompt, value: T) -> anyhow::Result<()> {
149 let node = if let Some(node) = guess_node_mut::<N>(prompt, None) {
150 node
151 } else {
152 return Err(anyhow!("Failed to find node"));
153 };
154 self.set_value(node, value)
155 }
156
157 /// Finds a `Node` leading into the given `output_node` and sets the value on it.
158 ///
159 /// # Inputs
160 ///
161 /// * `prompt` - A mutable reference to a `Prompt`.
162 ///
163 /// # Returns
164 ///
165 /// `Ok(())` on success, or an error if the node could not be found.
166 fn set_from(
167 &self,
168 prompt: &mut crate::models::Prompt,
169 output_node: &str,
170 value: T,
171 ) -> anyhow::Result<()> {
172 let node = if let Some(node) = Self::find_node(prompt, Some(output_node)) {
173 prompt
174 .get_node_by_id_mut(&node)
175 .context("Failed to find node")?
176 } else {
177 return Err(anyhow!("Failed to find node"));
178 };
179 self.set_value(node, value)
180 }
181
182 /// Sets the value on the node with id `node`.
183 ///
184 /// # Inputs
185 ///
186 /// * `prompt` - A mutable reference to a `Prompt`.
187 /// * `node` - The id of the node to set the value on.
188 ///
189 /// # Returns
190 ///
191 /// `Ok(())` on success, or an error if the node could not be found.
192 fn set_node(
193 &self,
194 prompt: &mut crate::models::Prompt,
195 node: &str,
196 value: T,
197 ) -> anyhow::Result<()> {
198 let node = prompt.get_node_by_id_mut(node).unwrap();
199 self.set_value(node, value)
200 }
201
202 /// Sets the value on the given `Node`.
203 ///
204 /// # Inputs
205 ///
206 /// * `node` - A mutable reference to a `Node`.
207 ///
208 /// # Returns
209 ///
210 /// `Ok(())` on success, or an error if the node could not be found.
211 fn set_value(&self, node: &mut dyn Node, value: T) -> anyhow::Result<()> {
212 *self.get_value_mut(node)? = value;
213 Ok(())
214 }
215}
216
217impl<G, T, N> Setter<T, N> for G
218where
219 G: Getter<T, N>,
220 N: Node + 'static,
221{
222}