File size: 3,193 Bytes
82ea528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import torch
import numpy as np
import cv2
def create_image_from_color(width, height, color="#FFFFFF"):
"""Create a solid color image with specified dimensions."""
if isinstance(color, str) and color.startswith('#'):
# Convert hex color to BGR (OpenCV format)
color = tuple(int(color[i:i+2], 16) for i in (5, 3, 1))[::-1]
return np.full((height, width, 3), color, dtype=np.uint8)
class AddMaskForICLora:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"first_image": ("IMAGE",),
"patch_mode": (["grid", "patch_right", "patch_bottom"], {
"default": "grid",
}),
"patch_number": ([2, 3, 4], {
"default": 4,
}),
"patch_color": (["#FF0000", "#00FF00", "#0000FF", "#FFFFFF"], {
"default": "#FF0000",
}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "INT", "INT", "INT", "INT")
RETURN_NAMES = ("IMAGE", "MASK", "x_offset", "y_offset", "target_width", "target_height", "total_width", "total_height")
FUNCTION = "add_mask"
OUTPUT_NODE = True
CATEGORY = "ICLoraUtils/AddMaskForICLora"
def add_mask(self, first_image, patch_mode, patch_number, patch_color):
# Get the first image from batch
image = first_image[0]
height, width = image.shape[:2]
# Calculate output dimensions
target_width = width
output_width = width * patch_number
# Ensure output width is divisible by 64
output_width = output_width - (output_width % 64) if output_width % 64 != 0 else output_width
# Create colored patch
colored_patch = create_image_from_color(target_width, height, patch_color)
colored_patch = torch.from_numpy(colored_patch)
# Create base masks
image_mask = torch.zeros((height, width))
patch_mask = torch.ones((height, width))
# Concatenate images and masks horizontally
patches = [image] + [colored_patch] * (patch_number - 1)
masks = [image_mask] + [patch_mask] * (patch_number - 1)
concatenated_image = np.hstack(patches)
concatenated_mask = np.hstack(masks)
# Calculate x offset (min_x)
x_offset = int((100 / patch_number) / 100.0 * concatenated_image.shape[1])
y_offset = 0
# Normalize image to [0, 1] range
concatenated_image = np.clip(255. * concatenated_image, 0, 255).astype(np.float32) / 255.0
# Convert to torch tensors with batch dimension
return_images = torch.from_numpy(concatenated_image)[None,]
return_masks = torch.from_numpy(concatenated_mask)[None,]
return (
return_images,
return_masks,
x_offset,
y_offset,
target_width,
height,
concatenated_image.shape[1],
concatenated_image.shape[0]
)
NODE_CLASS_MAPPINGS = {
"AddMaskForICLora": AddMaskForICLora,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AddMaskForICLora": "Add Mask For IC Lora x",
}
|