Skip to content

Commit 2a89781

Browse files
authored
Combobox support for control panel (#18)
Add support for combo-box in control-panel
1 parent a4c564d commit 2a89781

File tree

2 files changed

+180
-73
lines changed

2 files changed

+180
-73
lines changed

src/comfystream/client.py

Lines changed: 69 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ async def run_prompt(self, prompt_index: int):
4545

4646
async def cleanup(self):
4747
async with self.cleanup_lock:
48-
for task in self.running_prompts.values():
48+
tasks_to_cancel = list(self.running_prompts.values())
49+
for task in tasks_to_cancel:
4950
task.cancel()
5051
try:
5152
await task
@@ -54,7 +55,11 @@ async def cleanup(self):
5455
self.running_prompts.clear()
5556

5657
if self.comfy_client.is_running:
57-
await self.comfy_client.__aexit__()
58+
try:
59+
await self.comfy_client.__aexit__()
60+
except Exception as e:
61+
logger.error(f"Error during ComfyClient cleanup: {e}")
62+
5863

5964
await self.cleanup_queues()
6065
logger.info("Client cleanup complete")
@@ -110,77 +115,105 @@ async def get_available_nodes(self):
110115
for node_id, node in prompt.items()
111116
}
112117
nodes_info = {}
113-
118+
114119
# Only process nodes until we've found all the ones we need
115120
for class_type, node_class in nodes.NODE_CLASS_MAPPINGS.items():
116121
if not remaining_nodes: # Exit early if we've found all needed nodes
117122
break
118-
123+
119124
if class_type not in needed_class_types:
120125
continue
121-
126+
122127
# Get metadata for this node type (same as original get_node_metadata)
123128
input_data = node_class.INPUT_TYPES() if hasattr(node_class, 'INPUT_TYPES') else {}
124129
input_info = {}
125-
130+
126131
# Process required inputs
127132
if 'required' in input_data:
128133
for name, value in input_data['required'].items():
129-
if isinstance(value, tuple) and len(value) == 2:
130-
input_type, config = value
131-
input_info[name] = {
132-
'type': input_type,
133-
'required': True,
134-
'min': config.get('min', None),
135-
'max': config.get('max', None),
136-
'widget': config.get('widget', None)
137-
}
134+
if isinstance(value, tuple):
135+
if len(value) == 1 and isinstance(value[0], list):
136+
# Handle combo box case where value is ([option1, option2, ...],)
137+
input_info[name] = {
138+
'type': 'combo',
139+
'value': value[0], # The list of options becomes the value
140+
}
141+
elif len(value) == 2:
142+
input_type, config = value
143+
input_info[name] = {
144+
'type': input_type,
145+
'required': True,
146+
'min': config.get('min', None),
147+
'max': config.get('max', None),
148+
'widget': config.get('widget', None)
149+
}
150+
elif len(value) == 1:
151+
# Handle simple type case like ('IMAGE',)
152+
input_info[name] = {
153+
'type': value[0]
154+
}
138155
else:
139156
logger.error(f"Unexpected structure for required input {name}: {value}")
140-
141-
# Process optional inputs
157+
158+
# Process optional inputs with same logic
142159
if 'optional' in input_data:
143160
for name, value in input_data['optional'].items():
144-
if isinstance(value, tuple) and len(value) == 2:
145-
input_type, config = value
146-
input_info[name] = {
147-
'type': input_type,
148-
'required': False,
149-
'min': config.get('min', None),
150-
'max': config.get('max', None),
151-
'widget': config.get('widget', None)
152-
}
161+
if isinstance(value, tuple):
162+
if len(value) == 1 and isinstance(value[0], list):
163+
# Handle combo box case where value is ([option1, option2, ...],)
164+
input_info[name] = {
165+
'type': 'combo',
166+
'value': value[0], # The list of options becomes the value
167+
}
168+
elif len(value) == 2:
169+
input_type, config = value
170+
input_info[name] = {
171+
'type': input_type,
172+
'required': False,
173+
'min': config.get('min', None),
174+
'max': config.get('max', None),
175+
'widget': config.get('widget', None)
176+
}
177+
elif len(value) == 1:
178+
# Handle simple type case like ('IMAGE',)
179+
input_info[name] = {
180+
'type': value[0]
181+
}
153182
else:
154183
logger.error(f"Unexpected structure for optional input {name}: {value}")
155-
184+
156185
# Now process any nodes in our prompt that use this class_type
157186
for node_id in list(remaining_nodes):
158187
node = prompt[node_id]
159188
if node.get('class_type') != class_type:
160189
continue
161-
190+
162191
node_info = {
163192
'class_type': class_type,
164193
'inputs': {}
165194
}
166-
195+
167196
if 'inputs' in node:
168197
for input_name, input_value in node['inputs'].items():
198+
input_metadata = input_info.get(input_name, {})
169199
node_info['inputs'][input_name] = {
170200
'value': input_value,
171-
'type': input_info.get(input_name, {}).get('type', 'unknown'),
172-
'min': input_info.get(input_name, {}).get('min', None),
173-
'max': input_info.get(input_name, {}).get('max', None),
174-
'widget': input_info.get(input_name, {}).get('widget', None)
201+
'type': input_metadata.get('type', 'unknown'),
202+
'min': input_metadata.get('min', None),
203+
'max': input_metadata.get('max', None),
204+
'widget': input_metadata.get('widget', None)
175205
}
176-
206+
# For combo type inputs, include the list of options
207+
if input_metadata.get('type') == 'combo':
208+
node_info['inputs'][input_name]['value'] = input_metadata.get('value', [])
209+
177210
nodes_info[node_id] = node_info
178211
remaining_nodes.remove(node_id)
179212

