Skip to content

Commit be04c77

Browse files
jfcantujfcantu
committed
Allow using ckpt_name from Autoselect node, fix list refreshing
1 parent 6507679 commit be04c77

File tree

3 files changed

+136
-6
lines changed

3 files changed

+136
-6
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "ComfyUI-Prompt-Companion"
7-
version = "0.0.7"
7+
version = "0.0.8"
88
description = "A node that lets you save and reuse parts of prompts (embeddings, quality keywords, and so on.)"
99
authors = [
1010
{name = "John Cantu", email = "jfcantu@gmail.com"}

‎src/prompt_companion_node.py‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -773,9 +773,9 @@ class PromptCompanionAutoselectGroups:
773773
"""
774774

775775
# ComfyUI node metadata
776-
RETURN_TYPES = ("PROMPT_ADDITION",)
777-
RETURN_NAMES = ("prompt_addition",)
778-
OUTPUT_TOOLTIPS = ("Prompt addition data that can be connected to other nodes",)
776+
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints") if folder_paths else ["test_model.safetensors"], "PROMPT_ADDITION")
777+
RETURN_NAMES = ("ckpt_name", "prompt_addition")
778+
OUTPUT_TOOLTIPS = ("Validated checkpoint name", "Prompt addition data that can be connected to other nodes")
779779
FUNCTION = "autoselect_groups"
780780
CATEGORY = "jfc"
781781

@@ -813,7 +813,7 @@ def autoselect_groups(
813813
combine_mode: str,
814814
ckpt_name: str,
815815
prompt_addition: Optional[PromptAdditionInput] = None,
816-
) -> Tuple['PromptAdditionInput']:
816+
) -> Tuple[str, 'PromptAdditionInput']:
817817
"""Auto-select prompt groups based on checkpoint trigger words."""
818818

819819
# Find matching groups based on trigger words
@@ -875,7 +875,38 @@ def autoselect_groups(
875875
elif input_negative:
876876
final_negative = input_negative
877877

878-
return (PromptAdditionInput(final_positive, final_negative),)
878+
# Validate and ensure ckpt_name is compatible with Load Checkpoint
879+
validated_ckpt_name = self._validate_checkpoint_name(ckpt_name)
880+
881+
return (validated_ckpt_name, PromptAdditionInput(final_positive, final_negative))
882+
883+
def _validate_checkpoint_name(self, ckpt_name: str) -> str:
884+
"""
885+
Validate that the checkpoint name is in the available checkpoints list.
886+
This ensures compatibility with Load Checkpoint nodes.
887+
888+
Args:
889+
ckpt_name: The checkpoint name to validate
890+
891+
Returns:
892+
The validated checkpoint name, or the first available if invalid
893+
"""
894+
if not folder_paths:
895+
# Fallback for testing environments
896+
return ckpt_name or "test_model.safetensors"
897+
898+
available_checkpoints = folder_paths.get_filename_list("checkpoints")
899+
900+
if not available_checkpoints:
901+
# No checkpoints available, return as-is
902+
return ckpt_name
903+
904+
if ckpt_name in available_checkpoints:
905+
# Valid checkpoint name
906+
return ckpt_name
907+
else:
908+
# Invalid checkpoint name, return the first available one
909+
return available_checkpoints[0]
879910

880911
def _group_matches_checkpoint(self, group, ckpt_name: str) -> bool:
881912
"""Check if a group's trigger words match the checkpoint name."""

‎src/web/extension.js‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,13 +98,47 @@ app.registerExtension({
9898
// Update prompt addition names
9999
if (node.promptAdditionNameWidget) {
100100
const additionNames = Object.keys(this.promptAdditions);
101+
const currentValue = node.promptAdditionNameWidget.value;
101102
node.promptAdditionNameWidget.options.values = additionNames.length > 0 ? additionNames : [""];
103+
104+
// Trigger a redraw of the widget dropdown
105+
if (node.promptAdditionNameWidget.inputEl) {
106+
// Clear existing options
107+
node.promptAdditionNameWidget.inputEl.innerHTML = "";
108+
// Rebuild options
109+
node.promptAdditionNameWidget.options.values.forEach(value => {
110+
const option = document.createElement("option");
111+
option.value = value;
112+
option.textContent = value;
113+
if (value === currentValue) {
114+
option.selected = true;
115+
}
116+
node.promptAdditionNameWidget.inputEl.appendChild(option);
117+
});
118+
}
102119
}
103120

104121
// Update prompt group names
105122
if (node.promptAdditionGroupWidget) {
106123
const groupNames = this.promptGroups.map(g => g.name);
124+
const currentValue = node.promptAdditionGroupWidget.value;
107125
node.promptAdditionGroupWidget.options.values = groupNames.length > 0 ? groupNames : [""];
126+
127+
// Trigger a redraw of the widget dropdown
128+
if (node.promptAdditionGroupWidget.inputEl) {
129+
// Clear existing options
130+
node.promptAdditionGroupWidget.inputEl.innerHTML = "";
131+
// Rebuild options
132+
node.promptAdditionGroupWidget.options.values.forEach(value => {
133+
const option = document.createElement("option");
134+
option.value = value;
135+
option.textContent = value;
136+
if (value === currentValue) {
137+
option.selected = true;
138+
}
139+
node.promptAdditionGroupWidget.inputEl.appendChild(option);
140+
});
141+
}
108142
}
109143

110144
// Handle widget visibility based on addition_type
@@ -540,6 +574,57 @@ app.registerExtension({
540574
node.promptAdditions = {};
541575
node.promptAdditionNames = {};
542576
node.promptGroups = [];
577+
578+
// Add refreshWidgetOptions method to all nodes that might have dropdown widgets
579+
if (!node.refreshWidgetOptions) {
580+
node.refreshWidgetOptions = function (node) {
581+
// Update prompt addition names (for nodes with prompt_addition_name dropdown)
582+
if (node.promptAdditionNameWidget) {
583+
const additionNames = Object.keys(this.promptAdditions);
584+
const currentValue = node.promptAdditionNameWidget.value;
585+
node.promptAdditionNameWidget.options.values = additionNames.length > 0 ? additionNames : [""];
586+
587+
// Trigger a redraw of the widget dropdown
588+
if (node.promptAdditionNameWidget.inputEl) {
589+
// Clear existing options
590+
node.promptAdditionNameWidget.inputEl.innerHTML = "";
591+
// Rebuild options
592+
node.promptAdditionNameWidget.options.values.forEach(value => {
593+
const option = document.createElement("option");
594+
option.value = value;
595+
option.textContent = value;
596+
if (value === currentValue) {
597+
option.selected = true;
598+
}
599+
node.promptAdditionNameWidget.inputEl.appendChild(option);
600+
});
601+
}
602+
}
603+
604+
// Update prompt group names (for nodes with prompt_addition_group dropdown)
605+
if (node.promptAdditionGroupWidget) {
606+
const groupNames = this.promptGroups.map(g => g.name);
607+
const currentValue = node.promptAdditionGroupWidget.value;
608+
node.promptAdditionGroupWidget.options.values = groupNames.length > 0 ? groupNames : [""];
609+
610+
// Trigger a redraw of the widget dropdown
611+
if (node.promptAdditionGroupWidget.inputEl) {
612+
// Clear existing options
613+
node.promptAdditionGroupWidget.inputEl.innerHTML = "";
614+
// Rebuild options
615+
node.promptAdditionGroupWidget.options.values.forEach(value => {
616+
const option = document.createElement("option");
617+
option.value = value;
618+
option.textContent = value;
619+
if (value === currentValue) {
620+
option.selected = true;
621+
}
622+
node.promptAdditionGroupWidget.inputEl.appendChild(option);
623+
});
624+
}
625+
}
626+
};
627+
}
543628

544629
// Load prompt data
545630
const promptAdditions = await ApiOperations.getPromptAdditions();
@@ -579,13 +664,27 @@ app.registerExtension({
579664
node.promptAdditions = newPromptAdditionsObject;
580665
node.promptAdditionNames = Object.keys(newPromptAdditionsObject);
581666
node.promptGroups = event.detail.prompt_groups || [];
667+
668+
// Refresh dropdown widgets for nodes without updatePromptData method
669+
if (node.refreshWidgetOptions) {
670+
node.refreshWidgetOptions(node);
671+
}
582672
}
583673
};
584674

585675
api.addEventListener("prompt-companion.addition-list", eventListener);
586676

587677
// Store the event listener for cleanup
588678
node.eventListener = eventListener;
679+
680+
// Set up widget references for nodes that have dropdown widgets
681+
if (node.comfyClass === "PromptCompanionSingleAddition" || node.comfyClass === "PromptCompanion") {
682+
node.promptAdditionNameWidget = node.widgets?.find((w) => w.name === "prompt_addition_name");
683+
}
684+
685+
if (node.comfyClass === "PromptCompanionPromptGroup" || node.comfyClass === "PromptCompanion") {
686+
node.promptAdditionGroupWidget = node.widgets?.find((w) => w.name === "prompt_addition_group");
687+
}
589688
}
590689

591690
if (node.comfyClass == "PromptCompanion") {

0 commit comments

Comments
 (0)