From c3a0a87755bb65a5cc73b75568784f0652a232db Mon Sep 17 00:00:00 2001 From: Satadal Dhara Date: Thu, 18 Sep 2025 09:34:20 +0530 Subject: [PATCH] updated --- nodes/prompt_machine.py | 163 ++++++++++++++++++--------- nodes/resolution_machine.py | 45 +++++--- web/prompt_machine.js | 68 ++++++++---- web/resolution_machine.js | 212 ++++++++++++------------------------ 4 files changed, 256 insertions(+), 232 deletions(-) diff --git a/nodes/prompt_machine.py b/nodes/prompt_machine.py index 2ca7110..e98cfc2 100644 --- a/nodes/prompt_machine.py +++ b/nodes/prompt_machine.py @@ -1,49 +1,115 @@ +# prompt_machine_node.py import os import csv from server import PromptServer from aiohttp import web +# Directory for CSV files (adjust if needed) CSV_DIR = os.path.join(os.path.dirname(__file__), "..", "prompts") def list_csv_files(): - return sorted([f for f in os.listdir(CSV_DIR) if f.endswith(".csv")]) + try: + return sorted([f for f in os.listdir(CSV_DIR) if f.endswith(".csv")]) + except Exception as e: + print(f"[PromptMachine] list_csv_files error: {e}") + return [] def read_names_from_csv(filename): + """Return list of names (stripped) from CSV's 'name' column.""" + if not filename: + return [] path = os.path.join(CSV_DIR, filename) if not os.path.exists(path): + print(f"[PromptMachine] CSV not found: {path}") return [] - with open(path, "r", encoding="utf-8-sig") as f: - reader = csv.DictReader(f) - if not reader.fieldnames: - return [] - field_map = {fn.strip().lower(): fn for fn in reader.fieldnames} - if "name" not in field_map: - return [] - name_field = field_map["name"] - return [row[name_field] for row in reader if row.get(name_field)] + try: + with open(path, "r", encoding="utf-8-sig") as f: + reader = csv.DictReader(f) + if not reader.fieldnames: + print(f"[PromptMachine] CSV has no header: {filename}") + return [] + field_map = {fn.strip().lower(): fn for fn in reader.fieldnames} + if "name" not in field_map: + print(f"[PromptMachine] CSV missing 'name' column: {filename}") + return [] + name_field = field_map["name"] + names = [] + for row in reader: + raw = row.get(name_field) + if raw is None: + continue + val = raw.strip() + if val: + names.append(val) + print(f"[PromptMachine] Loaded {len(names)} names from {filename}") + return names + except Exception as e: + print(f"[PromptMachine] read_names_from_csv error for {filename}: {e}") + return [] + +def read_prompt_row(csv_file, name): + """Return (positive, negative, note) for a given csv_file and name. + Matching is whitespace-trimmed and exact on content (case-sensitive by default).""" + if not csv_file or csv_file == "None" or not name or name == "None": + return ("", "", "") + path = os.path.join(CSV_DIR, csv_file) + if not os.path.exists(path): + print(f"[PromptMachine] CSV file not found: {path}") + return ("", "", "") + try: + with open(path, "r", encoding="utf-8-sig") as f: + reader = csv.DictReader(f) + if not reader.fieldnames: + print(f"[PromptMachine] CSV has no header: {csv_file}") + return ("", "", "") + field_map = {fn.strip().lower(): fn for fn in reader.fieldnames} + if "name" not in field_map: + print(f"[PromptMachine] 'name' column missing in CSV: {csv_file}") + return ("", "", "") + name_field = field_map["name"] + pos_field = field_map.get("positive", "positive") + neg_field = field_map.get("negative", "negative") + note_field = field_map.get("note", "note") + + target = name.strip() + for row in reader: + cell = row.get(name_field) + if cell is None: + continue + if cell.strip() == target: + pos = (row.get(pos_field) or "").strip() + neg = (row.get(neg_field) or "").strip() + note = (row.get(note_field) or "").strip() + return (pos, neg, note) + except Exception as e: + print(f"[PromptMachine] read_prompt_row error for {csv_file}, {name}: {e}") + return ("", "", "") + class Prompt_Machine: - def __init__(self): - csvs = list_csv_files() - if csvs: - self.default_csv = csvs[0] - names = read_names_from_csv(self.default_csv) - self.default_name = names[0] if names else "None" - else: - self.default_csv = "None" - self.default_name = "None" - + """ + NOTE: name is a free STRING (not a static dropdown) to avoid 'Value not in list' validation errors. + The frontend helper will still provide dropdown-like selection and write the chosen name into this STRING. + """ @classmethod def INPUT_TYPES(cls): csvs = list_csv_files() default_csv = csvs[0] if csvs else "None" - names = read_names_from_csv(default_csv) if csvs else [] + # default_name set to first name of default CSV if available + names = read_names_from_csv(default_csv) if default_csv != "None" else [] default_name = names[0] if names else "None" + print(f"[PromptMachine] INIT - default_csv={default_csv}, default_name={default_name}") + return { "required": { - "csv_file": (csvs if csvs else ["None"], {"default": default_csv}), - "name": (names if names else ["None"], {"default": default_name}), + "csv_file": (csvs if csvs else ["None"], { + "default": default_csv, + # not using get_dynamic_inputs; frontend helper writes the STRING + "refresh": False, + }), + # make 'name' a free text field (STRING) — avoids list validation problems + "name": ("STRING", {"default": default_name}), } } @@ -53,39 +119,36 @@ class Prompt_Machine: CATEGORY = "SATA_UtilityNode" def get_prompts(self, csv_file, name): - if csv_file == "None" or name == "None": - return ("", "", "null") - path = os.path.join(CSV_DIR, csv_file) - if not os.path.exists(path): - print(f"[PromptMachine] CSV file not found: {path}") - return ("", "", "null") - with open(path, "r", encoding="utf-8-sig") as f: - reader = csv.DictReader(f) - if not reader.fieldnames: - print(f"[PromptMachine] CSV has no header: {csv_file}") - return ("", "", "null") - print(f"[PromptMachine] CSV header in {csv_file}: {reader.fieldnames}") - field_map = {fn.strip().lower(): fn for fn in reader.fieldnames} - if "name" not in field_map: - print(f"[PromptMachine] 'name' column missing in CSV: {csv_file}") - return ("", "", "null") - name_field = field_map["name"] - pos_field = field_map.get("positive", "positive") - neg_field = field_map.get("negative", "negative") - note_field = field_map.get("note", "note") - for row in reader: - if row.get(name_field) == name: - return (row.get(pos_field, ""), row.get(neg_field, ""), row.get(note_field, "")) - return ("", "", "","null") + print(f"[PromptMachine] get_prompts - csv_file={csv_file}, name={name}") + pos, neg, note = read_prompt_row(csv_file, name) + print(f"[PromptMachine] get_prompts -> pos_len={len(pos)}, neg_len={len(neg)}, note_len={len(note)}") + return (pos, neg, note) -# REST endpoints for frontend refresh +# ---------------- REST API ---------------- + @PromptServer.instance.routes.get("/sata/prompt_machine/csvs") async def list_csvs(request): - return web.json_response({"csvs": list_csv_files()}) + """Return all available CSV files""" + files = list_csv_files() + return web.json_response({"csvs": files}) @PromptServer.instance.routes.get("/sata/prompt_machine/names") async def list_names(request): + """Return names from a selected CSV""" csv_file = request.query.get("csv") - return web.json_response({"names": read_names_from_csv(csv_file)} if csv_file else {"names": []}) + names = read_names_from_csv(csv_file) if csv_file else [] + return web.json_response({"names": names}) + + +@PromptServer.instance.routes.get("/sata/prompt_machine/get") +async def get_prompt(request): + """ + Return the prompt triple for immediate frontend consumption. + Example: /sata/prompt_machine/get?csv=characters.csv&name=Bob + """ + csv_file = request.query.get("csv") + name = request.query.get("name") + pos, neg, note = read_prompt_row(csv_file, name) + return web.json_response({"positive": pos, "negative": neg, "note": note}) \ No newline at end of file diff --git a/nodes/resolution_machine.py b/nodes/resolution_machine.py index b7c8535..a999e17 100644 --- a/nodes/resolution_machine.py +++ b/nodes/resolution_machine.py @@ -1,10 +1,15 @@ -# resolution_machine.py from server import PromptServer from aiohttp import web class Resolution_Machine: CATEGORY = "SATA_UtilityNode" + # persist last-used selections + last_model = None + last_resolution = None + last_width = None + last_height = None + @classmethod def _get_config(cls): cfg = { @@ -102,31 +107,32 @@ class Resolution_Machine: @classmethod def INPUT_TYPES(cls, context=None): - """Expose model list and a default resolution list (for the default model). - Frontend will dynamically set the resolution options for other models using the REST API. - """ cfg = cls._get_config() models = [m["model"] for m in cfg["models"]] - default_model = models[0] if models else "" - # Use default_model resolutions so widget is created at instantiation + # restore last used model or default + default_model = cls.last_model if cls.last_model in models else (models[0] if models else "") + model_obj = next((m for m in cfg["models"] if m["model"] == default_model), None) if model_obj: resolutions = [r["name"] for r in model_obj.get("resolutions", [])] else: resolutions = ["Custom (manual)"] - # pick a default resolution that's non-custom if possible - default_res = next((n for n in resolutions if n != "Custom (manual)"), resolutions[0]) + # restore last used resolution if valid + if cls.last_resolution and cls.last_resolution in resolutions: + default_res = cls.last_resolution + else: + default_res = next((n for n in resolutions if n != "Custom (manual)"), resolutions[0]) - # determine defaults for width/height from default_res - default_w = 512 - default_h = 512 - if model_obj: - r_obj = next((r for r in model_obj["resolutions"] if r["name"] == default_res), None) + # determine defaults for width/height + default_w, default_h = 512, 512 + if cls.last_width and cls.last_height: + default_w, default_h = cls.last_width, cls.last_height + else: + r_obj = next((r for r in model_obj["resolutions"] if r["name"] == default_res), None) if model_obj else None if r_obj: - default_w = r_obj["width"] - default_h = r_obj["height"] + default_w, default_h = r_obj["width"], r_obj["height"] return { "required": { @@ -147,7 +153,14 @@ class Resolution_Machine: if model_obj and resolution and resolution != "Custom (manual)": res_obj = next((r for r in model_obj["resolutions"] if r["name"] == resolution), None) if res_obj: - return res_obj["width"], res_obj["height"] + width, height = res_obj["width"], res_obj["height"] + + # store last selections for persistence + Resolution_Machine.last_model = model + Resolution_Machine.last_resolution = resolution + Resolution_Machine.last_width = width + Resolution_Machine.last_height = height + return width, height diff --git a/web/prompt_machine.js b/web/prompt_machine.js index 5b878fe..01b9adf 100644 --- a/web/prompt_machine.js +++ b/web/prompt_machine.js @@ -1,30 +1,54 @@ +// prompt_machine_frontend.js import { app } from "../../../scripts/app.js"; app.registerExtension({ - name: "SATA_UtilityNode.PromptMachine", + name: "SATA_UtilityNode.PromptMachineFrontend", async nodeCreated(node) { - if (node.comfyClass === "Prompt_Machine") { - // When csv_file changes, refresh names list - const widget = node.widgets.find(w => w.name === "csv_file"); - if (widget) { - widget.callback = async (value) => { - try { - const resp = await fetch(`/sata/prompt_machine/names?csv=${encodeURIComponent(value)}`); - const data = await resp.json(); - const names = Array.isArray(data.names) && data.names.length > 0 ? data.names : ["None"]; - const nameWidget = node.widgets.find(w => w.name === "name"); - if (nameWidget) { - if (!nameWidget.options) nameWidget.options = {}; - nameWidget.options.values = names; - nameWidget.value = names[0]; - if (typeof node.setDirtyCanvas === "function") node.setDirtyCanvas(true); - if (typeof node.onWidgetChanged === "function") node.onWidgetChanged(nameWidget, nameWidget.value, null, 0); - } - } catch (e) { - console.error("[PromptMachine] Failed to fetch names:", e); - } - }; + if (node.comfyClass !== "Prompt_Machine") return; + + let csvWidget = node.widgets.find(w => w.name === "csv_file"); + let oldNameWidget = node.widgets.find(w => w.name === "name"); + + if (!csvWidget || !oldNameWidget) { + console.warn("[PromptMachineFrontend] widgets not found"); + return; + } + + // Remove the old STRING widget + const index = node.widgets.indexOf(oldNameWidget); + node.widgets.splice(index, 1); + + // Add a proper dropdown instead + let nameWidget = node.addWidget("combo", "name", "None", () => {}, { + values: ["None"] + }); + + async function refreshNames(csvValue) { + try { + const res = await fetch(`/sata/prompt_machine/names?csv=${encodeURIComponent(csvValue)}`); + const data = await res.json(); + const names = data.names || ["None"]; + + nameWidget.options.values = names; + if (!names.includes(nameWidget.value)) { + nameWidget.value = names[0]; + } + + if (node.graph) node.graph.setDirtyCanvas(true); + console.log("[PromptMachineFrontend] refreshed names:", names); + } catch (err) { + console.error("[PromptMachineFrontend] error fetching names:", err); } } + + // Initial load + refreshNames(csvWidget.value); + + // Reload names when CSV changes + const originalCsvCallback = csvWidget.callback; + csvWidget.callback = (value) => { + refreshNames(value); + if (originalCsvCallback) originalCsvCallback(value); + }; } }); \ No newline at end of file diff --git a/web/resolution_machine.js b/web/resolution_machine.js index 750d9f1..762f4f0 100644 --- a/web/resolution_machine.js +++ b/web/resolution_machine.js @@ -8,192 +8,116 @@ async function fetchConfig() { console.warn("ResolutionMachine: fetch failed", resp.status); return null; } - const json = await resp.json(); - return json; + return await resp.json(); } catch (e) { console.warn("ResolutionMachine: fetch exception", e); return null; } } -/* cross-version widget helpers */ -function setComboOptions(widget, options) { - if (!widget) return; - if (typeof widget.setOptions === "function") { - widget.setOptions(options); - return; - } - if (widget.options && Object.prototype.hasOwnProperty.call(widget.options, "values")) { - widget.options.values = options; - return; - } - widget.options = { values: options }; -} -function setWidgetValue(widget, v) { - if (!widget) return; - if (typeof widget.setValue === "function") return widget.setValue(v); - if (typeof widget.value === "function") return widget.value(v); - widget.value = v; -} -function getWidgetValue(widget) { - if (!widget) return undefined; - if (typeof widget.value === "function") return widget.value(); - return widget.value; -} -function triggerWidgetChange(widget, value) { - // call onChange or callback if present to simulate user selection - if (!widget) return; - if (typeof widget.onChange === "function") { - try { widget.onChange(value); return; } catch (e) { /* continue */ } - } - if (typeof widget.callback === "function") { - try { widget.callback(value); return; } catch (e) { /* continue */ } - } - // fallback: set value - setWidgetValue(widget, value); -} - -/* attach handler in tolerant way */ -function attachHandler(widget, node, handler) { - if (!widget) return; - if (typeof widget.onChange === "function") { - widget.onChange(handler); - return; - } - // preserve existing callback if present - const orig = widget.callback; - widget.callback = function (v) { - try { if (orig) orig.call(node, v); } catch (e) { /* ignore */ } - try { handler(v); } catch (e) { console.warn("handler error", e); } - }; -} - app.registerExtension({ name: "SATA_UtilityNode.ResolutionMachine", async nodeCreated(node) { if (node.comfyClass !== "Resolution_Machine") return; - const modelW = node.widgets.find(w => w.name === "model"); - const resW = node.widgets.find(w => w.name === "resolution"); - const wW = node.widgets.find(w => w.name === "width"); - const hW = node.widgets.find(w => w.name === "height"); + const cfg = await fetchConfig(); + if (!cfg) return; - console.log("ResolutionMachine nodeCreated widgets:", { modelW, resW, wW, hW }); + // find widgets + let modelW = node.widgets.find(w => w.name === "model"); + let resW = node.widgets.find(w => w.name === "resolution"); + let wW = node.widgets.find(w => w.name === "width"); + let hW = node.widgets.find(w => w.name === "height"); if (!modelW || !resW || !wW || !hW) { - console.warn("ResolutionMachine: missing widgets - ensure backend's INPUT_TYPES exposes model,resolution,width,height"); + console.warn("ResolutionMachine: widgets missing, check INPUT_TYPES"); return; } - const cfg = await fetchConfig(); - if (!cfg) { - console.warn("ResolutionMachine: no config received from backend"); - return; - } + // --- rebuild widgets as dropdowns to force combo type --- + const idxModel = node.widgets.indexOf(modelW); + node.widgets.splice(idxModel, 1); + modelW = node.addWidget("combo", "model", "None", () => {}, { + values: cfg.models.map(m => m.model) + }); - const models = cfg.models.map(m => m.model); - if (models.length > 0) setComboOptions(modelW, models); + const idxRes = node.widgets.indexOf(resW); + node.widgets.splice(idxRes, 1); + resW = node.addWidget("combo", "resolution", "Custom (manual)", () => {}, { + values: ["Custom (manual)"] + }); - // helper to get resolution list for a model (ensures Custom present) + // helper lookups function getResNamesForModel(modelName) { const m = cfg.models.find(x => x.model === modelName); if (!m) return ["Custom (manual)"]; - const names = m.resolutions.map(r => r.name.slice()); // copy + const names = m.resolutions.map(r => r.name); if (!names.includes("Custom (manual)")) names.push("Custom (manual)"); return names; } - - // helper to find resolution object function findResObj(modelName, resName) { const m = cfg.models.find(x => x.model === modelName); if (!m) return null; return m.resolutions.find(r => r.name === resName) || null; } - // guard to avoid manual-edit handler firing during programmatic updates - let suppressManualWH = false; + let suppressWH = false; function programmaticSetWH(wVal, hVal) { - suppressManualWH = true; - setWidgetValue(wW, wVal); - setWidgetValue(hW, hVal); + suppressWH = true; + wW.value = wVal; + hW.value = hVal; if (wW.inputEl) wW.inputEl.value = wVal; if (hW.inputEl) hW.inputEl.value = hVal; - // clear suppression on next tick - setTimeout(() => { suppressManualWH = false; }, 0); + setTimeout(() => { suppressWH = false; }, 0); } - // Model change handler: update resolution options and pick first preset - const onModelChange = (newModel) => { - try { - const names = getResNamesForModel(newModel); - setComboOptions(resW, names); - - // pick first non-custom if available - const pick = names[0] || "Custom (manual)"; - setWidgetValue(resW, pick); - // call resolution handler (simulate user selection) - triggerWidgetChange(resW, pick); - - // ensure node re-validation / UI update - try { node.setDirtyCanvas(true, true); } catch (e) { try { node.setDirtyCanvas(true); } catch(e){} } - } catch (e) { - console.warn("ResolutionMachine: onModelChange error", e); + // model handler + modelW.callback = (newModel) => { + const names = getResNamesForModel(newModel); + resW.options.values = names; + if (!names.includes(resW.value)) { + resW.value = names[0]; } + resW.callback(resW.value); // trigger res change + node.graph?.setDirtyCanvas(true); }; - // Resolution change handler: set width/height for presets, lock/unlock inputs accordingly - const onResolutionChange = (newRes) => { - try { - const modelName = getWidgetValue(modelW); - const resObj = findResObj(modelName, newRes); - if (resObj && newRes !== "Custom (manual)") { - programmaticSetWH(resObj.width, resObj.height); - // Lock width/height - if (wW.inputEl) wW.inputEl.disabled = true; - if (hW.inputEl) hW.inputEl.disabled = true; - } else { - // Custom: allow manual edits - if (wW.inputEl) wW.inputEl.disabled = false; - if (hW.inputEl) hW.inputEl.disabled = false; + // resolution handler + resW.callback = (newRes) => { + const modelName = modelW.value; + const resObj = findResObj(modelName, newRes); + if (resObj && newRes !== "Custom (manual)") { + programmaticSetWH(resObj.width, resObj.height); + if (wW.inputEl) wW.inputEl.disabled = true; + if (hW.inputEl) hW.inputEl.disabled = true; + } else { + if (wW.inputEl) wW.inputEl.disabled = false; + if (hW.inputEl) hW.inputEl.disabled = false; + } + node.graph?.setDirtyCanvas(true); + }; + + // manual WH handler + function onManualWH() { + if (suppressWH) return; + if (resW.value !== "Custom (manual)") { + const names = getResNamesForModel(modelW.value); + if (!names.includes("Custom (manual)")) { + names.push("Custom (manual)"); + resW.options.values = names; } - try { node.setDirtyCanvas(true, true); } catch (e) { try { node.setDirtyCanvas(true); } catch(e){} } - } catch (e) { - console.warn("ResolutionMachine: onResolutionChange error", e); + resW.value = "Custom (manual)"; + resW.callback("Custom (manual)"); } - }; + } + wW.callback = onManualWH; + hW.callback = onManualWH; - // Manual width/height change handler: set resolution to Custom (manual) - const onManualWH = (_newVal) => { - if (suppressManualWH) return; - try { - const cur = getWidgetValue(resW); - if (cur !== "Custom (manual)") { - const names = getResNamesForModel(getWidgetValue(modelW)); - if (!names.includes("Custom (manual)")) { - names.push("Custom (manual)"); - setComboOptions(resW, names); - } - setWidgetValue(resW, "Custom (manual)"); - try { triggerWidgetChange(resW, "Custom (manual)"); } catch(e){} - try { node.setDirtyCanvas(true, true); } catch (e) { try { node.setDirtyCanvas(true); } catch(e){} } - } - } catch (e) { - console.warn("ResolutionMachine: onManualWH error", e); - } - }; - - // Attach handlers (tolerant) - attachHandler(modelW, node, onModelChange); - attachHandler(resW, node, onResolutionChange); - attachHandler(wW, node, onManualWH); - attachHandler(hW, node, onManualWH); - - // Initialize model + resolution UI - const initialModel = getWidgetValue(modelW) || (models.length ? models[0] : undefined); - if (initialModel) { - setWidgetValue(modelW, initialModel); - // ensure options & width/height set for initial model - onModelChange(initialModel); + // --- initialize --- + const firstModel = cfg.models[0]?.model; + if (firstModel) { + modelW.value = firstModel; + modelW.callback(firstModel); } } }); \ No newline at end of file