diff --git a/__init__.py b/__init__.py index 543ed61..083a67e 100644 --- a/__init__.py +++ b/__init__.py @@ -15,6 +15,7 @@ from .src.load_lora_from_folder import LoadLoraFolderNode from .src.textbox import TextboxNode from .src.simple_sampler_custom import SimpleSamplerCustom from .src.show_sigmas import ShowSigmasNode +from .src.combine_schedulers import CombineSchedulersNode NODE_CLASS_MAPPINGS = { "Find SDXL Resolution": FindResolutionNode, @@ -33,6 +34,7 @@ NODE_CLASS_MAPPINGS = { "Copy/Paste Textbox": TextboxNode, "Simple Sampler Custom": SimpleSamplerCustom, "Show Sigmas": ShowSigmasNode, + "Combine Sigmas": CombineSchedulersNode, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -51,7 +53,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Load Lora Folder": "Load Lora Folder", "Copy/Paste Textbox": "Textbox", "Simple Sampler Custom": "Simple Sampler Custom", - "Show Sigmas": "Show Sigmas", + "Combine Sigmas": "Combine Sigmas", } WEB_DIRECTORY = "./web" diff --git a/src/combine_schedulers.py b/src/combine_schedulers.py new file mode 100644 index 0000000..299851b --- /dev/null +++ b/src/combine_schedulers.py @@ -0,0 +1,35 @@ +from comfy_extras.nodes_custom_sampler import SplitSigmasDenoise, SplitSigmas + +class CombineSchedulersNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sigmas_a": ("SIGMAS",), + "sigmas_b": ("SIGMAS",), + "split_at": ("INT", {"default": 10, "min": 10, "max": 10000}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}), + }, + } + + RETURN_TYPES = ("SIGMAS", ) + RETURN_NAMES = ("SIGMAS", ) + + FUNCTION = "get_combined_sigmas" + + CATEGORY = "sampling/custom_sampling" + + def get_combined_sigmas(self, sigmas_a, sigmas_b, split_at, denoise): + split_sigmas_a = SplitSigmas.get_sigmas(self, sigmas_a, split_at) + split_sigmas_b = SplitSigmas.get_sigmas(self, sigmas_b, split_at) + + sigmas = split_sigmas_a[0] + split_sigmas_b[1] + sigmas = self.get_denoised_sigmas(sigmas, denoise) + + return (sigmas,) + + def get_denoised_sigmas(self, sigmas, denoise): + if denoise == 1: + return (None, sigmas, ) + + return SplitSigmasDenoise.get_sigmas(self, sigmas, denoise) \ No newline at end of file diff --git a/src/show_sigmas.py b/src/show_sigmas.py index 5ca90da..df52b57 100644 --- a/src/show_sigmas.py +++ b/src/show_sigmas.py @@ -22,9 +22,9 @@ class ShowSigmasNode: FUNCTION = "run" CATEGORY = "utils" - - def run(self, sigmas, unique_id): - # Convert tensor to a numpy array if it's not already + + def sigmas_to_list(self, sigmas): + # Convert tensor to a numpy array if it's not already if isinstance(sigmas, np.ndarray): temporary_sigmas = sigmas else: @@ -42,13 +42,10 @@ class ShowSigmasNode: sigmas_list = [[float(value)] for value in temporary_sigmas] else: sigmas_list = temporary_sigmas.tolist() - - PromptServer.instance.send_sync("sn0w_get_sigmas", { - "id": unique_id, - "sigmas": sigmas_list, - }) - outputs = MessageHolder.waitForMessage(unique_id) + return sigmas_list + + def image_to_tensor(self, outputs): output_data = outputs['output'] if ',' in output_data: image_data = output_data.split(',', 1)[1] @@ -78,8 +75,17 @@ class ShowSigmasNode: print("Tensor shape for SaveImage node:", image_tensor.shape) - # Print final tensor shape - final_tensor = image_tensor - print("Final tensor shape to pass to SaveImage:", final_tensor.shape) + return image_tensor + + def run(self, sigmas, unique_id): + sigmas_list = self.sigmas_to_list(sigmas) + + PromptServer.instance.send_sync("sn0w_get_sigmas", { + "id": unique_id, + "sigmas": sigmas_list, + }) + outputs = MessageHolder.waitForMessage(unique_id) + + final_tensor = self.image_to_tensor(outputs) return (sigmas, final_tensor,) diff --git a/src/simple_sampler_custom.py b/src/simple_sampler_custom.py index 5cbfe8c..baf3f98 100644 --- a/src/simple_sampler_custom.py +++ b/src/simple_sampler_custom.py @@ -4,6 +4,7 @@ from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler from server import PromptServer from ..sn0w import Logger, MessageHolder from .custom_schedulers.custom_schedulers import CustomSchedulers +from .show_sigmas import ShowSigmasNode import comfy.samplers class SimpleSamplerCustom: @@ -17,7 +18,7 @@ class SimpleSamplerCustom: @classmethod def build_scheduler_list(cls): custom_scheduler_names = list(cls.custom_schedulers.get_scheduler_settings().keys()) - print(custom_scheduler_names) + cls.logger.log(custom_scheduler_names, "DEBUG") cls.scheduler_list = comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp"] + custom_scheduler_names cls.scheduler_settings = {name: settings for name, settings in cls.custom_schedulers.get_scheduler_settings().items()} @@ -55,8 +56,8 @@ class SimpleSamplerCustom: }, } - RETURN_TYPES = ("IMAGE", "LATENT", "CONDITIONING", "CONDITIONING") - RETURN_NAMES = ("IMAGE", "LATENT", "POSITIVE", "NEGATIVE") + RETURN_TYPES = ("IMAGE", "LATENT", "CONDITIONING", "CONDITIONING", "IMAGE") + RETURN_NAMES = ("IMAGE", "LATENT", "POSITIVE", "NEGATIVE", "SIGMAS GRAPH") FUNCTION = "sample" @@ -94,7 +95,17 @@ class SimpleSamplerCustom: else: image = (None, ) - return (image[0], samples[1], positive_prompt, negative_prompt) + sigmas_list = ShowSigmasNode.sigmas_to_list(self, sigmas) + + PromptServer.instance.send_sync("sampler_get_sigmas", { + "id": kwargs["unique_id"], + "sigmas": sigmas_list, + }) + outputs = MessageHolder.waitForMessage(kwargs["unique_id"]) + + final_tensor = ShowSigmasNode.image_to_tensor(self, outputs) + + return (image[0], samples[1], positive_prompt, negative_prompt, final_tensor,) def get_sampler(self, sampler_name): sampler = comfy.samplers.sampler_object(sampler_name) @@ -174,5 +185,4 @@ class SimpleSamplerCustom: "id": unique_id, }) outputs = MessageHolder.waitForMessage(unique_id) - print(outputs) return outputs diff --git a/web/show_sigmas.js b/web/show_sigmas.js index 69b9524..4c36f74 100644 --- a/web/show_sigmas.js +++ b/web/show_sigmas.js @@ -6,55 +6,13 @@ app.registerExtension({ name: "sn0w.ShowSigmas", async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeData.name === "Show Sigmas") { - // The populate function now returns the Base64 encoded image - function populate(sigmas) { - // Define the size of the canvas - const width = 800; - const height = 400; - - // Create a canvas element - const canvas = document.createElement('canvas'); - canvas.width = width; - canvas.height = height; - const ctx = canvas.getContext('2d'); - - // Background - ctx.fillStyle = 'white'; - ctx.fillRect(0, 0, width, height); - - // Normalize sigma values to fit within the canvas height - const maxSigma = Math.max(...sigmas.map(s => s[0])); - const minSigma = Math.min(...sigmas.map(s => s[0])); - const range = maxSigma - minSigma; - - // Function to scale a sigma value to the canvas coordinates - const scaleY = sigma => (height - 20) - (((sigma - minSigma) / range) * (height - 40)); - const scaleX = index => (index / (sigmas.length - 1)) * width; - - // Draw the graph line - ctx.strokeStyle = 'black'; - ctx.lineWidth = 2; - ctx.beginPath(); - ctx.moveTo(scaleX(0), scaleY(sigmas[0][0])); - - sigmas.forEach((sigma, index) => { - ctx.lineTo(scaleX(index), scaleY(sigma[0])); - }); - - ctx.stroke(); - - // Return the canvas content as a Base64 encoded image - return canvas.toDataURL('image/png'); - } - nodeType.prototype.onNodeCreated = function () { - console.log(this) api.addEventListener('sn0w_get_sigmas', (event) => { const data = event.detail; console.log(`Recieved ${data}`); if (this.id == data.id) { console.log(`Recieved ${data.sigmas}`); - const imageBase64 = populate(data.sigmas); + const imageBase64 = SettingUtils.drawSigmas(data.sigmas); // Send the generated image data back to the server api.fetchApi(`${SettingUtils.API_PREFIX}/get_sigmas`, { diff --git a/web/simpleSamplerCustom.js b/web/simpleSamplerCustom.js index 6f541b3..bd253e0 100644 --- a/web/simpleSamplerCustom.js +++ b/web/simpleSamplerCustom.js @@ -15,7 +15,6 @@ app.registerExtension({ nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { const slot = type[0] const inputType = type[1] - console.log(JSON.parse(JSON.stringify(targetSlot.inputs[slot]))) const initialCount = countWidgetsOfType(targetSlot.widgets, "converted-widget"); const originalSize = [targetSlot.size[0], targetSlot.size[1]]; @@ -88,7 +87,6 @@ app.registerExtension({ api.addEventListener('get_scheduler_values', (event) => { const data = event.detail const output = getWidgetOutputs(this, data.widgets_needed); - console.log(output); if (this.id == data.id) { api.fetchApi(`${SettingUtils.API_PREFIX}/scheduler_values`, { method: "POST", @@ -108,7 +106,6 @@ app.registerExtension({ api.addEventListener('should_decode_image', (event) => { const data = event.detail const output = checkImageOutput(this); - console.log(output); if (this.id == data.id) { api.fetchApi(`${SettingUtils.API_PREFIX}/should_decode_image`, { method: "POST", @@ -124,6 +121,29 @@ app.registerExtension({ }) } }) + + api.addEventListener('sampler_get_sigmas', (event) => { + const data = event.detail; + if (this.id == data.id) { + const imageBase64 = SettingUtils.drawSigmas(data.sigmas); + + // Send the generated image data back to the server + api.fetchApi(`${SettingUtils.API_PREFIX}/get_sigmas`, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify( + { + node_id: data.id, + outputs: { + output: imageBase64 + } + } + ), + }) + } + }); } nodeType.prototype.onConfigure = function () { @@ -196,25 +216,42 @@ app.registerExtension({ } function rearrangeWidgets(node, moveWidgetsBehind, totalWidgetsToMove) { - const widgets = node.widgets; - // Calculate the starting index for removing the last three widgets + let widgets = node.widgets; + + // Calculate the starting index for removing the last N widgets let startIndex = widgets.length - totalWidgetsToMove; - - // Check if there are at least three widgets to move + + // Check if there are enough widgets to move if (startIndex < 0) { console.warning("Not enough widgets to move."); return widgets; } - + + // Extract the last N widgets to be moved let widgetsToMove = widgets.splice(startIndex, totalWidgetsToMove); - + + // Find the scheduler widget and ensure it is not duplicated + const schedulerWidget = findWidget(node, "scheduler_name"); + + // If the schedulerWidget is found, remove existing occurrences from the current widgets array + if (schedulerWidget) { + widgets = widgets.filter(widget => widget !== schedulerWidget); + } + + // Ensure the insertion index is within the bounds of the current array + moveWidgetsBehind = Math.max(0, Math.min(moveWidgetsBehind, widgets.length)); + + // Prepare the new widgets array by inserting schedulerWidget followed by widgetsToMove let newWidgets = [ ...widgets.slice(0, moveWidgetsBehind), + schedulerWidget, ...widgetsToMove, ...widgets.slice(moveWidgetsBehind) ]; + // Update the node's widgets with the newly ordered widgets node.widgets = newWidgets; + return newWidgets; } @@ -295,8 +332,6 @@ app.registerExtension({ if (node.inputs.some(obj => obj.name === "height")) { node.inputs = node.inputs.filter(obj => obj.name !== "height"); } - - console.log(node) } function showAllLatentWidgets(node) { @@ -328,7 +363,7 @@ app.registerExtension({ const originalSize = [node.size[0], node.size[1]]; - rearrangeWidgets(node, 8, 10); + rearrangeWidgets(node, 7, allSettings.size); const nodesHidden = hideAllSchedulerWidgets(node); node.size[0] = originalSize[0]; diff --git a/web/sn0w.js b/web/sn0w.js index 022184c..bf43277 100644 --- a/web/sn0w.js +++ b/web/sn0w.js @@ -1,3 +1,5 @@ +import { api } from '../../scripts/api.js'; + export class SettingUtils { static API_PREFIX = '/sn0w'; @@ -176,4 +178,75 @@ export class SettingUtils { } } } + + static drawSigmas(sigmas) { + // Define the size of the canvas + const width = 800; + const height = 400; + + // Create a canvas element + const canvas = document.createElement('canvas'); + canvas.width = width; + canvas.height = height; + const ctx = canvas.getContext('2d'); + + // Background + ctx.fillStyle = '#333'; // Dark background + ctx.fillRect(0, 0, width, height); + + // Normalize sigma values to fit within the canvas height + const maxSigma = Math.max(...sigmas.map(s => s[0])); + const minSigma = Math.min(...sigmas.map(s => s[0])); + const range = maxSigma - minSigma; + + // Function to scale a sigma value to the canvas coordinates + const scaleY = sigma => (height - 20) - (((sigma - minSigma) / range) * (height - 40)); + const scaleX = index => (index / (sigmas.length - 1)) * width; + + // Draw horizontal grid lines and labels + const gridColor = '#555'; + const textColor = '#ccc'; + const numGridLines = 10; + const gridSpacing = (height - 40) / numGridLines; + ctx.strokeStyle = gridColor; + ctx.lineWidth = 1; + ctx.font = '10px Arial'; + ctx.fillStyle = textColor; + + for (let i = 0; i <= numGridLines; i++) { + let y = 20 + i * gridSpacing; + ctx.beginPath(); + ctx.moveTo(0, y); + ctx.lineTo(width, y); + ctx.stroke(); + + // Add text labels for grid lines + let gridValue = maxSigma - ((i / numGridLines) * range); + ctx.fillText(gridValue.toFixed(2), 5, y - 2); + } + + // Draw the graph line + ctx.strokeStyle = 'white'; // Line color + ctx.lineWidth = 2; + ctx.beginPath(); + ctx.moveTo(scaleX(0), scaleY(sigmas[0][0])); + + sigmas.forEach((sigma, index) => { + ctx.lineTo(scaleX(index), scaleY(sigma[0])); + }); + + ctx.stroke(); + + // Draw a dot at each sigma point + ctx.fillStyle = 'red'; // Dot color + const dotRadius = 4; + sigmas.forEach((sigma, index) => { + ctx.beginPath(); + ctx.arc(scaleX(index), scaleY(sigma[0]), dotRadius, 0, 2 * Math.PI); + ctx.fill(); + }); + + // Return the canvas content as a Base64 encoded image + return canvas.toDataURL('image/png'); + } }