This commit is contained in:
var1ableX
2024-09-13 21:23:12 -04:00
parent 56e932e997
commit d6bd841d7f
4 changed files with 197 additions and 55 deletions
+9 -2
View File
@@ -1,11 +1,18 @@
from .maskToCrop import GetMaskDimensions
from .unmasked import GetMaskDimensions, isMaskEmpty, GetRandomDimensions
# Node class mappings
NODE_CLASS_MAPPINGS = {
"GetMaskDimensions": GetMaskDimensions,
"isMaskEmpty": isMaskEmpty,
"GetRandomDimensions": GetRandomDimensions
}
# Node display name mappings
NODE_DISPLAY_NAME_MAPPINGS = {
"GetMaskDimensions": "Get Mask Dimensions",
}
"isMaskEmpty": "Is Mask Empty",
"GetRandomDimensions": "Get Random Dimensions"
}
WEB_DIRECTORY = "./web"
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
-53
View File
@@ -1,53 +0,0 @@
import torch
class GetMaskDimensions:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("INT", "INT", "INT", "INT",)
RETURN_NAMES = ("width", "height", "x", "y",)
FUNCTION = "execute"
CATEGORY = "image"
def execute(self, mask):
def get_mask_dimensions(mask):
# Assuming mask_tensor is of shape (1, H, W)
mask_array = mask[0]
# Find all coordinates where the mask is non-zero
coordinates = torch.nonzero(mask_array, as_tuple=False)
if coordinates.numel() == 0:
# If there are no non-zero coordinates, return a zero-sized rectangle
return 0, 0, 0, 0
# Get the minimum and maximum x and y coordinates
min_coords = torch.min(coordinates, dim=0)[0]
max_coords = torch.max(coordinates, dim=0)[0]
# Extract bounding box coordinates
top_left_x = min_coords[1].item()
top_left_y = min_coords[0].item()
bottom_right_x = max_coords[1].item()
bottom_right_y = max_coords[0].item()
# Calculate width and height
width = bottom_right_x - top_left_x + 1
height = bottom_right_y - top_left_y + 1
return width, height, top_left_x, top_left_y
# Process the mask tensor to find cropping dimensions
width, height, top_left_x, top_left_y = get_mask_dimensions(mask)
# Return the results as node outputs
return width, height, top_left_x, top_left_y
+120
View File
@@ -0,0 +1,120 @@
import torch
class GetMaskDimensions:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("INT", "INT", "INT", "INT",)
RETURN_NAMES = ("width", "height", "x", "y",)
FUNCTION = "execute"
CATEGORY = "image"
def execute(self, mask):
def get_mask_dimensions(mask):
# Assuming mask_tensor is of shape (1, H, W)
mask_array = mask[0]
# Find all coordinates where the mask is non-zero
coordinates = torch.nonzero(mask_array, as_tuple=False)
if coordinates.numel() == 0:
# If there are no non-zero coordinates, return a zero-sized rectangle
return 0, 0, 0, 0
# Get the minimum and maximum x and y coordinates
min_coords = torch.min(coordinates, dim=0)[0]
max_coords = torch.max(coordinates, dim=0)[0]
# Extract bounding box coordinates
top_left_x = min_coords[1].item()
top_left_y = min_coords[0].item()
bottom_right_x = max_coords[1].item()
bottom_right_y = max_coords[0].item()
# Calculate width and height
width = bottom_right_x - top_left_x + 1
height = bottom_right_y - top_left_y + 1
return width, height, top_left_x, top_left_y
# Process the mask tensor to find cropping dimensions
width, height, top_left_x, top_left_y = get_mask_dimensions(mask)
# Return the results as node outputs
return width, height, top_left_x, top_left_y
class isMaskEmpty:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("BOOLEAN",)
RETURN_NAMES = ("boolean",)
FUNCTION = "execute"
CATEGORY = "image"
def execute(self, mask):
def is_mask_empty(mask):
return not torch.any(mask == 1).item()
isEmpty = is_mask_empty(mask)
return (isEmpty,) # Return the boolean value as a tuple
import random
class GetRandomDimensions:
def __init__(self):
self.is_changed_enabled = True
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"min_width": ("INT", {"default": 768, "step": 16, "display": "number"}),
"min_height": ("INT", {"default": 768, "step": 16, "display": "number"}),
"max_width": ("INT", {"default": 1280, "step": 16, "display": "number"}),
"max_height": ("INT", {"default": 1280, "step": 16, "display": "number"}),
"randomize": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("INT", "INT")
RETURN_NAMES = ("width", "height",)
FUNCTION = "execute"
CATEGORY = "image"
OUTPUT_NODE = True
def execute(self, min_width, min_height, max_width, max_height, randomize ):
def get_random_dimensions(min_val, max_val):
min_div_16 = min_val // 16
max_div_16 = max_val // 16
if min_div_16 > max_div_16:
raise ValueError("No numbers in the specified range are divisible by 16.")
return random.randint(min_div_16, max_div_16) * 16
if randomize:
width = get_random_dimensions(min_width, max_width)
height = get_random_dimensions(min_height, max_height)
else:
width, height = max_width, max_height
# Setting the resolution string
text = f"{width}x{height}"
#return width, height
return {"ui": {"text": text},
"result": (width, height, text)}
+68
View File
@@ -0,0 +1,68 @@
import { app } from "../../../scripts/app.js";
import { ComfyWidgets } from "../../../scripts/widgets.js";
// Add new colors to LGraphCanvas.node_colors
function addNewColors() {
// Completely replacing the node_colors object
LGraphCanvas.node_colors = {
red: { color: "#C62828", bgcolor: "#E53935", groupcolor: "#B71C1C" }, // Red: 800, 600, 900
pink: { color: "#AD1457", bgcolor: "#D81B60", groupcolor: "#880E4F" }, // Pink: 800, 600, 900
purple: { color: "#6A1B9A", bgcolor: "#8E24AA", groupcolor: "#4A148C" }, // Purple: 800, 600, 900
indigo: { color: "#283593", bgcolor: "#3949AB", groupcolor: "#1A237E" }, // Indigo: 800, 600, 900
blue: { color: "#1565C0", bgcolor: "#1E88E5", groupcolor: "#0D47A1" }, // Blue: 800, 600, 900
cyan: { color: "#00838F", bgcolor: "#00ACC1", groupcolor: "#006064" }, // Cyan: 800, 600, 900
teal: { color: "#00695C", bgcolor: "#00897B", groupcolor: "#004D40" }, // Teal: 800, 600, 900
green: { color: "#2E7D32", bgcolor: "#43A047", groupcolor: "#1B5E20" }, // Green: 800, 600, 900
"light green": { color: "#558B2F", bgcolor: "#7CB342", groupcolor: "#33691E" }, // Light Green: 800, 600, 900
"deep orange": { color: "#D84315", bgcolor: "#F4511E", groupcolor: "#BF360C" }, // Deep Orange: 800, 600, 900
grey: { color: "#424242", bgcolor: "#757575", groupcolor: "#212121" }, // Grey: 800, 600, 900
// Adding the 6 extra colors
"blue grey": { color: "#455A64", bgcolor: "#607D8B", groupcolor: "#263238" }, // Blue Grey
lime: { color: "#9E9D24", bgcolor: "#C0CA33", groupcolor: "#827717" }, // Lime
"light blue": { color: "#0277BD", bgcolor: "#039BE5", groupcolor: "#01579B" }, // Light Blue
"steel blue": { color: "#4682B4", bgcolor: "#5A9BD3", groupcolor: "#2B5A8A" }, // Steel Blue
"turquoise": { color: "#00796B", bgcolor: "#009688", groupcolor: "#004D40" }, // Turquoise
navy: { color: "#003366", bgcolor: "#004080", groupcolor: "#001f3f" }, // Navy Blue
};
}
// Function to enhance nodes with text display
function addTextDisplay(nodeType) {
const onNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
const r = onNodeCreated?.apply(this, arguments);
const w = ComfyWidgets["STRING"](this, "display", ["STRING", { multiline: true, placeholder: " " }], app).widget;
w.inputEl.readOnly = true;
w.inputEl.style.opacity = 0.7;
w.inputEl.style.cursor = "auto";
return r;
};
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
onExecuted?.apply(this, arguments);
for (const widget of this.widgets) {
if (widget.type === "customtext" && widget.name === "display" && widget.inputEl.readOnly === true) {
widget.value = message.text.join('');
}
}
this.onResize?.(this.size);
};
}
app.registerExtension({
name: "rka.GetRandomDimensions",
setup() {
// Call the function to add new colors
addNewColors();
},
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (["GetRandomDimensions"].includes(nodeData.name)) {
addTextDisplay(nodeType);
}
},
});