nodes/sn0w12-ComfyUI-Sn0w-Scripts
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Simple sampler custom now kind of works

Browse Source
This commit is contained in:
Lucas Sköld
2024-05-09 03:19:40 +02:00
parent 7b2f06b1da
commit e2acf4da82
5 changed files with 181 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sn0w.py
+27 -1
View File
@@ -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__))
src/custom_schedulers/get_sigmas_sigmoid.py
+36
View File
@@ -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, )
src/simple_sampler_custom.py
+19 -6
View File
@@ -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):
web/simpleSamplerCustom.js
+98 -43
View File
@@ -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;
}
})
);
}
}
},
web/textbox.js
+1 -1
View File
@@ -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