From e2acf4da82ee3d2f724d82bf37c8f9e6b6ab9c79 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lucas=20=20Sk=C3=B6ld?= Date: Thu, 9 May 2024 03:19:40 +0200 Subject: [PATCH] Simple sampler custom now kind of works --- sn0w.py | 28 +++- src/custom_schedulers/get_sigmas_sigmoid.py | 36 +++++ src/simple_sampler_custom.py | 25 +++- web/simpleSamplerCustom.js | 141 ++++++++++++++------ web/textbox.js | 2 +- 5 files changed, 181 insertions(+), 51 deletions(-) create mode 100644 src/custom_schedulers/get_sigmas_sigmoid.py diff --git a/sn0w.py b/sn0w.py index 2244523..4378f7a 100644 --- a/sn0w.py +++ b/sn0w.py @@ -43,7 +43,6 @@ class Logger: def reload_config(cls): cls.enabled_levels = ConfigReader.get_setting('sn0w.LoggingLevel', ["INFORMATIONAL", "WARNING"]) - @classmethod def print_sn0w(cls, message, color): print(f"{color}{cls.PREFIX}{cls.RESET_TEXT}{message}") @@ -61,6 +60,33 @@ class Logger: if level.upper() in self.enabled_levels or level.upper() in ["EMERGENCY", "ALERT", "CRITICAL", "ERROR"]: self.print_sn0w(message, color) + def print_sigmas_differences(self, name, sigmas): + """ + Takes a tensor of sigmas and prints each sigma along with the difference to the next sigma + and the percentage difference to the next sigma. + + Args: + sigmas (torch.Tensor): A 1D tensor of sigmas with a zero appended at the end. + """ + if ("DEBUG" in self.enabled_levels): + self.print_sn0w(f"Scheduler: {name}", self.PURPLE_TEXT) + print("Index | Sigma Value | Difference to Next | % Difference") + print("-" * 65) + + # Compute the differences + differences = sigmas[1:] - sigmas[:-1] + + # Iterate over sigmas and differences + for i in range(len(sigmas) - 1): + if sigmas[i] != 0: + percent_diff = (differences[i] / sigmas[i]) * 100 + else: + percent_diff = float('inf') # To handle division by zero in a meaningful way + print(f"{i:<5} | {sigmas[i]:<11.4f} | {differences[i]:<18.4f} | {percent_diff:<12.2f}") + + # Print the last sigma value without a difference (since it's the appended zero) + print(f"{len(sigmas) - 1:<5} | {sigmas[-1]:<11.4f} | {'N/A':<18} | {'N/A':<12}") + class Utility: logger = Logger() ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) diff --git a/src/custom_schedulers/get_sigmas_sigmoid.py b/src/custom_schedulers/get_sigmas_sigmoid.py new file mode 100644 index 0000000..74aa8ad --- /dev/null +++ b/src/custom_schedulers/get_sigmas_sigmoid.py @@ -0,0 +1,36 @@ +import torch + +def append_zero(x): + return torch.cat([x, x.new_zeros([1])]) + +def get_sigmas_sigmoid(self, n, sigma_max, sigma_min, steepness=10, midpoint_ratio=0.5, device='cpu'): + """ + Generates 'n' sigmas on an S-curve from sigma_max to sigma_min, with a controllable midpoint. + + Args: + n (int): Number of sigmas to generate. + sigma_min (float): Minimum value of sigma. + sigma_max (float): Maximum value of sigma. + steepness (float): Controls the steepness of the S-curve transition. + midpoint_ratio (float): Determines where the midpoint of the S-curve is (0 to 1). + device (str): Device to generate the sigmas on ('cpu' or 'cuda'). + + Returns: + torch.Tensor: A tensor of 'n' sigmas in descending order. + """ + # Ensure midpoint_ratio is within bounds + if not (0 <= midpoint_ratio <= 1): + raise ValueError("midpoint_ratio must be between 0 and 1.") + + # Adjust the range of x to skew the sigmoid's midpoint + midpoint_shift = torch.linspace(-steepness, steepness, n, device=device) + midpoint_location = -steepness + 2 * steepness * midpoint_ratio + x = midpoint_shift + (midpoint_location) + + # Sigmoid function values + sigmoids = 1 / (1 + torch.exp(-steepness * (x / steepness))) + + # Map the sigmoid output to the range [sigma_max, sigma_min] + sigmas = append_zero(sigma_max - (sigma_max - sigma_min) * sigmoids) + + return (sigmas, ) \ No newline at end of file diff --git a/src/simple_sampler_custom.py b/src/simple_sampler_custom.py index 050f6b0..9f4e9f6 100644 --- a/src/simple_sampler_custom.py +++ b/src/simple_sampler_custom.py @@ -3,11 +3,12 @@ from comfy_extras.nodes_custom_sampler import SamplerCustom, BasicScheduler, Pol from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler from server import PromptServer from ..sn0w import Logger, MessageHolder +from .custom_schedulers.get_sigmas_sigmoid import get_sigmas_sigmoid import comfy.samplers class SimpleSamplerCustom: logger = Logger() - scheduler_list = comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp"] + scheduler_list = comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp", "sigmoid"] @classmethod def INPUT_TYPES(cls): @@ -73,9 +74,10 @@ class SimpleSamplerCustom: sampler = comfy.samplers.sampler_object(sampler_name) return (sampler, ) - def get_scheduler_values(self, unique_id): + def get_scheduler_values(self, unique_id, widgets_needed): PromptServer.instance.send_sync("get_scheduler_values", { "id": unique_id, + "widgets_needed": widgets_needed, }) outputs = MessageHolder.waitForMessage(unique_id) return outputs @@ -84,16 +86,16 @@ class SimpleSamplerCustom: if scheduler_name == "align_your_steps": sigmas = AlignYourStepsScheduler.get_sigmas(self, model_type, steps, denoise)[0] elif scheduler_name == "polyexponential": - values = self.get_scheduler_values(unique_id) + values = self.get_scheduler_values(unique_id, 3) - sigma_max = values["sigma_max"]["value"] - sigma_min = values["sigma_min"]["value"] + sigma_max = values["sigma_max_poly"]["value"] + sigma_min = values["sigma_min_poly"]["value"] rho = values["rho"]["value"] self.logger.log(f"Sigma Max: {sigma_max}, Sigma Min: {sigma_min}, Rho: {rho}", "DEBUG") sigmas = PolyexponentialScheduler.get_sigmas(self, steps, sigma_max, sigma_min, rho)[0] elif scheduler_name == "vp": - values = self.get_scheduler_values(unique_id) + values = self.get_scheduler_values(unique_id, 3) beta_d = values["beta_d"]["value"] beta_min = values["beta_min"]["value"] @@ -101,9 +103,20 @@ class SimpleSamplerCustom: self.logger.log(f"Beta D: {beta_d}, Beta Min: {beta_min}, Eps S: {eps_s}", "DEBUG") sigmas = VPScheduler.get_sigmas(self, steps, beta_d, beta_min, eps_s)[0] + elif scheduler_name == "sigmoid": + values = self.get_scheduler_values(unique_id, 4) + + midpoint_ratio = values["midpoint_ratio"]["value"] + sigma_max = values["sigma_max_sig"]["value"] + sigma_min = values["sigma_min_sig"]["value"] + steepness = values["steepness"]["value"] + self.logger.log(f"Sigma Max: {sigma_max}, Sigma Min: {sigma_min}, Steepness: {steepness}, Midpoint Ratio: {midpoint_ratio}", "DEBUG") + + sigmas = get_sigmas_sigmoid(self, steps, sigma_max, sigma_min, steepness, midpoint_ratio)[0] else: sigmas = BasicScheduler.get_sigmas(self, model, scheduler_name, steps, denoise)[0] + self.logger.print_sigmas_differences(scheduler_name, sigmas) return sigmas def get_prompt(self, name, text_encode, clip, kwargs): diff --git a/web/simpleSamplerCustom.js b/web/simpleSamplerCustom.js index 8d3ae23..500c4dc 100644 --- a/web/simpleSamplerCustom.js +++ b/web/simpleSamplerCustom.js @@ -7,6 +7,26 @@ app.registerExtension({ name: "sn0w.SimpleSamplerCustom", async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeData.name === "Simple Sampler Custom") { + const widgets = { // Name: [Type, Default, Min, Max, Snap, Round] + "polyexponential": { + "sigma_max_poly": ["FLOAT", 14.614642, 0.0, 5000.0, 0.01, false], + "sigma_min_poly": ["FLOAT", 0.0291675, 0.0, 5000.0, 0.01, false], + "rho": ["FLOAT", 1.0, 0.0, 100.0, 0.01, false], + }, + "vp": { //TODO: fix default values + "beta_d": ["FLOAT", 14.0, 0.0, 5000.0, 0.01, false], + "beta_min": ["FLOAT", 0.05, 0.0, 5000.0, 0.01, false], + "eps_s": ["FLOAT", 0.075, 0.0, 1.0, 0.0001, false], + }, + "sigmoid": { + "sigma_max_sig": ["FLOAT", 25.0, 0.0, 5000.0, 0.01, false], + "sigma_min_sig": ["FLOAT", 0.0, 0.0, 5000.0, 0.01, false], + "steepness": ["FLOAT", 3.5, 0.0, 10.0, 0.01, false], + "midpoint_ratio": ["FLOAT", 0.8, 0.0, 1.0, 0.01, false], + }, + } + + // Set positive and negative connection colors const onConnectInput = nodeType.prototype.onConnectInput; nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { const slot = type[0] @@ -24,19 +44,27 @@ app.registerExtension({ onConnectInput?.apply(targetSlot, type, output, originNode, originSlot); } - const onNodeCreated = nodeType.prototype.onNodeCreated; nodeType.prototype.onNodeCreated = function () { const schedulerWidget = findWidget(this, "scheduler_name"); - + console.log(nodeType.prototype) + schedulerWidget.value = schedulerWidget.value; + + // Set up all inputs + createEverything(this); + console.log(app) + if (schedulerWidget && schedulerWidget.callback) { + const originalCallback = schedulerWidget.callback; schedulerWidget.callback = (newValue) => { setUpCustomInputs(this, newValue); + if (originalCallback) originalCallback.call(this, newValue); }; } api.addEventListener('get_scheduler_values', (event) => { const data = event.detail - const output = getWidgetOutputs(this, 3); + const output = getWidgetOutputs(this, data.widgets_needed); + console.log(output) if (this.id == data.id) { api.fetchApi(`${SettingUtils.API_PREFIX}/scheduler_values`, { method: "POST", @@ -52,13 +80,14 @@ app.registerExtension({ }) } }) - - onNodeCreated?.apply() } nodeType.prototype.onConfigure = function () { - const scheduler = findWidget(this, "scheduler_name"); - setUpCustomInputs(this, scheduler.value); + this.serializeCount++; + const inputName = findWidget(this, "scheduler_name"); + const desiredWidgets = widgets[inputName.value] ? Object.keys(widgets[inputName.value]) : []; + const removedWidgets = cleanUpInputs(this, desiredWidgets); + rearrangeWidgets(this, findWidgetIndex(this.widgets, inputName) + 1, 10 - removedWidgets); } function findWidget(node, name) { @@ -77,7 +106,6 @@ app.registerExtension({ function getWidgetOutputs(node, totalWidgetsToGet) { const widgets = node.widgets; let startIndex = widgets.length - (totalWidgetsToGet + 3); - if (startIndex < 0) { console.warning("Not enough widgets to move."); @@ -109,16 +137,32 @@ app.registerExtension({ return false } - function cleanUpInputs(node) { - const widgetsToRemove = ["sigma_max", "sigma_min", "rho", "beta_d", "beta_min", "eps_s"]; + function createWidgetsToRemove(widgets) { + // Use a set to avoid duplicates + const widgetNames = new Set(); + + // Iterate over each category in the widgets object + for (const category in widgets) { + // Get all keys from the category and add them to the set + Object.keys(widgets[category]).forEach(key => widgetNames.add(key)); + } + + // Convert the set back to an array + return Array.from(widgetNames); + } + + function cleanUpInputs(node, desiredWidgets) { + const widgetsToRemove = createWidgetsToRemove(widgets); let widgetsRemoved = 0; widgetsToRemove.forEach(widget => { - if (removeWidget(node, widget)) - widgetsRemoved++; + if (!desiredWidgets.includes(widget)) { + if (removeWidget(node, widget)) + widgetsRemoved++; + } }); return widgetsRemoved; - } + } function rearrangeWidgets(node, moveWidgetsBehind, totalWidgetsToMove) { const widgets = node.widgets; @@ -148,42 +192,53 @@ app.registerExtension({ function setUpCustomInputs(node, inputName) { const originalWidth = node.size[0]; const originalHeight = node.size[1]; - const widgetsRemoved = cleanUpInputs(node); + + // Determine which widgets should be present based on inputName + const desiredWidgets = widgets[inputName] ? Object.keys(widgets[inputName]) : []; + const existingWidgets = new Set(node.widgets.map(w => w.name)); + + // Clean up unwanted widgets: remove widgets not in the desired list + let widgetsRemoved = cleanUpInputs(node, desiredWidgets); + let widgetsAdded = 0; - - const widgets = { - "polyexponential": { - "sigma_max": ["FLOAT", 14.614642, 0.0, 5000.0, 0.01, false], - "sigma_min": ["FLOAT", 0.0291675, 0.0, 5000.0, 0.01, false], - "rho": ["FLOAT", 1.0, 0.0, 100.0, 0.01, false], - }, - "vp": { //TODO: fix default values - "beta_d": ["FLOAT", 14.0, 0.0, 5000.0, 0.01, false], - "beta_min": ["FLOAT", 0.05, 0.0, 5000.0, 0.01, false], - "eps_s": ["FLOAT", 0.075, 0.0, 1.0, 0.0001, false], - } - } - - const widget = widgets[inputName]; - if (widget != undefined) { - // Iterate over all widgets to add - Object.keys(widget).forEach(prop => { - const [type, defaultValue, min, max, step, round] = widget[prop]; - - if (type == "FLOAT") { - ComfyWidgets[type](node, prop, [type, { default: defaultValue, min: min, max: max, step: step, round: round }], app).widget; - widgetsAdded++; + + // Check and add new widgets only if they do not exist + if (widgets[inputName] !== undefined) { + Object.keys(widgets[inputName]).forEach(prop => { + // Check if this widget is already added + if (!existingWidgets.has(prop)) { + const [type, defaultValue, min, max, step, round] = widgets[inputName][prop]; + if (type === "FLOAT") { + let addedWidget = ComfyWidgets[type](node, prop, [type, { default: defaultValue, min: min, max: max, step: step, round: round }], app).widget; + console.log(addedWidget) + widgetsAdded++; + } } }); } + + // Rearrange widgets under 'scheduler_name' + if (widgetsAdded != 0) + rearrangeWidgets(node, findWidgetIndex(node.widgets, findWidget(node, "scheduler_name")) + 1, widgetsAdded); - // Put ne wwidgets under scheduler_name - rearrangeWidgets(node, findWidgetIndex(node.widgets, findWidget(node, "scheduler_name")) + 1, widgetsAdded) let totalWidgets = widgetsAdded - widgetsRemoved; - - // Add or remove height based on widgets added/ removed - node.size[0] = (originalWidth); - node.size[1] = (originalHeight + totalWidgets * (70 / 3)); + console.log(totalWidgets) + + // Adjust the size based on widgets added/removed + node.size[0] = originalWidth; + node.size[1] = originalHeight + totalWidgets * (70 / 3); + } + + function createEverything(node) { + Object.keys(widgets).forEach(widget => + Object.keys(widgets[widget]).forEach(prop => { + // Check if this widget is already added + const [type, defaultValue, min, max, step, round] = widgets[widget][prop]; + if (type === "FLOAT") { + ComfyWidgets[type](node, prop, [type, { default: defaultValue, min: min, max: max, step: step, round: round }], app).widget; + } + }) + ); } } }, diff --git a/web/textbox.js b/web/textbox.js index ad69ca9..1f4072c 100644 --- a/web/textbox.js +++ b/web/textbox.js @@ -15,7 +15,7 @@ app.registerExtension({ nodeType.prototype.populate = function() { // Adding a multiline text widget and storing a reference to the input element - let textWidget = ComfyWidgets["STRING"](this, "text", ["STRING", { default: "", multiline: true }], app).widget; + let textWidget = ComfyWidgets["STRING"](this, "text", ["STRING", { default: "hello", multiline: true }], app).widget; this.inputEl = textWidget.inputEl || textWidget.querySelector('input, textarea'); // Implement copy functionality with proper callback