180213
all_prompts_nodes_info[prompt_index] = nodes_info
181-
214+
182215
return all_prompts_nodes_info
183-
216+
184217
except Exception as e:
185218
logger.error(f"Error getting node info: {str(e)}")
186219
return {}

ui/src/components/control-panel.tsx

Lines changed: 111 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,15 @@ import React, { useState, useEffect } from "react";
44
import { usePeerContext } from "@/context/peer-context";
55
import { usePrompt } from "./settings";
66

7-
type InputValue = string | number | boolean;
7+
type InputValue = string | number | boolean | string[];
88

99
interface InputInfo {
1010
value: InputValue;
1111
type: string;
1212
min?: number;
1313
max?: number;
1414
widget?: string;
15+
options?: string[];
1516
}
1617

1718
interface NodeInfo {
@@ -45,19 +46,32 @@ const InputControl = ({
4546
value: string;
4647
onChange: (value: string) => void;
4748
}) => {
48-
if (input.widget === "combo") {
49+
if (input.widget === "combo" || input.type === "combo") {
50+
// Get options from either the options field or value field
51+
const options = input.options
52+
? input.options
53+
: Array.isArray(input.value)
54+
? input.value
55+
: typeof input.value === 'string'
56+
? [input.value]
57+
: [];
58+
59+
// If no value is selected, select the first option by default
60+
const currentValue = value || options[0] || '';
61+
4962
return (
5063
<select
51-
value={value}
52-
onChange={(e) => onChange(e.target.value)}
64+
value={currentValue}
65+
onChange={(e) => {
66+
onChange(e.target.value);
67+
}}
5368
className="p-2 border rounded w-full"
5469
>
55-
{Array.isArray(input.value) &&
56-
input.value.map((option: string) => (
57-
<option key={option} value={option}>
58-
{option}
59-
</option>
60-
))}
70+
{options.map((option: string) => (
71+
<option key={option} value={option}>
72+
{option}
73+
</option>
74+
))}
6175
</select>
6276
);
6377
}
@@ -100,6 +114,9 @@ const InputControl = ({
100114
className="p-2 border rounded w-full"
101115
/>
102116
);
117+
// Handle combo in the main combo block above
118+
case "combo":
119+
return InputControl({ input: { ...input, widget: "combo" }, value, onChange });
103120
default:
104121
console.warn(`Unhandled input type: ${input.type}`); // Debug log
105122
return (
@@ -196,12 +213,33 @@ export const ControlPanel = ({
196213
]
197214
: null;
198215
if (!currentInput || !currentPrompts) return;
216+
217+
// Don't send updates if this is a combo and we haven't selected a value yet
218+
if (currentInput.widget === "combo" && !panelState.value) return;
199219

200220
let isValidValue = true;
201221
let processedValue: InputValue = panelState.value;
202222

203-
// Validate and process value based on type
204-
switch (currentInput.type.toLowerCase()) {
223+
// For combo inputs, use the value directly
224+
if (currentInput.widget === "combo" || currentInput.type === "combo") {
225+
// Get options from either the options field or value field
226+
const options = currentInput.options
227+
? currentInput.options
228+
: Array.isArray(currentInput.value)
229+
? currentInput.value as string[]
230+
: typeof currentInput.value === 'string'
231+
? [currentInput.value as string]
232+
: [];
233+
234+
// If no value is selected and we have options, use the first option
235+
const validValue = panelState.value || options[0] || '';
236+
237+
// Validate that the value is in the options list
238+
isValidValue = options.includes(validValue);
239+
processedValue = validValue;
240+
} else {
241+
// Validate and process value based on type
242+
switch (currentInput.type.toLowerCase()) {
205243
case "number":
206244
isValidValue =
207245
/^-?\d*\.?\d*$/.test(panelState.value) && panelState.value !== "";
@@ -217,13 +255,9 @@ export const ControlPanel = ({
217255
processedValue = panelState.value;
218256
break;
219257
default:
220-
if (currentInput.widget === "combo") {
221-
isValidValue = panelState.value !== "";
222-
processedValue = panelState.value;
223-
} else {
224-
isValidValue = panelState.value !== "";
225-
processedValue = panelState.value;
226-
}
258+
isValidValue = panelState.value !== "";
259+
processedValue = panelState.value;
260+
}
227261
}
228262

229263
const hasRequiredFields =
@@ -261,9 +295,32 @@ export const ControlPanel = ({
261295
}
262296
const updatedPrompt = JSON.parse(JSON.stringify(prompt)); // Deep clone
263297
if (updatedPrompt[panelState.nodeId]?.inputs) {
264-
updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] =
265-
processedValue;
266-
hasUpdated = true;
298+
// Ensure we're not overwriting with an invalid value
299+
const currentVal = updatedPrompt[panelState.nodeId].inputs[panelState.fieldName];
300+
const input = availableNodes[promptIdxToUpdate][panelState.nodeId]?.inputs[panelState.fieldName];
301+
302+
if (input?.widget === 'combo' || input?.type === 'combo') {
303+
// Get options from either the options field or value field
304+
const options = input.options
305+
? input.options
306+
: Array.isArray(input.value)
307+
? input.value as string[]
308+
: typeof input.value === 'string'
309+
? [input.value as string]
310+
: [];
311+
312+
// If no value is selected and we have options, use the first option
313+
const validValue = (processedValue as string) || options[0] || '';
314+
315+
// Only update if it's a valid combo value
316+
if (options.includes(validValue)) {
317+
updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] = validValue;
318+
hasUpdated = true;
319+
}
320+
} else {
321+
updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] = processedValue;
322+
hasUpdated = true;
323+
}
267324
}
268325
return updatedPrompt;
269326
},
@@ -312,8 +369,13 @@ export const ControlPanel = ({
312369
if (input.type.toLowerCase() === "boolean") {
313370
return (!!input.value).toString();
314371
}
315-
if (input.widget === "combo" && Array.isArray(input.value)) {
316-
return input.value[0]?.toString() || "";
372+
if (input.widget === "combo") {
373+
const options = Array.isArray(input.value)
374+
? input.value as string[]
375+
: typeof input.value === 'string'
376+
? [input.value as string]
377+
: [];
378+
return options[0] || "";
317379
}
318380
return input.value?.toString() || "0";
319381
};
@@ -327,11 +389,19 @@ export const ControlPanel = ({
327389
selectedField
328390
];
329391
if (input) {
330-
const initialValue = getInitialValue(input);
331-
onStateChange({
332-
fieldName: selectedField,
333-
value: initialValue,
334-
});
392+
// For combo fields, don't set an initial value to prevent auto-update from firing
393+
if (input.widget === "combo") {
394+
onStateChange({
395+
fieldName: selectedField,
396+
value: "",
397+
});
398+
} else {
399+
const initialValue = getInitialValue(input);
400+
onStateChange({
401+
fieldName: selectedField,
402+
value: initialValue,
403+
});
404+
}
335405
} else {
336406
onStateChange({ fieldName: selectedField });
337407
}
@@ -357,7 +427,7 @@ export const ControlPanel = ({
357427
onStateChange({
358428
nodeId: e.target.value,
359429
fieldName: "",
360-
value: "0",
430+
value: "", // Start with empty value to prevent auto-update from firing
361431
});
362432
}}
363433
className="p-2 border rounded"
@@ -387,16 +457,20 @@ export const ControlPanel = ({
387457
typeof info.type === "string"
388458
? info.type.toLowerCase()
389459
: String(info.type).toLowerCase();
390-
return (
391-
["boolean", "number", "float", "int", "string"].includes(
392-
type,
393-
) || info.widget === "combo"
394-
);
460+
return [
461+
"boolean",
462+
"number",
463+
"float",
464+
"int",
465+
"string",
466+
"combo",
467+
].includes(
468+
type,
469+
) || info.widget === "combo";
395470
})
396471
.map(([field, info]) => (
397472
<option key={field} value={field}>
398-
{field} ({info.type}
399-
{info.widget ? ` - ${info.widget}` : ""})
473+
{field} ({info.type})
400474
</option>
401475
))}
402476
</select>

0 commit comments

Comments
 (0)